import torch
from torchvision import transforms
from PIL import Image
import io

MODEL_PATH = "model_checkpoint.pt"
NUM_CLASSES = 4  
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Load Faster R-CNN model
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

transform = transforms.Compose([
    transforms.Resize((640, 640)),
    transforms.ToTensor(),
])

model = load_model(MODEL_PATH, NUM_CLASSES)

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):
    try:
        if "image" not in payload:
            return {"error": "No image provided. Please send an image."}
        
        image_bytes = payload["image"].encode("latin1")  
        
        results = detect_objects(image_bytes)
        return results
    except Exception as e:
        return {"error": str(e)}