Fix for warning raised if model outputs no box predictions (microsoft…

…#988) * Fix for warning when no predictions are made * Same thing for boxes Co-authored-by: Adam J. Stewart <[email protected]>
yichiac · Dec 30, 2022 · b7df83e · b7df83e
1 parent 8a16a65
commit b7df83e
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/torchgeo/datasets/nasa_marine_debris.py b/torchgeo/datasets/nasa_marine_debris.py
@@ -233,10 +233,12 @@ def plot(
         """
         ncols = 1
 
-        image = draw_bounding_boxes(image=sample["image"], boxes=sample["boxes"])
+        image = sample["image"]
+        if "boxes" in sample and len(sample["boxes"]):
+            image = draw_bounding_boxes(image=sample["image"], boxes=sample["boxes"])
         image = image.permute((1, 2, 0)).numpy()
 
-        if "prediction_boxes" in sample:
+        if "prediction_boxes" in sample and len(sample["prediction_boxes"]):
             ncols += 1
             preds = draw_bounding_boxes(
                 image=sample["image"], boxes=sample["prediction_boxes"]