Fixed single bbox bug (#184)

opengeos · Aug 30, 2023 · 738f235 · 738f235
1 parent f66b3e4
commit 738f235
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/samgeo/hq_sam.py b/samgeo/hq_sam.py
@@ -573,7 +573,9 @@ def predict(
 
         self.boxes = input_boxes
 
-        if boxes is None or (not isinstance(boxes[0], list)):
+        if boxes is None or (len(boxes) == 1) or (len(boxes) == 4 and isinstance(boxes[0], float)):
+            if isinstance(boxes, list) and isinstance(boxes[0], list):
+                boxes = boxes[0]
             masks, scores, logits = predictor.predict(
                 point_coords,
                 point_labels,

diff --git a/samgeo/samgeo.py b/samgeo/samgeo.py
@@ -570,7 +570,9 @@ def predict(
 
         self.boxes = input_boxes
 
-        if boxes is None or (not isinstance(boxes[0], list)):
+        if boxes is None or (len(boxes) == 1) or (len(boxes) == 4 and isinstance(boxes[0], float)):
+            if isinstance(boxes, list) and isinstance(boxes[0], list):
+                boxes = boxes[0]
             masks, scores, logits = predictor.predict(
                 point_coords,
                 point_labels,