| import torch |
| import axengine as axe |
| import numpy as np |
| import math |
|
|
| def _get_source_mask(src_seq, valid_ratios) -> torch.Tensor: |
| """Generate mask for source sequence. |
| |
| Args: |
| src_seq (torch.Tensor): Image sequence. Shape :math:`(N, T, C)`. |
| valid_ratios (list[float]): The valid ratio of input image. For |
| example, if the width of the original image is w1 and the width |
| after padding is w2, then valid_ratio = w1/w2. Source mask is |
| used to cover the area of the padding region. |
| |
| Returns: |
| Tensor or None: Source mask. Shape :math:`(N, T)`. The region of |
| padding area are False, and the rest are True. |
| """ |
|
|
| N, T, _ = src_seq.size() |
| mask = None |
| if len(valid_ratios) > 0: |
| mask = src_seq.new_zeros((N, T), device=src_seq.device) |
| for i, valid_ratio in enumerate(valid_ratios): |
| valid_width = min(T, math.ceil(T * valid_ratio)) |
| mask[i, :valid_width] = 1 |
|
|
| return mask |
|
|
| onnx_bb_encoder = axe.InferenceSession("backbone_encoder.axmodel") |
| onnx_decoder = axe.InferenceSession("decoder.axmodel") |
|
|
| input_image = torch.tensor(np.load('input_tensor/input_image.npy')) |
| out_enc = onnx_bb_encoder.run(["output"], {"input": np.array(input_image.cpu())})[0] |
| out_enc = torch.tensor(out_enc) |
| data_samples = None |
|
|
| N = out_enc.size(0) |
|
|
| init_target_seq = torch.tensor(np.load('input_tensor/init_target_seq.npy')).to(torch.int32) |
| outputs = [] |
| max_seq_len = 25 |
| for step in range(0, max_seq_len): |
| valid_ratios = [1.0 for _ in range(out_enc.size(0))] |
| if data_samples is not None: |
| valid_ratios = [] |
| for data_sample in data_samples: |
| valid_ratios.append(data_sample.get('valid_ratio')) |
| |
| src_mask = _get_source_mask(out_enc, valid_ratios) |
| |
| step_result = onnx_decoder.run(["output"],{'init_target_seq':np.array(init_target_seq), |
| 'out_enc':np.array(out_enc), |
| 'src_mask':np.array(src_mask), |
| 'step':np.array([step]).astype(np.int32)})[0][0] |
| step_result = torch.tensor(step_result) |
| outputs.append(step_result) |
| _, step_max_index = torch.max(step_result, dim=-1) |
| init_target_seq[:, step + 1] = step_max_index |
| outputs = torch.stack(outputs, dim=1) |
| np.save('output_tensor/outputs.npy',outputs) |
|
|
| |
| |
| |
| |
| |
