Spaces:
Sleeping
Sleeping
| import spaces | |
| import gradio as gr | |
| import torch | |
| import torch.nn as nn | |
| from torchvision import models, transforms | |
| from PIL import Image | |
| from huggingface_hub import hf_hub_download | |
| # Order matches full_train.classes from the training notebook (ImageFolder's | |
| # alphabetical folder sort) — do not reorder without re-checking the notebook. | |
| CLASS_NAMES = [ | |
| "Black_Beans", "Black_EyedPea", "Brown_Chickpea", "Desi_chickpea", | |
| "Horse_Gram", "Husked_BlackGram", "Husked_RedLentil", "Kidney_Bean", | |
| "Red_Beans", "Red_KidneyBean", "Red_chickpea", "Soya_Bean", | |
| "White_Chickpea", "White_beans", "White_peabeans", "Whole_BlackGram", | |
| "Whole_GreenGram", "Whole_TurkishGram", "chana1", "chana2", | |
| ] | |
| # Weights live in a separate Model repo (not bundled in this Space), so the | |
| # Space repo stays small and never needs Git LFS. Downloaded once per | |
| # container start and cached. | |
| WEIGHTS_PATH = hf_hub_download( | |
| repo_id="silversur4/indian-lentils-image-classifier", | |
| filename="resnet_lentils.pth", | |
| ) | |
| # ZeroGPU only grants GPU access inside a function decorated with | |
| # @spaces.GPU — no CUDA context exists at import time, so the model is | |
| # loaded on CPU here and moved to CUDA per-call inside predict(). | |
| model = models.resnet18(weights=None) | |
| model.fc = nn.Linear(model.fc.in_features, len(CLASS_NAMES)) | |
| model.load_state_dict(torch.load(WEIGHTS_PATH, map_location="cpu")) | |
| model.eval() | |
| # Same eval_transform used during evaluation in the training notebook. | |
| transform = transforms.Compose([ | |
| transforms.Resize(256), | |
| transforms.CenterCrop(224), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), | |
| ]) | |
| def predict(image: Image.Image): | |
| if image is None: | |
| return {} | |
| device = torch.device("cuda") | |
| model.to(device) | |
| x = transform(image.convert("RGB")).unsqueeze(0).to(device) | |
| with torch.no_grad(): | |
| probs = torch.softmax(model(x), dim=1)[0] | |
| return {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))} | |
| demo = gr.Interface( | |
| fn=predict, | |
| inputs=gr.Image(type="pil", sources=["upload", "webcam"], label="Photo of a lentil / bean / chickpea"), | |
| outputs=gr.Label(num_top_classes=5, label="Prediction"), | |
| title="Indian Lentils Classifier", | |
| description=( | |
| "ResNet18, fine-tuned on the Indian Lentils dataset (20 classes). " | |
| "Upload a photo or use your camera to classify a lentil, bean, or chickpea variety." | |
| ), | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |