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

# Load YOLO model
model = YOLO("best.pt").to("cpu")

def predict(image):
    """Detect objects using YOLO and return cropped images."""
    image_np = np.array(image)
    
    # Perform inference
    results = model(image_np, conf=0.85, device='cpu')
    
    cropped_images = []
    for result in results:
        for box in result.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            cropped = image_np[y1:y2, x1:x2]  # Crop using bounding box
            cropped_images.append(Image.fromarray(cropped))
    
    return cropped_images if cropped_images else image  # Return original if no detections

# Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs="image",
    outputs=gr.Gallery(label="Cropped Objects"),  # Allows multiple images to be displayed
    title="YOLO Object Detection - Cropped Images"
)

iface.launch()