from ultralytics import YOLO
import gradio as gr
import cv2
import numpy as np
from PIL import Image

# Load your fine-tuned YOLOv8 model
model = YOLO("best.pt")

def detect_objects(image):
    # Convert PIL image to numpy array
    image_np = np.array(image)
    
    # Run YOLOv8 inference
    results = model(image_np)
    
    predictions = []
    
    # Process each result
    for result in results:
        boxes = result.boxes.xyxy.cpu().numpy()  # Bounding box coordinates
        confs = result.boxes.conf.cpu().numpy()    # Confidence scores
        classes = result.boxes.cls.cpu().numpy()   # Class indices
        
        # Iterate over each detection
        for box, conf, cls in zip(boxes, confs, classes):
            x1, y1, x2, y2 = map(int, box[:4])
            label = model.names[int(cls)]  # Map class index to label
            
            predictions.append({
                "label": label,
                "score": float(conf),
                "box": {
                    "xmin": x1,
                    "ymin": y1,
                    "xmax": x2,
                    "ymax": y2
                }
            })
            
            # Draw bounding box on the image (for visualization)
            cv2.rectangle(image_np, (x1, y1), (x2, y2), (0, 255, 0), 3)
    
    # Convert back to PIL Image for display
    image_out = Image.fromarray(image_np)
    
    # Return the structured predictions and the annotated image
    return predictions, image_out

# Create the Gradio interface
interface = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="pil"),
    outputs=[gr.JSON(), gr.Image(type="pil")],  # Returns both JSON predictions and the processed image
    title="YOLOv8 Object Detection",
    description="Upload an image and detect objects using a fine-tuned YOLOv8 model."
)

# Launch the app
interface.launch()