2020"""
2121
2222
23- from typing import List , Tuple , Union
23+ from tempfile import NamedTemporaryFile
24+ from typing import List , Optional , Tuple , Union
2425
2526import cv2
2627import numpy
28+ import onnx
2729import torch
2830
31+ from sparseml .onnx .utils import get_tensor_dim_shape , set_tensor_dim_shape
32+ from sparsezoo import Zoo
33+
2934# ultralytics/yolov5 imports
3035from utils .general import non_max_suppression
3136
3742]
3843
3944
40- def _get_grid (size : int ) -> torch .Tensor :
41- # adapted from yolov5.yolo.Detect._make_grid
42- coords_y , coords_x = torch .meshgrid ([torch .arange (size ), torch .arange (size )])
43- grid = torch .stack ((coords_x , coords_y ), 2 )
44- return grid .view (1 , 1 , size , size , 2 ).float ()
45-
46-
4745# Yolo V3 specific variables
4846_YOLO_V3_ANCHORS = [
4947 torch .Tensor ([[10 , 13 ], [16 , 30 ], [33 , 23 ]]),
5048 torch .Tensor ([[30 , 61 ], [62 , 45 ], [59 , 119 ]]),
5149 torch .Tensor ([[116 , 90 ], [156 , 198 ], [373 , 326 ]]),
5250]
5351_YOLO_V3_ANCHOR_GRIDS = [t .clone ().view (1 , - 1 , 1 , 1 , 2 ) for t in _YOLO_V3_ANCHORS ]
54- _YOLO_V3_OUTPUT_SHAPES = [80 , 40 , 20 ]
55- _YOLO_V3_GRIDS = [_get_grid (grid_size ) for grid_size in _YOLO_V3_OUTPUT_SHAPES ]
5652
5753
5854def load_image (
@@ -70,28 +66,53 @@ def load_image(
7066 return img
7167
7268
73- def pre_nms_postprocess ( outputs : List [ numpy . ndarray ]) -> torch . Tensor :
69+ class YoloPostprocessor :
7470 """
75- :param outputs: raw outputs of a YOLOv3 model before anchor grid processing
76- :return: post-processed model outputs without NMS.
77- """
78- # postprocess and transform raw outputs into single torch tensor
79- processed_outputs = []
80- for idx , pred in enumerate (outputs ):
81- pred = torch .from_numpy (pred )
82- pred = pred .sigmoid ()
71+ Class for performing postprocessing of YOLOv3 model predictions
8372
84- # get grid and stride
85- grid = _YOLO_V3_GRIDS [idx ]
86- anchor_grid = _YOLO_V3_ANCHOR_GRIDS [idx ]
87- stride = 640 / _YOLO_V3_OUTPUT_SHAPES [idx ]
73+ :param image_size: size of input image to model. used to calculate stride based on
74+ output shapes
75+ """
8876
89- # decode xywh box values
90- pred [..., 0 :2 ] = (pred [..., 0 :2 ] * 2.0 - 0.5 + grid ) * stride
91- pred [..., 2 :4 ] = (pred [..., 2 :4 ] * 2 ) ** 2 * anchor_grid
92- # flatten anchor and grid dimensions -> (bs, num_predictions, num_classes + 5)
93- processed_outputs .append (pred .view (pred .size (0 ), - 1 , pred .size (- 1 )))
94- return torch .cat (processed_outputs , 1 )
77+ def __init__ (self , image_size : Tuple [int ]):
78+ self ._image_size = image_size
79+ self ._grids = {} # Dict[Tuple[int], torch.Tensor]
80+
81+ def pre_nms_postprocess (self , outputs : List [numpy .ndarray ]) -> torch .Tensor :
82+ """
83+ :param outputs: raw outputs of a YOLOv3 model before anchor grid processing
84+ :return: post-processed model outputs without NMS.
85+ """
86+ # postprocess and transform raw outputs into single torch tensor
87+ processed_outputs = []
88+ for idx , pred in enumerate (outputs ):
89+ pred = torch .from_numpy (pred )
90+ pred = pred .sigmoid ()
91+
92+ # get grid and stride
93+ grid_shape = pred .shape [2 :4 ]
94+ grid = self ._get_grid (grid_shape )
95+ anchor_grid = _YOLO_V3_ANCHOR_GRIDS [idx ]
96+ stride = self ._image_size [0 ] / grid_shape [0 ]
97+
98+ # decode xywh box values
99+ pred [..., 0 :2 ] = (pred [..., 0 :2 ] * 2.0 - 0.5 + grid ) * stride
100+ pred [..., 2 :4 ] = (pred [..., 2 :4 ] * 2 ) ** 2 * anchor_grid
101+ # flatten anchor and grid dimensions -> (bs, num_predictions, num_classes + 5)
102+ processed_outputs .append (pred .view (pred .size (0 ), - 1 , pred .size (- 1 )))
103+ return torch .cat (processed_outputs , 1 )
104+
105+ def _get_grid (self , grid_shape : Tuple [int ]) -> torch .Tensor :
106+ if grid_shape not in self ._grids :
107+ # adapted from yolov5.yolo.Detect._make_grid
108+ coords_y , coords_x = torch .meshgrid (
109+ [torch .arange (grid_shape [0 ]), torch .arange (grid_shape [1 ])]
110+ )
111+ grid = torch .stack ((coords_x , coords_y ), 2 )
112+ self ._grids [grid_shape ] = grid .view (
113+ 1 , 1 , grid_shape [0 ], grid_shape [1 ], 2
114+ ).float ()
115+ return self ._grids [grid_shape ]
95116
96117
97118def postprocess_nms (outputs : torch .Tensor ) -> List [numpy .ndarray ]:
@@ -102,3 +123,60 @@ def postprocess_nms(outputs: torch.Tensor) -> List[numpy.ndarray]:
102123 # run nms in PyTorch, only post-process first output
103124 nms_outputs = non_max_suppression (outputs )
104125 return [output .cpu ().numpy () for output in nms_outputs ]
126+
127+
128+ def modify_yolo_onnx_input_shape (
129+ model_path : str , image_shape : Tuple [int ]
130+ ) -> Tuple [str , Optional [NamedTemporaryFile ]]:
131+ """
132+ Creates a new YOLOv3 ONNX model from the given path that accepts the given input
133+ shape. If the given model already has the given input shape no modifications are
134+ made. Uses a tempfile to store the modified model file.
135+
136+ :param model_path: file path to YOLOv3 ONNX model or SparseZoo stub of the model
137+ to be loaded
138+ :param image_shape: 2-tuple of the image shape to resize this yolo model to
139+ :return: filepath to an onnx model reshaped to the given input shape will be the
140+ original path if the shape is the same. Additionally returns the
141+ NamedTemporaryFile for managing the scope of the object for file deletion
142+ """
143+ original_model_path = model_path
144+ if model_path .startswith ("zoo:" ):
145+ # load SparseZoo Model from stub
146+ model = Zoo .load_model_from_stub (model_path )
147+ model_path = model .onnx_file .downloaded_path ()
148+ print (f"Downloaded { original_model_path } { model_path } )
149+
150+ model = onnx .load (model_path )
151+ model_input = model .graph .input [0 ]
152+
153+ initial_x = get_tensor_dim_shape (model_input , 2 )
154+ initial_y = get_tensor_dim_shape (model_input , 3 )
155+
156+ if not (isinstance (initial_x , int ) and isinstance (initial_y , int )):
157+ return model_path , None # model graph does not have static integer input shape
158+
159+ if (initial_x , initial_y ) == tuple (image_shape ):
160+ return model_path , None # no shape modification needed
161+
162+ scale_x = initial_x / image_shape [0 ]
163+ scale_y = initial_y / image_shape [1 ]
164+ set_tensor_dim_shape (model_input , 2 , image_shape [0 ])
165+ set_tensor_dim_shape (model_input , 3 , image_shape [1 ])
166+
167+ for model_output in model .graph .output :
168+ output_x = get_tensor_dim_shape (model_output , 2 )
169+ output_y = get_tensor_dim_shape (model_output , 3 )
170+ set_tensor_dim_shape (model_output , 2 , int (output_x / scale_x ))
171+ set_tensor_dim_shape (model_output , 3 , int (output_y / scale_y ))
172+
173+ tmp_file = NamedTemporaryFile () # file will be deleted after program exit
174+ onnx .save (model , tmp_file .name )
175+
176+ print (
177+ f"Overwriting original model shape { (initial_x , initial_y )} { image_shape } \n "
178+ f"Original model path: { original_model_path }
179+ f"{ tmp_file .name }
180+ )
181+
182+ return tmp_file .name , tmp_file
0 commit comments