import gradio as gr
from PIL import Image
import numpy as np
import cv2
from ultralytics import YOLO
from transformers import pipeline
import matplotlib.pyplot as plt

# Load models
yolo_model = YOLO("yolov8n.pt")
damage_pipe = pipeline("image-classification", model="beingamit99/car_damage_detection")

# License Plate Detection
def detect_license_plate(img: Image.Image):
    results = yolo_model.predict(img)
    result = results[0]
    img_array = np.array(img.convert("RGB"))

    for box in result.boxes.xyxy:
        x1, y1, x2, y2 = map(int, box)
        img_array = cv2.rectangle(img_array, (x1, y1), (x2, y2), (0, 255, 0), 2)

    return Image.fromarray(img_array)

# Car Damage Classification
def classify_damage(image: Image.Image):
    results = damage_pipe(image)
    labels = [res["label"] for res in results]
    scores = [res["score"] for res in results]

    fig, ax = plt.subplots()
    ax.barh(labels, scores, color="crimson")
    ax.set_xlabel("Confidence")
    ax.set_xlim(0, 1)
    ax.set_title("Damage Type Classification")
    plt.tight_layout()

    return fig

# Build UI
with gr.Blocks() as app:
    gr.Markdown("# 🚘 Car Analyzer\nChoose a tool below:")

    with gr.Tab("🛂 License Plate Detection"):
        with gr.Row():
            with gr.Column():
                lp_input = gr.Image(type="pil", label="Upload Car Image")
                lp_btn = gr.Button("Detect")
            with gr.Column():
                lp_output = gr.Image(label="Detected License Plate")
        lp_btn.click(detect_license_plate, inputs=lp_input, outputs=lp_output)

    with gr.Tab("💥 Car Damage Classification"):
        with gr.Row():
            with gr.Column():
                dmg_input = gr.Image(type="pil", label="Upload Damaged Car Image")
                dmg_btn = gr.Button("Classify")
            with gr.Column():
                dmg_output = gr.Plot(label="Classification Results")
        dmg_btn.click(classify_damage, inputs=dmg_input, outputs=dmg_output)

app.launch()