@@ -78,3 +78,160 @@ def get_label_pos(contour):
7878def draw_contours (img , contours , color = "info" ):
7979 if color == "info" :
8080 cv2 .drawContours (img , contours , - 1 , (255 , 145 , 30 ), 3 )
81+
82+
83+ def yolox_preprocess (img , input_size , swap = (2 , 0 , 1 )):
84+ if len (img .shape ) == 3 :
85+ padded_img = numpy .ones ((input_size [0 ], input_size [1 ], 3 ), dtype = numpy .uint8 ) * 114
86+ else :
87+ padded_img = numpy .ones (input_size , dtype = numpy .uint8 ) * 114
88+ r = min (input_size [0 ] / img .shape [0 ], input_size [1 ] / img .shape [1 ])
89+ resized_img = cv2 .resize (
90+ img ,
91+ (int (img .shape [1 ] * r ), int (img .shape [0 ] * r )),
92+ interpolation = cv2 .INTER_LINEAR ,
93+ ).astype (numpy .uint8 )
94+ padded_img [: int (img .shape [0 ] * r ), : int (img .shape [1 ] * r )] = resized_img
95+ padded_img = padded_img .transpose (swap )
96+ padded_img = numpy .ascontiguousarray (padded_img , dtype = numpy .float32 )
97+ return padded_img , r
98+
99+
100+ def yolox_postprocess (outputs , img_size , p6 = False ):
101+ grids = []
102+ expanded_strides = []
103+ if not p6 :
104+ strides = [8 , 16 , 32 ]
105+ else :
106+ strides = [8 , 16 , 32 , 64 ]
107+ hsizes = [img_size [0 ] // stride for stride in strides ]
108+ wsizes = [img_size [1 ] // stride for stride in strides ]
109+ for hsize , wsize , stride in zip (hsizes , wsizes , strides ):
110+ xv , yv = numpy .meshgrid (numpy .arange (wsize ), numpy .arange (hsize ))
111+ grid = numpy .stack ((xv , yv ), 2 ).reshape (1 , - 1 , 2 )
112+ grids .append (grid )
113+ shape = grid .shape [:2 ]
114+ expanded_strides .append (numpy .full ((* shape , 1 ), stride ))
115+ grids = numpy .concatenate (grids , 1 )
116+ expanded_strides = numpy .concatenate (expanded_strides , 1 )
117+ outputs [..., :2 ] = (outputs [..., :2 ] + grids ) * expanded_strides
118+ outputs [..., 2 :4 ] = numpy .exp (outputs [..., 2 :4 ]) * expanded_strides
119+ return outputs
120+
121+
122+ def nms (boxes , scores , nms_thr ):
123+ """Single class NMS implemented in Numpy."""
124+ x1 = boxes [:, 0 ]
125+ y1 = boxes [:, 1 ]
126+ x2 = boxes [:, 2 ]
127+ y2 = boxes [:, 3 ]
128+
129+ areas = (x2 - x1 + 1 ) * (y2 - y1 + 1 )
130+ order = scores .argsort ()[::- 1 ]
131+
132+ keep = []
133+ while order .size > 0 :
134+ i = order [0 ]
135+ keep .append (i )
136+ xx1 = numpy .maximum (x1 [i ], x1 [order [1 :]])
137+ yy1 = numpy .maximum (y1 [i ], y1 [order [1 :]])
138+ xx2 = numpy .minimum (x2 [i ], x2 [order [1 :]])
139+ yy2 = numpy .minimum (y2 [i ], y2 [order [1 :]])
140+
141+ w = numpy .maximum (0.0 , xx2 - xx1 + 1 )
142+ h = numpy .maximum (0.0 , yy2 - yy1 + 1 )
143+ inter = w * h
144+ ovr = inter / (areas [i ] + areas [order [1 :]] - inter )
145+
146+ inds = numpy .where (ovr <= nms_thr )[0 ]
147+ order = order [inds + 1 ]
148+
149+ return keep
150+
151+
152+ def multiclass_nms (boxes , scores , nms_thr , score_thr , class_agnostic = True ):
153+ """Multiclass NMS implemented in Numpy"""
154+ if class_agnostic :
155+ nms_method = multiclass_nms_class_agnostic
156+ else :
157+ nms_method = multiclass_nms_class_aware
158+ return nms_method (boxes , scores , nms_thr , score_thr )
159+
160+
161+ def multiclass_nms_class_agnostic (boxes , scores , nms_thr , score_thr ):
162+ """Multiclass NMS implemented in Numpy. Class-agnostic version."""
163+ cls_inds = scores .argmax (1 )
164+ cls_scores = scores [numpy .arange (len (cls_inds )), cls_inds ]
165+
166+ valid_score_mask = cls_scores > score_thr
167+ if valid_score_mask .sum () == 0 :
168+ return None
169+ valid_scores = cls_scores [valid_score_mask ]
170+ valid_boxes = boxes [valid_score_mask ]
171+ valid_cls_inds = cls_inds [valid_score_mask ]
172+ keep = nms (valid_boxes , valid_scores , nms_thr )
173+ if keep :
174+ dets = numpy .concatenate (
175+ [valid_boxes [keep ], valid_scores [keep , None ], valid_cls_inds [keep , None ]], 1
176+ )
177+ return dets
178+
179+
180+ def multiclass_nms_class_aware (boxes , scores , nms_thr , score_thr ):
181+ """Multiclass NMS implemented in Numpy. Class-aware version."""
182+ final_dets = []
183+ num_classes = scores .shape [1 ]
184+ for cls_ind in range (num_classes ):
185+ cls_scores = scores [:, cls_ind ]
186+ valid_score_mask = cls_scores > score_thr
187+ if valid_score_mask .sum () == 0 :
188+ continue
189+ else :
190+ valid_scores = cls_scores [valid_score_mask ]
191+ valid_boxes = boxes [valid_score_mask ]
192+ keep = nms (valid_boxes , valid_scores , nms_thr )
193+ if len (keep ) > 0 :
194+ cls_inds = numpy .ones ((len (keep ), 1 )) * cls_ind
195+ dets = numpy .concatenate (
196+ [valid_boxes [keep ], valid_scores [keep , None ], cls_inds ], 1
197+ )
198+ final_dets .append (dets )
199+ if len (final_dets ) == 0 :
200+ return None
201+ return numpy .concatenate (final_dets , 0 )
202+
203+
204+ def img_show (img , boxes , scores , cls_ids , conf = 0.5 , class_names = None ):
205+ _COLORS = numpy .array ([255 , 0 , 0 ,
206+ 195 , 123 , 40 ,
207+ 110 , 176 , 23 ]).astype (numpy .float32 ).reshape (- 1 , 3 )
208+ for i in range (len (boxes )):
209+ box = boxes [i ]
210+ cls_id = int (cls_ids [i ])
211+ score = scores [i ]
212+ if score < conf :
213+ continue
214+ x0 = int (box [0 ])
215+ y0 = int (box [1 ])
216+ x1 = int (box [2 ])
217+ y1 = int (box [3 ])
218+
219+ color = _COLORS [cls_id ].astype (numpy .uint8 ).tolist ()
220+ text = '{}:{:.1f}%' .format (class_names [cls_id ], score * 100 )
221+ txt_color = (0 , 0 , 0 ) if numpy .mean (_COLORS [cls_id ]) > 128 else (255 , 255 , 255 )
222+ font = cv2 .FONT_HERSHEY_SIMPLEX
223+
224+ txt_size = cv2 .getTextSize (text , font , 0.4 , 1 )[0 ]
225+ cv2 .rectangle (img , (x0 , y0 ), (x1 , y1 ), color , 3 )
226+
227+ txt_bk_color = (_COLORS [cls_id ] * 0.7 ).astype (numpy .uint8 ).tolist ()
228+ cv2 .rectangle (
229+ img ,
230+ (x0 , y0 + 1 ),
231+ (x0 + txt_size [0 ] + 1 , y0 + int (1.5 * txt_size [1 ])),
232+ txt_bk_color ,
233+ - 1
234+ )
235+ cv2 .putText (img , text , (x0 , y0 + txt_size [1 ]), font , 0.4 , txt_color , thickness = 1 )
236+
237+ return img
0 commit comments