Spaces:
Sleeping
Sleeping
File size: 1,968 Bytes
3cd3c9a 33566fb 3cd3c9a 33566fb 3cd3c9a 33566fb 3cd3c9a 33566fb 3cd3c9a 33566fb 3cd3c9a 33566fb 3cd3c9a 33566fb |
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 |
import gradio as gr
from transformers import AutoModelForImageClassification
import torch
import torchvision.transforms as transforms
from PIL import Image
import traceback
import sys
# Load model from Hub instead of local file
def load_model():
try:
model = AutoModelForImageClassification.from_pretrained(
"nragrawal/resnet-imagenet",
trust_remote_code=True
)
model.eval()
return model
except Exception as e:
print(f"Error loading model: {str(e)}")
print(traceback.format_exc())
raise e
# Preprocessing
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])
])
# Inference function
def predict(image):
try:
model = load_model()
# Preprocess image
img = Image.fromarray(image)
img = transform(img).unsqueeze(0)
# Inference
with torch.no_grad():
output = model(img)
probabilities = torch.nn.functional.softmax(output[0], dim=0)
# Get top 5 predictions
top5_prob, top5_catid = torch.topk(probabilities, 5)
return {f"Class {i}": float(prob) for i, prob in zip(top5_catid, top5_prob)}
except Exception as e:
print(f"Error during prediction: {str(e)}")
print(traceback.format_exc())
return {"error": str(e)}
# Create Gradio interface with error handling
iface = gr.Interface(
fn=predict,
inputs=gr.Image(),
outputs=gr.Label(num_top_classes=5),
title="ResNet Image Classification",
description="Upload an image to classify it using ResNet",
allow_flagging="never"
)
# Add error handling to launch
try:
iface.launch()
except Exception as e:
print(f"Error launching interface: {str(e)}")
print(traceback.format_exc()) |