Create medsam2_model.py

medsam2_model.py

@@ -0,0 +1,32 @@
+import torch
+import numpy as np
+from skimage import transform
+# from sam2_train.build_sam import build_sam2
+# from sam2_train.sam2_image_predictor import SAM2ImagePredictor
+from sam2.build_sam import build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+class MedSAM2:
+    def __init__(self, model_path, device="cpu"):
+        self.device = device
+        self.model = build_sam2("sam2_hiera_t", model_path, device=device)
+        self.predictor = SAM2ImagePredictor(self.model)
+
+    def predict(self, image: np.ndarray, box: list[float]) -> np.ndarray:
+        image_3c = image if image.shape[2] == 3 else np.repeat(image[:, :, None], 3, axis=-1)
+        img_1024 = transform.resize(image_3c, (1024, 1024), preserve_range=True).astype(np.uint8)
+
+        box_np = np.array(box)
+        box_1024 = box_np / np.array([image.shape[1], image.shape[0], image.shape[1], image.shape[0]]) * 1024
+        box_1024 = box_1024[None, :]
+
+        with torch.inference_mode(), torch.autocast(self.device, dtype=torch.bfloat16):
+            self.predictor.set_image(img_1024)
+            masks, _, _ = self.predictor.predict(
+                point_coords=None,
+                point_labels=None,
+                box=box_1024,
+                multimask_output=False
+            )
+
+        return masks[0].astype(np.uint8)