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

4
tests/test_cli.py
View file

@ -75,6 +75,7 @@ def test_fastsam(task='segment', model=WEIGHTS_DIR / 'FastSAM-s.pt', data='coco8
from ultralytics import FastSAM
from ultralytics.models.fastsam import FastSAMPrompt
from ultralytics.models.sam import Predictor
# Create a FastSAM model
sam_model = FastSAM(model) # or FastSAM-x.pt
@ -82,6 +83,9 @@ def test_fastsam(task='segment', model=WEIGHTS_DIR / 'FastSAM-s.pt', data='coco8
# Run inference on an image
everything_results = sam_model(source, device='cpu', retina_masks=True, imgsz=1024, conf=0.4, iou=0.9)
# Remove small regions
new_masks, _ = Predictor.remove_small_regions(everything_results[0].masks.data, min_area=20)
# Everything prompt
prompt_process = FastSAMPrompt(source, everything_results, device='cpu')
ann = prompt_process.everything_prompt()