import gradio as gr
from transformers import ConvNextForImageClassification, AutoImageProcessor
import torch
from PIL import Image

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load model and processor
model = ConvNextForImageClassification.from_pretrained("todorristov/car_classification_model")
processor = AutoImageProcessor.from_pretrained("todorristov/car_classification_model")

model.to(device)  # ****** Make sure model is on the right device

def classify_car(image):
    # Process image and move all input tensors to the device
    inputs = processor(images=image, return_tensors="pt")
    
    # ****** Move each tensor in inputs dict to device individually
    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=1)
        predicted_class_idx = probs.argmax().item()

        # ****** Safer lookup of label with fallback
        key_str = str(predicted_class_idx)
        if key_str in model.config.id2label:
            label = model.config.id2label[key_str]
        elif predicted_class_idx in model.config.id2label:
            label = model.config.id2label[predicted_class_idx]
        else:
            label = "Unknown class"

        confidence = probs[0][predicted_class_idx].item()
    return f"{label} ({confidence:.2%})"

title = "Car Classification - Brand, Model & Model Year"
description = "Upload a car image and the model will predict its brand, model, and year of production."

demo = gr.Interface(
    fn=classify_car,
    inputs=gr.Image(type="pil"),
    outputs=gr.Textbox(label="Prediction"),
    title=title,
    description=description,
    allow_flagging="never"
)

demo.launch()