import gradio as gr
import torch
import numpy as np
import cv2
from PIL import Image, ImageDraw
from transformers import YolosImageProcessor, YolosForObjectDetection

# Load model
processor = YolosImageProcessor.from_pretrained(
    "nickmuchi/yolos-small-finetuned-license-plate-detection"
)
model = YolosForObjectDetection.from_pretrained(
    "nickmuchi/yolos-small-finetuned-license-plate-detection"
)
model.eval()

# -------- Plate Color Classifier -------- #
def classify_plate_color(plate_img):
    img = np.array(plate_img)
    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)

    green = cv2.inRange(hsv, (35, 40, 40), (85, 255, 255))
    yellow = cv2.inRange(hsv, (15, 50, 50), (35, 255, 255))
    white = cv2.inRange(hsv, (0, 0, 200), (180, 30, 255))

    g = np.sum(green)
    y = np.sum(yellow)
    w = np.sum(white)

    if g > y and g > w:
        return "   ELECTRIC Vehicle (Green Plate)   "
    elif y > g and y > w:
        return "    COMMERCIAL Vehicle (Yellow Plate) "
    else:
        return "   PERSONAL Vehicle (White Plate)  "

# -------- Main Pipeline -------- #
def process_image(img):
    image = Image.fromarray(img)

    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    target_sizes = torch.tensor([[image.size[1], image.size[0]]])
    results = processor.post_process_object_detection(
        outputs, threshold=0.3, target_sizes=target_sizes
    )[0]

    draw = ImageDraw.Draw(image)

    if len(results["boxes"]) == 0:
        return image, "No license plate detected"

    box = results["boxes"][0].tolist()
    x1, y1, x2, y2 = map(int, box)

    plate = image.crop((x1, y1, x2, y2))
    vehicle_type = classify_plate_color(plate)

    draw.rectangle([x1, y1, x2, y2], outline="yellow", width=3)
    draw.text((x1, y1 - 10), vehicle_type, fill="black")

    return image, vehicle_type

# -------- Gradio UI -------- #
with gr.Blocks() as demo:
    gr.Markdown("# 🚗 Vehicle Classification using License Plate")
    gr.Markdown("Upload or take a photo of a car. The AI detects the license plate and classifies the vehicle.")

    with gr.Row():
        input_img = gr.Image(type="numpy", sources=["upload", "webcam"])
        output_img = gr.Image()
    
    result = gr.Textbox(label="Vehicle Type")

    btn = gr.Button("Detect Vehicle")
    btn.click(process_image, input_img, [output_img, result])

demo.launch()