Update handler.py

handler.py

@@ -1,64 +1,72 @@
 import torch
-from torchvision import transforms
+from model import get_model 
+from torchvision.transforms import ToTensor
 from PIL import Image
 import io
-import os
 
-from model import get_model
+# Constants
+NUM_CLASSES = 4
+CONFIDENCE_THRESHOLD = 0.5
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 class EndpointHandler:
     def __init__(self, path: str = ""):
         """
-        Initialize the handler. Load the Faster R-CNN model.
+        Initialize the handler: load the model.
         """
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model_weights_path = os.path.join(path, "model.pt")  # Adjust path
-
-        # Load model
-        self.model = get_model(num_classes=4) 
-        checkpoint = torch.load(self.model_weights_path, map_location=self.device)
+        # Load the model
+        self.model_weights_path = os.path.join(path, "model.pt")
+        self.model = get_model(NUM_CLASSES).to(DEVICE)
+        checkpoint = torch.load(self.model_weights_path, map_location=DEVICE)
         self.model.load_state_dict(checkpoint["model_state_dict"])
-        self.model.to(self.device)
         self.model.eval()
 
-        # Image preprocessing
-        self.transform = transforms.Compose([
-            transforms.Resize((640, 640)), 
-            transforms.ToTensor(),
-        ])
+        # Preprocessing function
+        self.preprocess = ToTensor()
+
+        # Class labels
+        self.label_map = {1: "yellow", 2: "red", 3: "blue"}
+
+    def preprocess_frame(self, image_bytes):
+        """
+        Convert raw binary image data to a tensor.
+        """
+        # Load image from binary data
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        image_tensor = self.preprocess(image).unsqueeze(0).to(DEVICE)
+        return image_tensor
 
     def __call__(self, data):
         """
-        Process incoming binary image data and return object detection results.
+        Process incoming raw binary image data.
         """
         try:
-            # Read raw binary data (image file)
-            image_bytes = data.get("body", b"") 
-            if not image_bytes:
+            if "body" not in data:
                 return {"error": "No image data provided in request."}
 
-            
-            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-
-            
-            input_tensor = self.transform(image).unsqueeze(0).to(self.device)
-
+            image_bytes = data["body"]  
+            image_tensor = self.preprocess_frame(image_bytes)
 
+            # Perform inference
             with torch.no_grad():
-                predictions = self.model(input_tensor)
+                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()
 
-            
-            threshold = 0.5
-            results = [
-                {"box": box, "label": label, "score": score}
-                for box, label, score in zip(boxes, labels, scores)
-                if score > threshold
-            ]
+            # 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)
+                    })
 
             return {"predictions": results}
         except Exception as e: