|
|
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) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
