onnx模型加入resize op

my code: 

# -*- coding: utf-8 -*-
import numpy as np
import onnx
from onnx import helper
from onnx import AttributeProto, TensorProto, GraphProto

import onnxruntime.backend as backend
import cv2
import onnxruntime as ort
import matplotlib.pyplot as plt



def create_resize_sub_graph(dst_height, dst_width):
    ## === 定义动态images
    INPUT_images = helper.make_tensor_value_info(
        name='images',
        elem_type=TensorProto.FLOAT,
        shape=["batch",3,"height","width"])
    #    shape=[1,3,"height","width"])
    #    shape=[1,3,240,240])
    
    ## === 获取输入images的shape
    VALUE_images_shape = helper.make_tensor_value_info(
        name='images_shape',
        elem_type=TensorProto.INT64,
        shape=[4])
    
    NODE_get_images_shape = helper.make_node(
        op_type='Shape',
        inputs=['images'],
        outputs=['images_shape'])
    
    ## === 获取输入images的batch
    VALUE_images_batch = helper.make_tensor_value_info(
        name='images_batch',
        elem_type=TensorProto.INT64,
        shape=[1])
    
    CONSTANT_indices_pspl0 = helper.make_tensor(
        name='indices_pspl0',
        data_type=TensorProto.INT64,
        dims=[1],
        vals=[0])
    
    NODE_get_images_batch = helper.make_node(
        op_type='Gather',
        inputs=['images_shape', 'indices_pspl0'],
        outputs=['images_batch'],
        axis=0)
    
    ## === 把batch和[3, 320, 640]进行concat
    CONSTANT_input_sizes_not_batch = helper.make_tensor(
        name='input_sizes_not_batch',
        data_type=TensorProto.INT64,
        dims=[3],
        vals=[3, dst_height, dst_width])
    
    VALUE_input_sizes = helper.make_tensor_value_info(
        name='input_sizes',
        elem_type=TensorProto.INT64,
        shape=[4])
    
    NODE_get_input_sizes = helper.make_node(
        op_type='Concat',
        inputs=['images_batch', 'input_sizes_not_batch'],
        outputs=['input_sizes'],
        axis=0)
    
    ## === 把images resize到指定大小
    VALUE_images_resized = helper.make_tensor_value_info(
        name='images_resized',
        elem_type=TensorProto.FLOAT,
        shape=['batch', 3, dst_height, dst_width])

    
    NODE_resize = helper.make_node(
        name='resize_input_images',
        op_type='Resize',
    #    inputs=['images', 'roi_0','scales_0','input_sizes'],
        inputs=['images', '','','input_sizes'],
        outputs=['images_resized'],
        mode='linear')
    
    
    nodes = [NODE_get_images_shape, NODE_get_images_batch, NODE_get_input_sizes, NODE_resize]
    inputs = [INPUT_images]
    outputs = [VALUE_images_resized]
    constants = [CONSTANT_indices_pspl0, CONSTANT_input_sizes_not_batch]
    return nodes, inputs, outputs, constants



nodes, inputs, outputs, constants = create_resize_sub_graph(320, 640)




## 到这里可以生成临时模型测试
graph_def = helper.make_graph(
    nodes=nodes,
    name='test-model',
    inputs=inputs,
    outputs=outputs)

## === 常量放进initializer
graph_def.initializer.extend(constants)

model_def = helper.make_model(
    graph=graph_def,
    producer_name='onnx_example')
onnx.checker.check_model(model_def)

onnx.save(model_def, "test_with_resize.onnx")

raise ValueError


# model_def 是 model_def = helper.make_model() 和 onnx.checker.check_model(model_def) 的
rt = backend.prepare(model_def, 'CPU')
out = rt.run(np.ones([1, 3, 250, 250], dtype=np.float32))
print(out[0].shape)

raise ValueError





### === example: 加入到Yolov5 onnx模型中
# === 读取原模型graph
model_path = r"xxx.onnx"
model_onnx = onnx.load(model_path)

graph = model_onnx.graph


# === 修改原模型 前面使用到"images"变量的节点
old_nodes = list(graph.node)

# 原来images变量的下游节点的输入改为images_resized
for i in range(len(old_nodes)):
    n = old_nodes[i]
    if n.name in ('Slice_4', 'Slice_14',
                  'Slice_24', 'Slice_34'):
        n.input[0] = "images_resized"


# === 创建新的graph把原模型graph上的 node/outputs/initializer 加进去
# === (这里测试过在原模型的graph上修改进行保存, 然而会报错 Resize节点的roi和scales必须指定, 测试版本onnx==1.9/1.10)
graph_def = helper.make_graph(
    nodes=nodes+old_nodes,
    name='test-model',
    inputs=inputs,
    outputs=list(graph.output))

# === 常量放进initializer
graph_def.initializer.extend(
    constants + list(graph.initializer))

# === 生成模型并测试/保存
model_def = helper.make_model(
    graph=graph_def,
    producer_name='onnx_example')
onnx.checker.check_model(model_def)

onnx.save(model_def, "test_with_resize.onnx")

raise ValueError


# model_def 是 model_def = helper.make_model() 和 onnx.checker.check_model(model_def) 的
rt = backend.prepare(model_def, 'CPU')
out = rt.run(np.ones([1, 3, 250, 250], dtype=np.float32))

raise ValueError

end