fix no yolox detections

mikel-brostrom · Jul 11, 2023 · 02a2c48 · 02a2c48
1 parent 43e6b48
commit 02a2c48
Showing 1 changed file with 22 additions and 19 deletions.
diff --git a/examples/multi_yolo_backend.py b/examples/multi_yolo_backend.py
@@ -110,26 +110,29 @@ def __call__(self, im, im0s):
                 nms_thre=0.45, class_agnostic=True
             )[0]
 
-            # (x, y, x, y, conf, obj, cls) --> (x, y, x, y, conf, cls)
-            preds[:, 4] = preds[:, 4] * preds[:, 5]
-            preds = preds[:, [0, 1, 2, 3, 4, 6]]
-
-            # calculate factor for predictions
-            im0_w = im0s[0].shape[1]
-            im0_h = im0s[0].shape[0]
-            im_w = im[0].shape[2]
-            im_h = im[0].shape[1]
-            w_r = im0_w / im_w
-            h_r = im0_h / im_h
-
-            # scale to original image
-            preds[:, [0, 2]] = preds[:, [0, 2]] * w_r
-            preds[:, [1, 3]] = preds[:, [1, 3]] * h_r
+            if preds is None:
+                return torch.empty(0, 6)
+            else:
+                # (x, y, x, y, conf, obj, cls) --> (x, y, x, y, conf, cls)
+                preds[:, 4] = preds[:, 4] * preds[:, 5]
+                preds = preds[:, [0, 1, 2, 3, 4, 6]]
 
-            preds = torch.clip(preds, min=0)
+                # calculate factor for predictions
+                im0_w = im0s[0].shape[1]
+                im0_h = im0s[0].shape[0]
+                im_w = im[0].shape[2]
+                im_h = im[0].shape[1]
+                w_r = im0_w / im_w
+                h_r = im0_h / im_h
 
-            if self.args.classes:  # Filter boxes by classes
-                preds = preds[torch.isin(preds[:, 5].cpu(), torch.as_tensor(self.args.classes))]
+                # scale to original image
+                preds[:, [0, 2]] = preds[:, [0, 2]] * w_r
+                preds[:, [1, 3]] = preds[:, [1, 3]] * h_r
+
+                preds = torch.clip(preds, min=0)
+
+                if self.args.classes:  # Filter boxes by classes
+                    preds = preds[torch.isin(preds[:, 5].cpu(), torch.as_tensor(self.args.classes))]
 
         elif 'yolov8' in self.model_name:
             preds = self.model(
@@ -183,4 +186,4 @@ def postprocess(self, path, preds, im, im0s, predictor):
 if __name__ == "__main__":
     yolo = MultiYolo(model='YOLO_NAS_S', device='cuda:0')
     rgb = np.random.randint(255, size=(640, 640, 3), dtype=np.uint8)
-    yolo(rgb, rgb)
+    yolo(rgb, rgb)