import gradio as gr
import torch
from PIL import Image
import io
from ultralytics import YOLO

# --- Load YOLO Model ---
MODEL_PATH = 'model/char.pt'
try:
    model = YOLO(MODEL_PATH)
    print(f"Model loaded successfully from: {MODEL_PATH}")
except Exception as e:
    print(f"Error loading model: {e}")
    model = None

# --- Prediction Function for Gradio ---
def predict(image):
    if model is None or image is None:
        return None

    try:
        img = Image.fromarray(image).convert('RGB')
        results = model(img)

        predictions = []
        for result in results:
            for box in result.boxes:
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                label = model.model.names[int(box.cls)]
                confidence = float(box.conf[0])
                predictions.append({'label': label, 'confidence': confidence, 'bbox': (x1, y1, x2, y2)})

        # Draw bounding boxes on the image
        draw = ImageDraw.Draw(img)
        for pred in predictions:
            x1, y1, x2, y2 = pred['bbox']
            label = f"{pred['label']} ({pred['confidence']:.2f})"
            draw.rectangle([x1, y1, x2, y2], outline="green", width=2)
            draw.text((x1, y1 - 10), label, fill="red")

        return img

    except Exception as e:
        return f"Error during prediction: {e}"

# --- Gradio Interface ---
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(label="Upload an Image"),
    outputs=gr.Image(label="Image with Predictions"),
    title="YOLO Object Detection",
    description="Upload an image to see object detection predictions using a YOLO model.",
)

iface.launch()