import gradio as gr
import torch
from PIL import Image, ImageDraw
from transformers import AutoImageProcessor, ViTForImageClassification
from ultralytics import YOLO
import os
from dotenv import load_dotenv

load_dotenv()

HF_TOKEN = os.getenv("HF_TOKEN") 

processor = AutoImageProcessor.from_pretrained("Abduqayum/Vehicle-Color-Recognition", token=HF_TOKEN)
model = ViTForImageClassification.from_pretrained("Abduqayum/Vehicle-Color-Recognition", token=HF_TOKEN)
model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

detection_model = YOLO("yolov8n.pt")  

CLASSES_ID = [2, 3, 5, 7]

def classify_image(image):
    image = image.convert("RGB")

    results = detection_model(image, classes=CLASSES_ID, conf=0.5, verbose=False)
    detections = results[0].boxes.data.cpu().numpy()  

    if len(detections) == 0:
        return "No vehicle detected (car, bus, or truck).", image

    largest = max(detections, key=lambda det: (det[2] - det[0]) * (det[3] - det[1]))
    x1, y1, x2, y2, conf, cls_id = largest
    x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])

    cropped = image.crop((x1, y1, x2, y2))

    inputs = processor(cropped, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.no_grad():
        outputs = model(**inputs)

    logits = outputs.logits
    probs = torch.nn.functional.softmax(logits, dim=-1)
    pred_idx = probs.argmax(dim=-1).item()
    label = model.config.id2label[pred_idx]
    confidence = probs[0, pred_idx].item()

    draw = ImageDraw.Draw(image)
    draw.rectangle([x1, y1, x2, y2], outline="red", width=4)
    draw.text((x1, y1 - 10), f"{label} ({confidence:.2%})", fill="red")

    return f"Prediction: {label} (confidence: {confidence:.2%})", image

gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=["text", "image"],
    title="YOLO Vehicle Detector + ViT Classifier to identify color of vehicles"
).launch()