Run a single prediction on the model
(
self,
input_image: Path = Input(description="Input image"),
model_name: str = Input(
description="choose a model",
choices=["FastSAM-x", "FastSAM-s"],
default="FastSAM-x",
),
iou: float = Input(
description="iou threshold for filtering the annotations", default=0.7
),
text_prompt: str = Input(
description='use text prompt eg: "a black dog"', default=None
),
conf: float = Input(description="object confidence threshold", default=0.25),
retina: bool = Input(
description="draw high-resolution segmentation masks", default=True
),
box_prompt: str = Input(default="[0,0,0,0]", description="[x,y,w,h]"),
point_prompt: str = Input(default="[[0,0]]", description="[[x1,y1],[x2,y2]]"),
point_label: str = Input(default="[0]", description="[1,0] 0:background, 1:foreground"),
withContours: bool = Input(
description="draw the edges of the masks", default=False
),
better_quality: bool = Input(
description="better quality using morphologyEx", default=False
),
)
| 17 | self.models = {k: YOLO(f"{k}.pt") for k in ["FastSAM-s", "FastSAM-x"]} |
| 18 | |
| 19 | def predict( |
| 20 | self, |
| 21 | input_image: Path = Input(description="Input image"), |
| 22 | model_name: str = Input( |
| 23 | description="choose a model", |
| 24 | choices=["FastSAM-x", "FastSAM-s"], |
| 25 | default="FastSAM-x", |
| 26 | ), |
| 27 | iou: float = Input( |
| 28 | description="iou threshold for filtering the annotations", default=0.7 |
| 29 | ), |
| 30 | text_prompt: str = Input( |
| 31 | description='use text prompt eg: "a black dog"', default=None |
| 32 | ), |
| 33 | conf: float = Input(description="object confidence threshold", default=0.25), |
| 34 | retina: bool = Input( |
| 35 | description="draw high-resolution segmentation masks", default=True |
| 36 | ), |
| 37 | box_prompt: str = Input(default="[0,0,0,0]", description="[x,y,w,h]"), |
| 38 | point_prompt: str = Input(default="[[0,0]]", description="[[x1,y1],[x2,y2]]"), |
| 39 | point_label: str = Input(default="[0]", description="[1,0] 0:background, 1:foreground"), |
| 40 | withContours: bool = Input( |
| 41 | description="draw the edges of the masks", default=False |
| 42 | ), |
| 43 | better_quality: bool = Input( |
| 44 | description="better quality using morphologyEx", default=False |
| 45 | ), |
| 46 | ) -> Path: |
| 47 | """Run a single prediction on the model""" |
| 48 | |
| 49 | # default params |
| 50 | |
| 51 | out_path = "output" |
| 52 | if os.path.exists(out_path): |
| 53 | shutil.rmtree(out_path) |
| 54 | os.makedirs(out_path, exist_ok=True) |
| 55 | |
| 56 | device = torch.device( |
| 57 | "cuda" |
| 58 | if torch.cuda.is_available() |
| 59 | else "mps" |
| 60 | if torch.backends.mps.is_available() |
| 61 | else "cpu" |
| 62 | ) |
| 63 | |
| 64 | args = argparse.Namespace( |
| 65 | better_quality=better_quality, |
| 66 | box_prompt=box_prompt, |
| 67 | conf=conf, |
| 68 | device=device, |
| 69 | img_path=str(input_image), |
| 70 | imgsz=1024, |
| 71 | iou=iou, |
| 72 | model_path="FastSAM-x.pt", |
| 73 | output=out_path, |
| 74 | point_label=point_label, |
| 75 | point_prompt=point_prompt, |
| 76 | randomcolor=True, |
nothing calls this directly
no test coverage detected