Phoenix3238/CalTech101-Classifier

1import os2import json3import time4import glob5import torch6import torch.nn as nn7from torchvision import models, transforms8from PIL import Image9import gradio as gr10import spaces 11 12device = torch.device("cuda" if torch.cuda.is_available() else "cpu")13 14with open("class_names.json", "r") as f:15    class_names = json.load(f)16 17def load_model():18    model = models.convnext_tiny(weights=None)19    in_features = model.classifier[2].in_features20    model.classifier[2] = nn.Sequential(21        nn.Dropout(p=0.5), # Standard dropout layer (index 0)22        nn.Linear(in_features, 102) # Linear layer (index 1)23    )24    25    state_dict = torch.load("convnext_caltech101.pth", map_location=device)26    model.load_state_dict(state_dict)27    28    model.to(device)29    model.eval()30    return model31 32model = load_model()33 34class ConvertToRGB:35    def __call__(self, img):36        return img.convert("RGB")37 38transform = transforms.Compose([39    ConvertToRGB(),40    transforms.Resize((224, 224)),41    transforms.ToTensor(),42    transforms.Normalize(43        mean=[0.485, 0.456, 0.406],44        std=[0.229, 0.224, 0.225]45    )46])47 48# Adding the ZeroGPU decorator to the inference function49@spaces.GPU50def predict(image):51    if image is None:52        return None, 0.053        54    start_time = time.perf_counter()55    img_tensor = transform(image).unsqueeze(0).to(device)56    57    with torch.inference_mode():58        logits = model(img_tensor)59        probs = torch.softmax(logits, dim=1).squeeze()60        61    # Ensure accurate timing on GPU62    if device.type == "cuda":63        torch.cuda.synchronize()64        65    latency = time.perf_counter() - start_time66    67    confidences = {68        class_names[i]: float(probs[i])69        for i in range(len(class_names))70    }71    72    return confidences, round(latency, 4)73 74# Building the UI75example_images = glob.glob("demo_images/*.jpg")76 77demo = gr.Interface(78    fn=predict,79    inputs=gr.Image(type="pil", label="Upload Input Image"),80    outputs=[81        gr.Label(num_top_classes=5, label="Top 5 Predictions"),82        gr.Number(label="Inference Latency (seconds)")83    ],84    title="CalTech-101 Image Classifier",85    description="Upload an image to see how the fine-tuned ConvNeXt model classifies it.",86    examples=example_images if example_images else None,87    flagging_mode="never"88)89 90if __name__ == "__main__":91    demo.launch()