File size: 1,557 Bytes
8193622 | 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 | import torch
from torchvision import models, transforms
from PIL import Image
import gradio as gr
# =======================
# Configuration
# =======================
device = "cuda" if torch.cuda.is_available() else "cpu"
MODEL_PATH = "cattle_breed_efficientnetb3_pytorch.pth" # Upload this to the Space
CLASS_NAMES = ["Gir", "Deoni", "Murrah"]
# =======================
# Load Model
# =======================
model = models.efficientnet_b3(pretrained=False)
model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, len(CLASS_NAMES))
model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
model.to(device)
model.eval()
# =======================
# Image Preprocessing
# =======================
transform = transforms.Compose([
transforms.Resize((300, 300)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])
# =======================
# Prediction Function
# =======================
def predict(image):
image = image.convert("RGB")
img_tensor = transform(image).unsqueeze(0).to(device)
with torch.no_grad():
output = model(img_tensor)
pred_idx = torch.argmax(output, dim=1).item()
return CLASS_NAMES[pred_idx]
# =======================
# Gradio Interface
# =======================
iface = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs="text",
title="Indian Bovine Breed Classifier",
description="Upload an image of a cow and the model will predict its breed."
)
iface.launch()
|