G-Sinha/Cattle_Breed_Classifier

1import io
2import base64
3import torch
4import torch.nn.functional as F
5import torchvision.transforms as transforms
6import timm
7from PIL import Image
8import gradio as gr
9import os
10
11# ------------------------
12# Class names
13# ------------------------
14class_names = [
15    "Alambadi", "Amritmahal", "Ayrshire", "Banni", "Bargur", "Bhadawari",
16    "Brown_Swiss", "Dangi", "Deoni", "Gir", "Guernsey", "Hallikar", "Hariana",
17    "Holstein_Friesian", "Jaffrabadi", "Jersey", "Kangayam", "Kankrej", 
18    "Kasargod", "Kenkatha", "Kherigarh", "Khillari", "Krishna_Valley", 
19    "Malnad_gidda", "Mehsana", "Murrah", "Nagori", "Nagpuri", "Nili_Ravi", 
20    "Nimari", "Ongole", "Pulikulam", "Rathi", "Red_Dane", "Red_Sindhi", 
21    "Sahiwal", "Surti", "Tharparkar", "Toda", "Umblachery", "Vechur"
22]
23NUM_CLASSES = len(class_names)
24
25# ------------------------
26# Load model
27# ------------------------
28DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
29
30MODEL_PATH = "best_model_final.pth"
31if not os.path.exists(MODEL_PATH):
32    raise FileNotFoundError(f"Model file not found at {MODEL_PATH}")
33
34model = timm.create_model("convnext_tiny", pretrained=False, num_classes=NUM_CLASSES)
35
36checkpoint = torch.load(MODEL_PATH, map_location="cpu")
37if isinstance(checkpoint, dict):
38    if "model_state_dict" in checkpoint:
39        state_dict = checkpoint["model_state_dict"]
40    elif "state_dict" in checkpoint:
41        state_dict = checkpoint["state_dict"]
42    else:
43        state_dict = checkpoint
44else:
45    state_dict = checkpoint
46
47model.load_state_dict(state_dict)
48model.to(DEVICE)
49model.eval()
50
51# ------------------------
52# Preprocessing
53# ------------------------
54transform = transforms.Compose([
55    transforms.Resize((224, 224)),
56    transforms.ToTensor(),
57    transforms.Normalize([0.485, 0.456, 0.406],
58                         [0.229, 0.224, 0.225])
59])
60
61# ------------------------
62# Inference function
63# ------------------------
64def predict(image: Image.Image):
65    img = image.convert("RGB")
66    input_tensor = transform(img).unsqueeze(0).to(DEVICE)
67
68    with torch.no_grad():
69        logits = model(input_tensor)
70        probs = F.softmax(logits, dim=1)
71        conf, pred_idx = torch.max(probs, dim=1)
72        confidence = float(conf.cpu().item())
73        pred_idx = int(pred_idx.cpu().item())
74
75    breed_name = class_names[pred_idx]
76    return {breed_name: confidence}
77
78# ------------------------
79# Gradio Interface
80# ------------------------
81demo = gr.Interface(
82    fn=predict,
83    inputs=gr.Image(type="pil"),
84    outputs=gr.Label(num_top_classes=3),
85    title="🐄 Cow Breed Classifier",
86    description="Upload a cow image to identify its breed."
87)
88
89if __name__ == "__main__":
90    demo.launch()
91