Update handler.py

handler.py

@@ -38,37 +38,39 @@ class EndpointHandler:
         return image_tensor
 
     def __call__(self, data):
-        """
-        Process incoming raw binary image data.
-        """
-        try:
-            if "body" not in data:
-                return {"error": "No image data provided in request."}
+    """
+    Process incoming raw binary image data.
+    """
+    try:
+        if "body" not in data:
+            return {"error": "No image data provided in request."}
 
-            image_bytes = data["body"]  
-            image_tensor = self.preprocess_frame(image_bytes)
+        image_bytes = data["body"] 
+        image_tensor = self.preprocess_frame(image_bytes)
 
-            # Perform inference
-            with torch.no_grad():
-                predictions = self.model(image_tensor)
+        with torch.no_grad():
+            predictions = self.model(image_tensor)
 
-            # Extract predictions
-            boxes = predictions[0]["boxes"].cpu().tolist()
-            labels = predictions[0]["labels"].cpu().tolist()
-            scores = predictions[0]["scores"].cpu().tolist()
+        boxes = predictions[0]["boxes"].cpu().tolist()
+        labels = predictions[0]["labels"].cpu().tolist()
+        scores = predictions[0]["scores"].cpu().tolist()
 
-            # Filter predictions by confidence threshold
-            results = []
-            for box, label, score in zip(boxes, labels, scores):
-                if score >= CONFIDENCE_THRESHOLD:
-                    x1, y1, x2, y2 = map(int, box)
-                    label_text = self.label_map.get(label, "unknown")
-                    results.append({
-                        "box": [x1, y1, x2, y2],
-                        "label": label_text,
-                        "score": round(score, 2)
-                    })
+        results = []
+        for box, label, score in zip(boxes, labels, scores):
+            if score >= CONFIDENCE_THRESHOLD:
+                x1, y1, x2, y2 = map(int, box) 
+                label_text = self.label_map.get(label, "unknown") 
+                results.append({
+                    "label": label_text,
+                    "score": round(score, 2), 
+                    "box": {
+                        "xmin": x1,
+                        "ymin": y1,
+                        "xmax": x2,
+                        "ymax": y2
+                    }
+                })
 
-            return {"predictions": results}
-        except Exception as e:
-            return {"error": str(e)}
+        return results 
+    except Exception as e:
+        return {"error": str(e)}