@@ -273,21 +273,24 @@ def draw(
273273 Returns:
274274 Output image with annotations.
275275 """
276- for detection in predictions .objects :
277- class_id = int (detection .id )
278- color = self .color_palette [class_id ]
279- det_label = self .color_palette [class_id ] if self .labels and len (self .labels ) >= class_id else f"#{ class_id } "
280- xmin , ymin , xmax , ymax = detection .xmin , detection .ymin , detection .xmax , detection .ymax
281- cv2 .rectangle (frame , (xmin , ymin ), (xmax , ymax ), color , 2 )
282- cv2 .putText (
283- frame ,
284- f"{ det_label } { detection .score :.1%} " ,
285- (xmin , ymin - 7 ),
286- cv2 .FONT_HERSHEY_COMPLEX ,
287- 0.6 ,
288- color ,
289- 1 ,
290- )
276+ if len (predictions .bboxes .shape ):
277+ for i , box in enumerate (predictions .bboxes ):
278+ class_id = int (predictions .labels [i ])
279+ color = self .color_palette [class_id ]
280+ det_label = (
281+ self .color_palette [class_id ] if self .labels and len (self .labels ) >= class_id else f"#{ class_id } "
282+ )
283+ xmin , ymin , xmax , ymax = box
284+ cv2 .rectangle (frame , (xmin , ymin ), (xmax , ymax ), color , 2 )
285+ cv2 .putText (
286+ frame ,
287+ f"{ det_label } { predictions .scores [i ]:.1%} " ,
288+ (xmin , ymin - 7 ),
289+ cv2 .FONT_HERSHEY_COMPLEX ,
290+ 0.6 ,
291+ color ,
292+ 1 ,
293+ )
291294
292295 return frame
293296
@@ -339,16 +342,10 @@ def draw(
339342 np.ndarray - The input frame with the instance segmentation results drawn on it.
340343 """
341344 result = frame .copy ()
342- output_objects = predictions .segmentedObjects
343- bboxes = [[output .xmin , output .ymin , output .xmax , output .ymax ] for output in output_objects ]
344- scores = [output .score for output in output_objects ]
345- masks = [output .mask for output in output_objects ]
346- label_names = [output .str_label for output in output_objects ]
347-
348- result = self ._overlay_masks (result , masks )
349- return self ._overlay_labels (result , bboxes , label_names , scores )
345+ result = self ._overlay_masks (result , predictions .masks )
346+ return self ._overlay_labels (result , predictions .bboxes , predictions .label_names , predictions .scores )
350347
351- def _overlay_masks (self , image : np .ndarray , masks : list [ np .ndarray ] ) -> np .ndarray :
348+ def _overlay_masks (self , image : np .ndarray , masks : np .ndarray ) -> np .ndarray :
352349 segments_image = image .copy ()
353350 aggregated_mask = np .zeros (image .shape [:2 ], dtype = np .uint8 )
354351 aggregated_colored_mask = np .zeros (image .shape , dtype = np .uint8 )
@@ -381,9 +378,9 @@ def _overlay_boxes(self, image: np.ndarray, boxes: list[np.ndarray], classes: li
381378 def _overlay_labels (
382379 self ,
383380 image : np .ndarray ,
384- boxes : list [ np .ndarray ] ,
381+ boxes : np .ndarray ,
385382 classes : list [str ],
386- scores : list [ float ] ,
383+ scores : np . ndarray ,
387384 ) -> np .ndarray :
388385 template = "{}: {:.2f}" if self .show_scores else "{}"
389386
0 commit comments