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

# Load model (it will be uploaded into the same repo)
model = YOLO('best.pt')
model.eval()

def predict_image(image):
    image = np.array(image)
    results = model(image)

    boxes = results[0].boxes.xyxy
    scores = results[0].boxes.conf
    labels = results[0].boxes.cls

    output_image = image.copy()

    for box, score, label in zip(boxes, scores, labels):
        x1, y1, x2, y2 = map(int, box)
        color = (0, 255, 0)
        thickness = 2

        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, thickness)
        label_text = f"Triangle {score:.2f}"
        cv2.putText(output_image, label_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    output_image_rgb = cv2.cvtColor(output_image, cv2.COLOR_BGR2RGB)
    num_triangles = len(boxes)
    return Image.fromarray(output_image_rgb), f"Detected {num_triangles} triangles!"

interface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil", label="Upload an image"),
    outputs=[gr.Image(label="Detection Result"), gr.Textbox(label="Detection Info")],
    title="Triangle Detection",
    description="Upload an image, and the model will detect triangles in it."
)

interface.launch(share=True)