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:
    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 path

        # Load model
        self.model = get_model(num_classes=4) 
        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()

        # Image preprocessing
        self.transform = transforms.Compose([
            transforms.Resize((640, 640)), 
            transforms.ToTensor(),
        ])

    def __call__(self, data):
        """
        Process incoming binary image data and return object detection results.
        """
        try:
            # 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 = Image.open(io.BytesIO(image_bytes)).convert("RGB")

            
            input_tensor = self.transform(image).unsqueeze(0).to(self.device)


            with torch.no_grad():
                predictions = self.model(input_tensor)

            
            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
            ]

            return {"predictions": results}
        except Exception as e:
            return {"error": str(e)}