from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from ultralytics import YOLO
from PIL import Image
import io
import base64

app = FastAPI()

# Load the YOLO model
model = YOLO(r'best.pt')

class ImageData(BaseModel):
    image_base64: str

@app.post("/process_image/")
async def process_image(data: ImageData):
    try:
        # Decode the base64 string to an image
        image_data = base64.b64decode(data.image_base64)
        image = Image.open(io.BytesIO(image_data))

        # Process the image with YOLO
        results = model(image)
        result = results[0]

        # Extract bounding boxes and confidence scores
        boxes = result.boxes.xyxy  # Bounding box coordinates
        scores = result.boxes.conf  # Confidence scores

        if len(boxes) > 0:
            # Get the index of the bounding box with the highest score
            highest_score_idx = scores.argmax()
            # Extract the bounding box with the highest score
            highest_score_box = boxes[highest_score_idx].tolist()
            x1, y1, x2, y2 = map(int, highest_score_box)  # Convert to integers
        else:
            # If no boxes, return the whole image dimensions
            x1, y1, x2, y2 = 0, 0, image.width, image.height

        return {"x1": x1, "y1": y1, "x2": x2, "y2": y2}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))