import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import gradio as gr

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

# Class names (order must match training)
class_names = ["Helmet", "Helmetless", "Number Plate"]

# Define model
num_classes = len(class_names)
model = models.inception_v3(weights=None, aux_logits=True, init_weights=True)
model.fc = nn.Linear(model.fc.in_features, num_classes)

# Load trained weights
model.load_state_dict(torch.load("inceptionv3_model.pth", map_location=device))
model = model.to(device)
model.eval()

# Preprocessing
transform = transforms.Compose([
    transforms.Resize((299, 299)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])
])

# Prediction function
def predict(image):
    if image is None:
        return {}
    image = Image.fromarray(image).convert("RGB")
    img_tensor = transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        outputs = model(img_tensor)
        if isinstance(outputs, tuple):  # inception returns (main, aux)
            outputs = outputs[0]
        probs = torch.nn.functional.softmax(outputs[0], dim=0)

    confidences = {class_names[i]: float(probs[i]) for i in range(num_classes)}
    return confidences

# Gradio Interface (upload OR camera)
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(source="camera", type="numpy", label="Take a Picture"),
    outputs=gr.Label(num_top_classes=3, label="Prediction"),
    title="Helmet, Helmetless & Number Plate Classifier",
    description="Take a picture using your camera and the model will classify it."
)

if __name__ == "__main__":
    demo.launch()