Update handler.py

handler.py

@@ -4,7 +4,6 @@ from torchvision import transforms
 from PIL import Image
 import io
 
-# Import your Faster R-CNN model definition
 from model import get_model
 
 class EndpointHandler:
@@ -13,10 +12,10 @@ 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")  # Adjust for your file name
+        self.model_weights_path = os.path.join(path, "model.pt") 
         
         # Load the model
-        self.model = get_model(num_classes=4)  # Modify for your num_classes
+        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"])
@@ -25,7 +24,7 @@ class EndpointHandler:
 
         # Define image preprocessing
         self.transform = transforms.Compose([
-            transforms.Resize((640, 640)),  # Adjust size to match your training setup
+            transforms.Resize((640, 640)),
             transforms.ToTensor(),
         ])
 
@@ -34,27 +33,21 @@ class EndpointHandler:
         Process the incoming request and return object detection predictions.
         """
         try:
-            # Expect input data to include a Base64-encoded image
             if "image" not in data:
                 return [{"error": "No 'image' provided in request."}]
 
-            # Convert Base64-encoded image to bytes
             image_bytes = data["image"].encode("latin1")
             image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
-            # Preprocess the image
             input_tensor = self.transform(image).unsqueeze(0).to(self.device)
 
-            # Run inference
             with torch.no_grad():
                 outputs = self.model(input_tensor)
 
-            # Extract results
             boxes = outputs[0]["boxes"].cpu().tolist()
             labels = outputs[0]["labels"].cpu().tolist()
             scores = outputs[0]["scores"].cpu().tolist()
 
-            # Confidence threshold
             threshold = 0.5
             predictions = [
                 {"box": box, "label": label, "score": score}