Spaces:
Sleeping
Sleeping
| 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 | |
| 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() |