fix multiclass_nms bug (#66)

demo/ONNXRuntime/onnx_inference.py

@@ -84,10 +84,11 @@ if __name__ == '__main__':
     boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
     boxes_xyxy /= ratio
     dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.65, score_thr=0.1)
-
-    final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
-    origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,
-                     conf=args.score_thr, class_names=COCO_CLASSES)
+    
+    if dets is not None:
+        final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
+        origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,
+                         conf=args.score_thr, class_names=COCO_CLASSES)
 
     mkdir(args.output_dir)
     output_path = os.path.join(args.output_dir, args.image_path.split("/")[-1])

demo/OpenVINO/python/openvino_inference.py

@@ -142,11 +142,12 @@ def main():
     boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
     boxes_xyxy /= ratio
     dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.65, score_thr=0.1)
-
-    final_boxes = dets[:, :4]
-    final_scores, final_cls_inds = dets[:, 4], dets[:, 5]
-    origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,
-                     conf=args.score_thr, class_names=COCO_CLASSES)
+    
+    if dets is not None:
+        final_boxes = dets[:, :4]
+        final_scores, final_cls_inds = dets[:, 4], dets[:, 5]
+        origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,
+                         conf=args.score_thr, class_names=COCO_CLASSES)
 
     mkdir(args.output_dir)
     output_path = os.path.join(args.output_dir, args.image_path.split("/")[-1])

yolox/utils/demo_utils.py

@@ -61,6 +61,8 @@ def multiclass_nms(boxes, scores, nms_thr, score_thr):
                 cls_inds = np.ones((len(keep), 1)) * cls_ind
                 dets = np.concatenate([valid_boxes[keep], valid_scores[keep, None], cls_inds], 1)
                 final_dets.append(dets)
+    if len(final_dets) == 0:
+        return None
     return np.concatenate(final_dets, 0)