35 lines
1005 B
Python
35 lines
1005 B
Python
from dataclasses import dataclass
|
|
from typing import Protocol
|
|
|
|
import numpy as np
|
|
|
|
|
|
@dataclass
|
|
class PredictResult:
|
|
"""Result of one bbox-prompted segmentation.
|
|
|
|
`polygon` is a list of (x, y) pixel coordinates describing the outer
|
|
contour of the largest mask blob. `score` is the model's confidence for
|
|
the selected mask (in [0, 1]).
|
|
"""
|
|
polygon: list[tuple[float, float]]
|
|
score: float
|
|
|
|
|
|
class BaseAdapter(Protocol):
|
|
"""Minimal adapter contract. Each model family implements this."""
|
|
|
|
def load(self, config: str, checkpoint: str, device: str) -> None:
|
|
"""Load weights into memory. Called once, lazily on first predict."""
|
|
...
|
|
|
|
def predict_bbox(
|
|
self, image_rgb: np.ndarray, bbox_xyxy: tuple[int, int, int, int]
|
|
) -> PredictResult:
|
|
"""Run segmentation with a bbox prompt, return the best polygon."""
|
|
...
|
|
|
|
def unload(self) -> None:
|
|
"""Free GPU memory. Called when the registry evicts this model."""
|
|
...
|