Update model.py

model.py

@@ -1,24 +1,35 @@
 import os
 from typing import List, Dict, Optional
-from uuid import uuid4
 
 from label_studio_converter import brush
 from label_studio_ml.model import LabelStudioMLBase
+from uuid import uuid4
 
 from sam_predictor import SAMPredictor
 
-SAM_CHOICE = os.getenv("SAM_CHOICE", "MobileSAM")
+SAM_CHOICE = os.environ.get("SAM_CHOICE", "MobileSAM")
 PREDICTOR = SAMPredictor(SAM_CHOICE)
 
 
 class SamMLBackend(LabelStudioMLBase):
+    def __init__(self, project_id=None, label_config=None, **kwargs):
+        # Make sure model_dir always exists, even if the backend package
+        # does not initialize it correctly.
+        self.model_dir = os.environ.get("MODEL_DIR", "/tmp/mlbackend")
+        os.makedirs(self.model_dir, exist_ok=True)
+        super().__init__(project_id=project_id, label_config=label_config)
+
     def setup(self):
-        # Explicitly mark the model as initialized for the current backend API
+        # Mark the model as initialized
         self.set("model_version", f"{SAM_CHOICE}-v1")
 
     def predict(self, tasks: List[Dict], context: Optional[Dict] = None, **kwargs) -> List[Dict]:
-        # Newer backend API: parsed config is exposed via self.label_interface
-        from_name, to_name, value = self.label_interface.get_first_tag_occurence("BrushLabels", "Image")
+        # Hard-code these to match your current Label Studio XML:
+        # <BrushLabels name="tag" toName="image">
+        # <Image name="image" value="$image" ... />
+        from_name = "tag"
+        to_name = "image"
+        value = "image"
 
         if not context or not context.get("result"):
             return []
@@ -32,11 +43,12 @@ class SamMLBackend(LabelStudioMLBase):
         selected_label = None
 
         for ctx in context["result"]:
-            x = ctx["value"]["x"] * image_width / 100.0
-            y = ctx["value"]["y"] * image_height / 100.0
             ctx_type = ctx["type"]
             selected_label = ctx["value"][ctx_type][0]
 
+            x = ctx["value"]["x"] * image_width / 100.0
+            y = ctx["value"]["y"] * image_height / 100.0
+
             if ctx_type == "keypointlabels":
                 point_labels.append(int(ctx["is_positive"]))
                 point_coords.append([int(x), int(y)])
@@ -55,20 +67,8 @@ class SamMLBackend(LabelStudioMLBase):
             input_box=input_box,
         )
 
-        return self.get_results(
-            masks=predictor_results["masks"],
-            probs=predictor_results["probs"],
-            width=image_width,
-            height=image_height,
-            from_name=from_name,
-            to_name=to_name,
-            label=selected_label,
-        )
-
-    def get_results(self, masks, probs, width, height, from_name, to_name, label):
         results = []
-
-        for mask, prob in zip(masks, probs):
+        for mask, prob in zip(predictor_results["masks"], predictor_results["probs"]):
             label_id = str(uuid4())[:8]
             mask = (mask * 255).astype("uint8")
             rle = brush.mask2rle(mask)
@@ -77,13 +77,13 @@ class SamMLBackend(LabelStudioMLBase):
                 "id": label_id,
                 "from_name": from_name,
                 "to_name": to_name,
-                "original_width": width,
-                "original_height": height,
+                "original_width": image_width,
+                "original_height": image_height,
                 "image_rotation": 0,
                 "value": {
                     "format": "rle",
                     "rle": rle,
-                    "brushlabels": [label],
+                    "brushlabels": [selected_label],
                 },
                 "score": float(prob),
                 "type": "brushlabels",