File size: 1,678 Bytes
a1fbc0c b991315 | 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 | import gradio as gr
import torch
import torchvision.transforms as transforms
from PIL import Image
import torchvision.models as models
import torch.nn as nn
# 🔹 Load your trained model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
model.fc = nn.Sequential(
nn.Linear(2048, 512),
nn.ReLU(),
nn.Dropout(0.4),
nn.Linear(512, 15)
)
model.load_state_dict(torch.load("best_model.pth", map_location=device)) # 🟡 Replace with your file path
model.to(device)
model.eval()
# 🔹 Preprocessing (must match training)
transform = transforms.Compose([
transforms.Resize(256),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
]) # Enclose the transforms in a list
# 🔹 Class labels (update to match your labels)
class_names = [
'Bear', 'Bird', 'Cat', 'Cow', 'Deer', 'Dog', 'Dolphin',
'Elephant', 'Giraffe', 'Horse', 'Kangaroo', 'Lion',
'Panda', 'Tiger', 'Zebra'
]
# 🔹 Inference function
def classify_image(img):
img = transform(img).unsqueeze(0).to(device)
with torch.no_grad():
outputs = model(img)
probs = torch.nn.functional.softmax(outputs, dim=1)
return {class_names[i]: float(probs[0][i]) for i in range(len(class_names))}
# 🔹 Gradio UI
interface = gr.Interface(
fn=classify_image,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=5),
title="Animal Image Classifier",
description="Upload an image of an animal and get the top predictions!"
)
# 🔹 Launch the app (use share=True in Colab to get a public link)
interface.launch() |