Update handler.py

handler.py

@@ -1,9 +1,9 @@
-import os
 import torch
 from torchvision import transforms
 from PIL import Image
 import io
 
+# Load the Faster R-CNN model
 from model import get_model
 
 class EndpointHandler:
@@ -12,49 +12,54 @@ class EndpointHandler:
         Initialize the handler. Load the Faster R-CNN model.
         """
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model_weights_path = os.path.join(path, "model.pt") 
-        
-        # Load the model
+        self.model_weights_path = os.path.join(path, "model.pt")  # Adjust path
+
+        # Load model
         self.model = get_model(num_classes=4) 
-        print(f"Loading weights from: {self.model_weights_path}")
         checkpoint = torch.load(self.model_weights_path, map_location=self.device)
         self.model.load_state_dict(checkpoint["model_state_dict"])
         self.model.to(self.device)
         self.model.eval()
 
-        # Define image preprocessing
+        # Image preprocessing
         self.transform = transforms.Compose([
-            transforms.Resize((640, 640)),
+            transforms.Resize((640, 640)), 
             transforms.ToTensor(),
         ])
 
     def __call__(self, data):
         """
-        Process the incoming request and return object detection predictions.
+        Process incoming binary image data and return object detection results.
         """
         try:
-            if "image" not in data:
-                return [{"error": "No 'image' provided in request."}]
+            # Read raw binary data (image file)
+            image_bytes = data.get("body", b"") 
+            if not image_bytes:
+                return {"error": "No image data provided in request."}
 
-            image_bytes = data["image"].encode("latin1")
+            
             image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
+            
             input_tensor = self.transform(image).unsqueeze(0).to(self.device)
 
+
             with torch.no_grad():
-                outputs = self.model(input_tensor)
+                predictions = self.model(input_tensor)
 
-            boxes = outputs[0]["boxes"].cpu().tolist()
-            labels = outputs[0]["labels"].cpu().tolist()
-            scores = outputs[0]["scores"].cpu().tolist()
+            
+            boxes = predictions[0]["boxes"].cpu().tolist()
+            labels = predictions[0]["labels"].cpu().tolist()
+            scores = predictions[0]["scores"].cpu().tolist()
 
+            
             threshold = 0.5
-            predictions = [
+            results = [
                 {"box": box, "label": label, "score": score}
                 for box, label, score in zip(boxes, labels, scores)
                 if score > threshold
             ]
 
-            return [{"predictions": predictions}]
+            return {"predictions": results}
         except Exception as e:
-            return [{"error": str(e)}]
+            return {"error": str(e)}