Create handler.py

handler.py

@@ -0,0 +1,64 @@
+# handler.py
+
+import torch
+import torchvision.transforms as T
+from PIL import Image
+import io
+import json
+
+# Define class labels (same order as training)
+CLASS_LABELS = [
+    "glove_outline",
+    "webbing",
+    "thumb",
+    "palm_pocket",
+    "hand",
+    "glove_exterior"
+]
+
+# Load model from disk
+def load_model():
+    model = torch.load("pytorch_model.bin", map_location="cpu")
+    model.eval()
+    return model
+
+model = load_model()
+
+# Preprocessing transform
+transform = T.Compose([
+    T.Resize((720, 1280)),   # or whatever input size model expects
+    T.ToTensor()
+])
+
+# Input: raw image bytes
+def preprocess(input_bytes):
+    image = Image.open(io.BytesIO(input_bytes)).convert("RGB")
+    tensor = transform(image).unsqueeze(0)  # [1, 3, H, W]
+    return tensor
+
+# Postprocess output: convert logits to mask
+def postprocess(output_tensor):
+    # Argmax over channel dimension (assumes shape [1, C, H, W])
+    pred = torch.argmax(output_tensor, dim=1)[0].cpu().numpy()
+    return pred.tolist()  # List of H x W values from 0 to 5
+
+# TorchServe/HF entrypoint
+def infer(payload):
+    # If input is multipart/form-data, raw bytes
+    if isinstance(payload, bytes):
+        image_tensor = preprocess(payload)
+    elif isinstance(payload, dict) and "inputs" in payload:
+        # Hugging Face Inference API passes {"inputs": "base64 image data"}
+        from base64 import b64decode
+        image_tensor = preprocess(b64decode(payload["inputs"]))
+    else:
+        raise ValueError("Unsupported input format")
+
+    with torch.no_grad():
+        output = model(image_tensor)
+
+    mask = postprocess(output)
+    return {
+        "mask": mask,
+        "classes": CLASS_LABELS
+    }