app.py

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

# Load YOLOv11 model
model = YOLO("best.pt")  # Ensure this matches your uploaded file

# Damage class labels
class_names = {
    0: 'no damage',
    1: 'lost parts',
    2: 'torn',
    3: 'dent',
    4: 'paint scratch',
    5: 'hole',
    6: 'broken glass',
    7: 'broken lamp'
}

def detect_damage(img, conf=0.31):
    img_np = np.array(img)
    results = model(img_np, conf=conf, augment=True)

    for r in results:
        boxes = r.boxes.xyxy.cpu().numpy()
        scores = r.boxes.conf.cpu().numpy()
        classes = r.boxes.cls.cpu().numpy().astype(int)

        for box, score, cls in zip(boxes, scores, classes):
            x1, y1, x2, y2 = map(int, box)
            label = f"{class_names.get(cls, 'Unknown')} {score:.2f}"
            cv2.rectangle(img_np, (x1, y1), (x2, y2), (0, 0, 255), 2)
            cv2.putText(img_np, label, (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
    return img_np

with gr.Blocks() as demo:
    gr.Markdown("## YOLOv11m – Vehicle Damage Detection")
    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Image")
        image_output = gr.Image(type="numpy", label="Detected Image")
    confidence = gr.Slider(0.1, 1.0, value=0.31, label="Confidence Threshold")
    submit_btn = gr.Button("Run Detection")

    submit_btn.click(fn=detect_damage, inputs=[image_input, confidence], outputs=image_output)

demo.launch()