Spaces:
Sleeping
Sleeping
File size: 3,057 Bytes
6a81b92 7036bb9 6a81b92 7036bb9 6a81b92 7036bb9 6a81b92 7036bb9 6a81b92 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
import gradio as gr
import torch
import torchvision.transforms as transforms
from PIL import Image
import requests
from io import BytesIO
import json
import torchvision.models as models
from transformers import AutoImageProcessor
# Load ImageNet class labels
LABELS_URL = "https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json"
response = requests.get(LABELS_URL)
labels = json.loads(response.text)
def load_model():
"""
Load model and processor from Hugging Face Hub
"""
model_id = "jatingocodeo/ImageNet"
# Initialize ResNet50 model
model = models.resnet50(weights=None)
model.fc = torch.nn.Linear(model.fc.in_features, 1000) # 1000 ImageNet classes
# Load model weights
checkpoint = torch.hub.load_state_dict_from_url(
f"https://huggingface.co/{model_id}/resolve/main/pytorch_model.bin",
map_location="cpu"
)
model.load_state_dict(checkpoint)
model.eval()
# Create processor
processor = AutoImageProcessor.from_pretrained(model_id)
return model, processor
def predict(image):
"""
Make prediction on input image
"""
if image is None:
return None
try:
# Load model and processor (with caching)
model, processor = load_model()
# Process image
inputs = processor(image, return_tensors="pt")
# Get predictions
with torch.no_grad():
outputs = model(inputs.pixel_values)
# Get probabilities and classes
probs = torch.nn.functional.softmax(outputs, dim=1)[0]
top_probs, top_indices = torch.topk(probs, k=5)
# Format results
results = {}
for prob, idx in zip(top_probs, top_indices):
label = labels[idx.item()]
confidence = prob.item() * 100
results[label] = confidence
return results
except Exception as e:
return f"Error processing image: {str(e)}"
# Create Gradio interface
title = "ImageNet Classifier"
description = """
## ResNet50 ImageNet Classifier
This model classifies images into 1000 ImageNet categories. Upload an image or use one of the example images to get predictions.
### Instructions:
1. Upload an image using the input box below
2. The model will predict the top 5 classes for the image
3. Results show class names and confidence scores
### Model Details:
- Architecture: ResNet50
- Dataset: ImageNet
- Input Size: 224x224
- Number of Classes: 1000
"""
# Example images
examples = [
"examples/dog.jpg",
"examples/cat.jpg",
"examples/bird.jpg",
"examples/car.jpg",
"examples/flower.jpg"
]
# Create the interface
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil", label="Upload Image"),
outputs=gr.Label(num_top_classes=5, label="Predictions"),
title=title,
description=description,
examples=examples,
theme=gr.themes.Soft(),
allow_flagging="never"
)
# Launch the app
iface.launch() |