batchno_classid_score_x1y1x2y2 のNMS出力を使用して画像から複数のROIを抽出して固定サイズにリサイズ後にスタックするONNX処理の模索

#! /usr/bin/env python

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import onnx
from onnxsim import simplify
import numpy as np
np.random.seed(0)


class Model(nn.Module):
    def __init__(
        self,
    ):
        super(Model, self).__init__()

        self.resize_height = 128
        self.resize_width = 64

    def forward(self, image: torch.Tensor, nms_output: torch.Tensor):
        break_idx = nms_output.shape[0]
        elements = nms_output.shape[1]
        n, c, h, w = image.shape
        # batchno_classid_score_x1y1x2y2
        result_tensor = torch.zeros([break_idx, 3, self.resize_height, self.resize_width])
        for counter in range(break_idx):
            batch = nms_output[counter, 0].to(torch.int64)
            classid = nms_output[counter, 1].to(torch.int64)
            score = nms_output[counter, 2]
            x1 = (nms_output[counter, 3] * w).to(torch.int64)
            y1 = (nms_output[counter, 4] * h).to(torch.int64)
            x2 = (nms_output[counter, 5] * w).to(torch.int64)
            y2 = (nms_output[counter, 6] * h).to(torch.int64)
            partial_image = image[:, :, y1:y2, x1:x2]
            resized_partial_image = \
                F.interpolate(
                    input=partial_image,
                    size=(self.resize_height, self.resize_width),
                    mode='bilinear',
                )
            result_tensor[counter:counter+1, ...] = resized_partial_image
        return result_tensor

def generate_tensor():
    # [3, 7] 形状のテンソル初期化
    tensor = torch.empty(3, 7)
    # batchnoとclassidにランダムな実数を割り当て (ここでは範囲を0.0から10.0、0.0から5.0とする)
    tensor[:, 0] = torch.randint(0, 11, (3,)).float()  # batchno
    tensor[:, 1] = torch.randint(0, 6, (3,)).float()    # classid
    # scoreに0.1から1.0のランダムな実数を割り当て
    tensor[:, 2] = 0.9 * torch.rand(3) + 0.1
    # x1, x2 および y1, y2 に対してランダムな実数を割り当て、適切な条件 x1 < x2, y1 < y2 を満たすように並べ替え
    for i in range(3):
        x1, x2 = sorted(0.9 * torch.rand(2) + 0.1)
        y1, y2 = sorted(0.9 * torch.rand(2) + 0.1)
        tensor[i, 3] = x1
        tensor[i, 4] = y1
        tensor[i, 5] = x2
        tensor[i, 6] = y2
    return tensor

if __name__ == "__main__":
    OPSET=13
    MODEL = f'NMS_Loop'
    model = Model()
    onnx_file = f"{MODEL}_{OPSET}.onnx"

    image = torch.randn([1, 3, 480, 640], dtype=torch.float32)
    # batchno_classid_score_x1y1x2y2
    nms_output = generate_tensor()
    torch.onnx.export(
        torch.jit.script(model),
        args=(image, nms_output),
        f=onnx_file,
        opset_version=OPSET,
        input_names=[
            f'input_image',
            f'input_batchno_classid_score_x1y1x2y2',
        ],
        output_names=[
            f'output_roi',
        ],
        dynamic_axes={
            f'input_batchno_classid_score_x1y1x2y2': {0: 'N'},
            f'output_roi': {0: 'N', 1: '3', 2: str(model.resize_height), 3: str(model.resize_width)},
        },
    )
    model_onnx1 = onnx.load(onnx_file)
    model_onnx1 = onnx.shape_inference.infer_shapes(model_onnx1)
    onnx.save(model_onnx1, onnx_file)
    model_onnx2 = onnx.load(onnx_file)
    model_simp, check = simplify(model_onnx2)
    onnx.save(model_simp, onnx_file)