sam_adapter/model.py at main · CBI-PITT/sam_adapter · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
from __future__ import annotations

import logging
from threading import Lock
from typing import Any

import numpy as np

from .config import Settings

try:
    import torch
except Exception:  # pragma: no cover - optional import guard for local smoke tests
    torch = None  # type: ignore[assignment]


logger = logging.getLogger(__name__)


class ModelLoadError(RuntimeError):
    pass


class Sam2Model:
    def __init__(self, settings: Settings):
        self._settings = settings
        self._lock = Lock()
        self._predictor: Any | None = None
        self._device = self._resolve_device()
        self._backend = "uninitialized"

    @property
    def backend(self) -> str:
        return self._backend

    @property
    def device(self) -> str:
        return self._device

    def is_ready(self) -> bool:
        return self._predictor is not None

    def ensure_loaded(self) -> None:
        if self._predictor is not None:
            return
        with self._lock:
            if self._predictor is not None:
                return
            self._predictor = self._load_predictor()

    def predict(
        self,
        image_rgb: np.ndarray,
        point: tuple[int, int] | None,
        bbox: tuple[int, int, int, int] | None,
    ) -> np.ndarray:
        self.ensure_loaded()
        assert self._predictor is not None

        predictor = self._predictor
        predictor.set_image(image_rgb)
        kwargs: dict[str, Any] = {"multimask_output": True}
        if point is not None:
            kwargs["point_coords"] = np.asarray([[point[0], point[1]]], dtype=np.float32)
            kwargs["point_labels"] = np.asarray([1], dtype=np.int32)
        if bbox is not None:
            kwargs["box"] = np.asarray([bbox[0], bbox[1], bbox[2], bbox[3]], dtype=np.float32)

        masks, scores, _ = predictor.predict(**kwargs)
        best_idx = int(np.argmax(scores))
        return np.asarray(masks[best_idx], dtype=np.float32)

    def _resolve_device(self) -> str:
        if self._settings.force_cpu:
            return "cpu"
        if self._settings.device:
            return self._settings.device
        if torch is not None and torch.cuda.is_available():
            return "cuda"
        return "cpu"

    def _load_predictor(self):
        try:
            from sam2.build_sam import build_sam2
            from sam2.sam2_image_predictor import SAM2ImagePredictor
        except Exception as exc:  # pragma: no cover - dependency missing
            raise ModelLoadError(
                "Could not import the SAM2 runtime. Install the package from facebookresearch/sam2."
            ) from exc

        logger.info(
            "Loading SAM2 model config=%s checkpoint=%s device=%s",
            self._settings.model_config,
            self._settings.model_checkpoint,
            self._device,
        )
        model = build_sam2(self._settings.model_config, self._settings.model_checkpoint, device=self._device)
        predictor = SAM2ImagePredictor(model)
        self._backend = "sam2"
        logger.info("SAM2 model loaded successfully")
        return predictor