Update handler.py

handler.py

@@ -6,29 +6,29 @@ import numpy as np
 # Здесь предполагается, что у вас есть функция segment(image: PIL.Image) -> np.ndarray (маска)
 from medsam2_model import MedSAM2
 
-# Загрузка модели
-model = MedSAM2("MedSAM2_pretrain_10ep_b1_AMD-SD_sam2_hiera_t.pth")
-
 def image_to_base64(image: Image.Image) -> str:
     buffered = io.BytesIO()
     image.save(buffered, format="PNG")
     return "data:image/png;base64," + base64.b64encode(buffered.getvalue()).decode()
 
-def predict(payload):
-    # CVAT будет отправлять изображение в виде base64
-    if isinstance(payload, dict) and "image" in payload:
-        image_data = payload["image"]
-        if image_data.startswith("data:image"):
-            header, base64_data = image_data.split(",", 1)
-            image = Image.open(io.BytesIO(base64.b64decode(base64_data)))
-
-            # Получаем маску
-            mask_array = model.predict(image_np, box)  # Предполагается бинарная маска (0 и 1)
-            mask_pil = Image.fromarray((mask_array * 255).astype(np.uint8))
-
-            return [{
-                "label": "mock-segmentation",
-                "mask": image_to_base64(mask_pil),
-                "score": 0.99
-            }]
-    return [{"label": "mock-segmentation", "mask": None, "score": 0.0}]
+class EndpointHandler():
+    def __init__(self, path=""):
+        model = MedSAM2("MedSAM2_pretrain_10ep_b1_AMD-SD_sam2_hiera_t.pth")
+    
+   def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+        if isinstance(data, dict) and "image" in data:
+            image_data = data["image"]
+            if image_data.startswith("data:image"):
+                header, base64_data = image_data.split(",", 1)
+                image = Image.open(io.BytesIO(base64.b64decode(base64_data)))
+    
+                # Получаем маску
+                mask_array = model.predict(image_np, box)  # Предполагается бинарная маска (0 и 1)
+                mask_pil = Image.fromarray((mask_array * 255).astype(np.uint8))
+    
+                return [{
+                    "label": "mock-segmentation",
+                    "mask": image_to_base64(mask_pil),
+                    "score": 0.99
+                }]
+        return [{"label": "mock-segmentation", "mask": None, "score": 0.0}]