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

# Load fine-tuned YOLOv8 model for car damage detection
model = YOLO("best.pt")  # Replace with your trained model file

def predict(input_img):
    # Convert PIL image to OpenCV format
    image = np.array(input_img)
    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Run inference
    results = model(image)
    
    # Draw bounding boxes
    for result in results:
        for box in result.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            conf = float(box.conf[0])
            label = f"Damage: {conf:.2f}"
            
            # Draw red bounding box
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255), 3)  # Red color (BGR: 0,0,255)
            cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
    
    # Convert back to PIL format
    output_img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    return output_img

# Gradio interface
gradio_app = gr.Interface(
    fn=predict,
    inputs=gr.Image(label="Upload a car image", sources=['upload', 'webcam'], type="pil"),
    outputs=gr.Image(label="Detected Damage"),
    title="Car Damage Detection",
    description="Upload an image of a car, and the model will detect and highlight damaged areas."
)

gradio_app.launch()