Upload folder using huggingface_hub

inference.py

@@ -1,4 +1,5 @@
-from typing import Dict
+import base64
+import io
 import torch
 import numpy as np
 from PIL import Image
@@ -6,34 +7,60 @@ 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(
+# Инициализация модели один раз
+device = "cuda" if torch.cuda.is_available() else "cpu"
+torch.autocast(device_type=device, dtype=torch.bfloat16).__enter__()
+if torch.cuda.is_available() and torch.cuda.get_device_properties(0).major >= 8:
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+# Загрузи путь к своей модели
+medsam2_model = build_sam2('sam2_hiera_t', 'MedSAM2_pretrain_10ep_b1_AMD-SD_sam2_hiera_t.pth', device=device)
+predictor = SAM2ImagePredictor(medsam2_model)
+
+# --- Вспомогательные функции ---
+def decode_image(base64_str):
+    img_bytes = base64.b64decode(base64_str)
+    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+    return np.array(img)
+
+def encode_mask_to_base64(mask_np):
+    mask_img = Image.fromarray((mask_np * 255).astype(np.uint8))
+    buffer = io.BytesIO()
+    mask_img.save(buffer, format="PNG")
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+# --- Главная функция вызова инференса ---
+def predict(input_dict):
+    try:
+        image_b64 = input_dict["image"]
+        box = input_dict["box"]  # [x1, y1, x2, y2]
+
+        image = decode_image(image_b64)
+
+        # Resize image to 1024x1024 for MedSAM2
+        img_3c = np.repeat(image[:, :, None], 3, axis=-1) if len(image.shape) == 2 else image
+        img_1024 = transform.resize(img_3c, (1024, 1024), order=3, preserve_range=True, anti_aliasing=True).astype(np.uint8)
+
+        # Normalize box
+        box_np = np.array(box)
+        box_1024 = box_np / np.array([img_3c.shape[1], img_3c.shape[0], img_3c.shape[1], img_3c.shape[0]]) * 1024
+
+        with torch.inference_mode(), torch.autocast(device, dtype=torch.bfloat16):
+            predictor.set_image(img_1024)
+            masks, _, _ = predictor.predict(
                 point_coords=None,
                 point_labels=None,
-                box=box_1024,
-                multimask_output=False
+                box=np.array([box_1024]),
+                multimask_output=False,
             )
 
-        mask = masks[0].astype(np.uint8)
-        return {"mask": mask.tolist()}
+        mask = masks[0]  # (H, W) np.uint8
+
+        return {
+            "mask": encode_mask_to_base64(mask.astype(np.uint8)),
+            "shape": mask.shape,
+        }
+
+    except Exception as e:
+        return {"error": str(e)}