zjt/SAMapi
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Eric Mintun 2023-04-12 09:48:07 -07:00
parent d4ecc68dea
commit dec4c12940
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
segment_anything/predictor.py
View File

@ -160,10 +160,10 @@ class SamPredictor:
return_logits=return_logits,
)
masks = masks[0].detach().cpu().numpy()
iou_predictions = iou_predictions[0].detach().cpu().numpy()
low_res_masks = low_res_masks[0].detach().cpu().numpy()
return masks, iou_predictions, low_res_masks
masks_np = masks[0].detach().cpu().numpy()
iou_predictions_np = iou_predictions[0].detach().cpu().numpy()
low_res_masks_np = low_res_masks[0].detach().cpu().numpy()
return masks_np, iou_predictions_np, low_res_masks_np
@torch.no_grad()
def predict_torch(