Upload app.py

app.py

@@ -0,0 +1,85 @@
+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.Linear(in_features, len(class_names))
+    
+    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()
+
+transform = transforms.Compose([
+    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",
+    live=True,
+    theme=gr.themes.Soft()
+)
+
+if __name__ == "__main__":
+    demo.launch()