Fix mypy thinking output of SamPredictor.predict is a tensor.

dec4c129 · Eric Mintun · d4ecc68d · dec4c129
Commit dec4c129 authored Apr 12, 2023 by Eric Mintun
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 4 deletions

segment_anything/predictor.py segment_anything/predictor.py +4 -4

No files found.
--- a/segment_anything/predictor.py
+++ b/segment_anything/predictor.py
@@ -160,10 +160,10 @@ class SamPredictor:
            return_logits=return_logits,
        )
-        masks = masks[0].detach().cpu().numpy()
+        masks_np = masks[0].detach().cpu().numpy()
-        iou_predictions = iou_predictions[0].detach().cpu().numpy()
+        iou_predictions_np = iou_predictions[0].detach().cpu().numpy()
-        low_res_masks = low_res_masks[0].detach().cpu().numpy()
+        low_res_masks_np = low_res_masks[0].detach().cpu().numpy()
-        return masks, iou_predictions, low_res_masks
+        return masks_np, iou_predictions_np, low_res_masks_np
    @torch.no_grad()
    def predict_torch(