Update handler.py

handler.py

@@ -4,58 +4,64 @@ from torchvision import transforms
 from PIL import Image
 import io
 
-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"
-
-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)
-    checkpoint = torch.load(model_path, map_location=DEVICE)
-    model.load_state_dict(checkpoint["model_state_dict"])
-    model.to(DEVICE)
-    model.eval()
-    return model
-
-model = load_model(MODEL_PATH, NUM_CLASSES)
-
-transform = transforms.Compose([
-    transforms.Resize((640, 640)),
-    transforms.ToTensor(),
-])
-
-def detect_objects(image_bytes):
-    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-    input_tensor = transform(image).unsqueeze(0).to(DEVICE)
-
-    with torch.no_grad():
-        predictions = model(input_tensor)
-
-    boxes = predictions[0]["boxes"].cpu().tolist()
-    labels = predictions[0]["labels"].cpu().tolist()
-    scores = predictions[0]["scores"].cpu().tolist()
-
-    confidence_threshold = 0.5
-    results = [
-        {"box": box, "label": label, "score": score}
-        for box, label, score in zip(boxes, labels, scores)
-        if score > confidence_threshold
-    ]
-
-    return {"predictions": results}
-
-def inference(payload):
-    import base64
-    try:
-        if "image" not in payload:
-            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:
-        return {"error": str(e)}
+# Import your Faster R-CNN model definition
+from model import get_model
+
+class EndpointHandler:
+    def __init__(self, path: str = ""):
+        """
+        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
+        
+        # Load the model
+        self.model = get_model(num_classes=4)  # Modify for your num_classes
+        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
+        self.transform = transforms.Compose([
+            transforms.Resize((640, 640)),  # Adjust size to match your training setup
+            transforms.ToTensor(),
+        ])
+
+    def __call__(self, data):
+        """
+        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}
+                for box, label, score in zip(boxes, labels, scores)
+                if score > threshold
+            ]
+
+            return [{"predictions": predictions}]
+        except Exception as e:
+            return [{"error": str(e)}]