Fix NanoDet post process. Deal with different package versions.

PINTO0309 · Feb 16, 2022 · 1f9bc43 · 1f9bc43
1 parent 257b815
commit 1f9bc43
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/072_NanoDet/demo/demo_onnx.py b/072_NanoDet/demo/demo_onnx.py
@@ -259,9 +259,14 @@ def _get_bboxes_single(
 
         # Check the number of cases after NMS processing
         if len(indexes) > 0:
-            bboxes = bboxes[indexes[:, 0]]
-            scores = scores[indexes[:, 0]]
-            class_ids = class_ids[indexes[:, 0]]
+            if indexes.ndim == 2:
+                bboxes = bboxes[indexes[:, 0]]
+                scores = scores[indexes[:, 0]]
+                class_ids = class_ids[indexes[:, 0]]
+            elif indexes.ndim == 1:
+                bboxes = bboxes[indexes]
+                scores = scores[indexes]
+                class_ids = class_ids[indexes]
         else:
             bboxes = np.array([])
             scores = np.array([])

diff --git a/072_NanoDet/demo/demo_tflite.py b/072_NanoDet/demo/demo_tflite.py
@@ -280,9 +280,14 @@ def _get_bboxes_single(
 
         # Check the number of cases after NMS processing
         if len(indexes) > 0:
-            bboxes = bboxes[indexes[:, 0]]
-            scores = scores[indexes[:, 0]]
-            class_ids = class_ids[indexes[:, 0]]
+            if indexes.ndim == 2:
+                bboxes = bboxes[indexes[:, 0]]
+                scores = scores[indexes[:, 0]]
+                class_ids = class_ids[indexes[:, 0]]
+            elif indexes.ndim == 1:
+                bboxes = bboxes[indexes]
+                scores = scores[indexes]
+                class_ids = class_ids[indexes]
         else:
             bboxes = np.array([])
             scores = np.array([])