fix when no objects detected

3springs · Oct 14, 2020 · d41aa24 · d41aa24
1 parent 60e6c1b
commit d41aa24
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 397 deletions.
diff --git a/notebooks/c08_Object_Detection/Object_Detection.ipynb b/notebooks/c08_Object_Detection/Object_Detection.ipynb
diff --git a/notebooks/c08_Object_Detection/Object_Detection.py b/notebooks/c08_Object_Detection/Object_Detection.py
@@ -288,9 +288,10 @@ def get_prediction(img, threshold=0.5):
 
     # Get the Prediction Score
     pred_score = list(pred[0]["scores"].cpu().detach().numpy())
-    pred_t = [pred_score.index(x) for x in pred_score if x > threshold][
-        -1
-    ]  
+    pred_ts = [pred_score.index(x) for x in pred_score if x > threshold]
+    if len(pred_ts)==0:
+        return [], [], []
+    pred_t = pred_ts[-1]  
 
     # Get list of index with score greater than threshold.
     pred_boxes = pred_boxes[: pred_t + 1]
@@ -568,3 +569,5 @@ def instance_segmentation_api(img_path, threshold=0.5, rect_th=1, text_size=1):
 # -
 
 
+
+