Update handler.py

handler.py

@@ -6,9 +6,7 @@ from PIL import Image
 import io
 import json
 
-from sam2_model_stub import SAM2Hierarchical  # 👈 stub class we define separately
-
-# Define class labels (same order as training)
+# Define class labels (must match training order)
 CLASS_LABELS = [
     "glove_outline",
     "webbing",
@@ -19,17 +17,10 @@ CLASS_LABELS = [
 ]
 
 # ----------------------------
-# Load model weights + class
+# Load model directly from full .bin
 # ----------------------------
 def load_model():
-    model = SAM2Hierarchical(
-        num_classes=len(CLASS_LABELS),
-        in_channels=3,
-        backbone="vit_b",              # <-- match your config.yaml
-        freeze_backbone=True,
-        use_cls_head=True
-    )
-    model.load_state_dict(torch.load("pytorch_model.bin", map_location="cpu"))
+    model = torch.load("pytorch_model.bin", map_location="cpu")
     model.eval()
     return model
 
@@ -48,11 +39,33 @@ def preprocess(input_bytes):
     tensor = transform(image).unsqueeze(0)  # [1, 3, H, W]
     return tensor
 
+# ----------------------------
+# Dummy input wrapper
+# ----------------------------
+class DummyInput:
+    def __init__(self, image_tensor):
+        B, C, H, W = image_tensor.shape
+        self.images = image_tensor
+        self.masks = [torch.zeros(B, H, W, dtype=torch.bool)]
+        self.num_frames = 1
+        self.original_size = [(H, W)]
+        self.target_size = [(H, W)]
+        self.point_coords = [None]
+        self.point_labels = [None]
+        self.boxes = [None]
+        self.mask_inputs = torch.zeros(B, 1, H, W)
+        self.video_mask = torch.zeros(B, 1, H, W)
+        self.flat_obj_to_img_idx = [[0]]
+
 # ----------------------------
 # Postprocessing
 # ----------------------------
 def postprocess(output_tensor):
-    pred = torch.argmax(output_tensor, dim=1)[0].cpu().numpy()
+    if isinstance(output_tensor, dict) and "masks" in output_tensor:
+        logits = output_tensor["masks"]
+    else:
+        logits = output_tensor
+    pred = torch.argmax(logits, dim=1)[0].cpu().numpy()
     return pred.tolist()
 
 # ----------------------------
@@ -67,8 +80,10 @@ def infer(payload):
     else:
         raise ValueError("Unsupported input format")
 
+    input_obj = DummyInput(image_tensor)
+
     with torch.no_grad():
-        output = model(image_tensor)
+        output = model(input_obj)
 
     mask = postprocess(output)
     return {