Update handler.py

handler.py

@@ -1,13 +1,16 @@
+import os
 import torch
 from torchvision import transforms
 from PIL import Image
 import io
 
-MODEL_PATH = "model.pt"
-NUM_CLASSES = 4  
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))  
+MODEL_FILENAME = "model.pt"
+MODEL_PATH = os.path.join(BASE_DIR, MODEL_FILENAME)
+
+NUM_CLASSES = 4
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load Faster R-CNN model
 def load_model(model_path, num_classes):
     from torchvision.models.detection import fasterrcnn_resnet50_fpn
     model = fasterrcnn_resnet50_fpn(pretrained=False, num_classes=num_classes)
@@ -17,13 +20,13 @@ def load_model(model_path, num_classes):
     model.eval()
     return model
 
+model = load_model(MODEL_PATH, NUM_CLASSES)
+
 transform = transforms.Compose([
     transforms.Resize((640, 640)),
     transforms.ToTensor(),
 ])
 
-model = load_model(MODEL_PATH, NUM_CLASSES)
-
 def detect_objects(image_bytes):
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
     input_tensor = transform(image).unsqueeze(0).to(DEVICE)
@@ -45,12 +48,13 @@ def detect_objects(image_bytes):
     return {"predictions": results}
 
 def inference(payload):
+    import base64
     try:
         if "image" not in payload:
-            return {"error": "No image provided. Please send an image."}
-        
-        image_bytes = payload["image"].encode("latin1")  
-        
+            return {"error": "No image provided. Please send a Base64-encoded image."}
+
+        image_bytes = base64.b64decode(payload["image"])
+
         results = detect_objects(image_bytes)
         return results
     except Exception as e: