Skip to content

Commit 28b6e46

Browse files
Simplify Mask Handling When Only One Object is Detected (#352)
* Fix empty mask overlay handling when no objects are detected. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * mask overlay corrected * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Handle case for single objects detected. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1 parent 7d1d648 commit 28b6e46

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

samgeo/text_sam.py

+4-1
Original file line numberDiff line numberDiff line change
@@ -337,7 +337,10 @@ def predict(
337337
masks = torch.tensor([])
338338
if len(boxes) > 0:
339339
masks = self.predict_sam(image_pil, boxes)
340-
if 1 in masks.shape:
340+
# If masks have 4 dimensions and the second dimension is 1 (e.g., [boxes, 1, height, width]),
341+
# squeeze that dimension to reduce it to 3 dimensions ([boxes, height, width]).
342+
# If boxes = 1, the mask's shape will be [1, height, width] after squeezing.
343+
if masks.ndim == 4 and masks.shape[1] == 1:
341344
masks = masks.squeeze(1)
342345

343346
if boxes.nelement() == 0: # No "object" instances found

0 commit comments

Comments
 (0)