import os
import json
import time
import glob
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import gradio as gr
import spaces 

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

with open("class_names.json", "r") as f:
    class_names = json.load(f)

def load_model():
    model = models.convnext_tiny(weights=None)
    in_features = model.classifier[2].in_features
    model.classifier[2] = nn.Sequential(
        nn.Dropout(p=0.5), # Standard dropout layer (index 0)
        nn.Linear(in_features, 102) # Linear layer (index 1)
    )
    
    state_dict = torch.load("convnext_caltech101.pth", map_location=device)
    model.load_state_dict(state_dict)
    
    model.to(device)
    model.eval()
    return model

model = load_model()

class ConvertToRGB:
    def __call__(self, img):
        return img.convert("RGB")

transform = transforms.Compose([
    ConvertToRGB(),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

# Adding the ZeroGPU decorator to the inference function
@spaces.GPU
def predict(image):
    if image is None:
        return None, 0.0
        
    start_time = time.perf_counter()
    img_tensor = transform(image).unsqueeze(0).to(device)
    
    with torch.inference_mode():
        logits = model(img_tensor)
        probs = torch.softmax(logits, dim=1).squeeze()
        
    # Ensure accurate timing on GPU
    if device.type == "cuda":
        torch.cuda.synchronize()
        
    latency = time.perf_counter() - start_time
    
    confidences = {
        class_names[i]: float(probs[i])
        for i in range(len(class_names))
    }
    
    return confidences, round(latency, 4)

# Building the UI
example_images = glob.glob("demo_images/*.jpg")

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Input Image"),
    outputs=[
        gr.Label(num_top_classes=5, label="Top 5 Predictions"),
        gr.Number(label="Inference Latency (seconds)")
    ],
    title="CalTech-101 Image Classifier",
    description="Upload an image to see how the fine-tuned ConvNeXt model classifies it.",
    examples=example_images if example_images else None,
    flagging_mode="never"
)

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