HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Add SAM Predictor remove_small_regions test (#4576)

Browse source 
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Laughing 2023-08-26 19:07:15 +08:00 • committed by GitHub
parent b4dca690d4
commit e9f596430f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 10 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

13
ultralytics/models/sam/predict.py
View file

@ -374,6 +374,10 @@ class Predictor(BasePredictor):
masks (torch.Tensor): Masks, (N, H, W).
min_area (int): Minimum area threshold.
nms_thresh (float): NMS threshold.
Returns:
new_masks (torch.Tensor): New Masks, (N, H, W).
keep (List[int]): The indices of the new masks, which can be used to filter
the corresponding boxes.
"""
if len(masks) == 0:
return masks
@ -382,7 +386,7 @@ class Predictor(BasePredictor):
new_masks = []
scores = []
for mask in masks:
mask = mask.cpu().numpy()
mask = mask.cpu().numpy().astype(np.uint8)
mask, changed = remove_small_regions(mask, min_area, mode='holes')
unchanged = not changed
mask, changed = remove_small_regions(mask, min_area, mode='islands')
@ -402,9 +406,4 @@ class Predictor(BasePredictor):
nms_thresh,
)
# Only recalculate masks for masks that have changed
for i in keep:
if scores[i] == 0.0:
masks[i] = new_masks[i]
return masks[keep]
return new_masks[keep].to(device=masks.device, dtype=masks.dtype), keep