Upload folder using huggingface_hub

MedSAM2_pretrain_10ep_b1_AMD-SD_sam2_hiera_t.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8018821b159767012d081db1f5cef2720aa495bad6861b8d6b1c63dce24a5a6e
+size 155972738

README.md

@@ -1,3 +0,0 @@
----
-license: apache-2.0
----

config.json

@@ -0,0 +1,5 @@
+{
+    "pipeline_tag": "image-segmentation",
+    "base_model": "medsam2"
+  }
+  

inference.py

@@ -0,0 +1,39 @@
+from typing import Dict
+import torch
+import numpy as np
+from PIL import Image
+from skimage import transform
+from sam2.build_sam import build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+class PreTrainedModel:
+    def __init__(self):
+        self.model = build_sam2(
+            "sam2_hiera_t",
+            "MedSAM2_pretrain_10ep_b1_AMD-SD_sam2_hiera_t.pth",
+            device="cuda" if torch.cuda.is_available() else "cpu"
+        )
+        self.predictor = SAM2ImagePredictor(self.model)
+
+    def __call__(self, inputs: Dict):
+        image = Image.open(inputs["image"]).convert("RGB")
+        box = list(map(float, inputs["box"]))
+
+        image_np = np.array(image)
+        img_3c = image_np if image_np.shape[2] == 3 else np.repeat(image_np[:, :, None], 3, axis=-1)
+        img_1024 = transform.resize(img_3c, (1024, 1024), preserve_range=True).astype(np.uint8)
+
+        box_1024 = np.array(box) / [image_np.shape[1], image_np.shape[0], image_np.shape[1], image_np.shape[0]] * 1024
+        box_1024 = box_1024[None, :]
+
+        with torch.inference_mode(), torch.autocast("cuda" if torch.cuda.is_available() else "cpu", 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
+            )
+
+        mask = masks[0].astype(np.uint8)
+        return {"mask": mask.tolist()}

requirements.txt

@@ -0,0 +1,6 @@
+torch
+numpy
+scikit-image
+gradio
+pillow
+git+https://github.com/facebookresearch/sam2.git