Update app.py

app.py

@@ -1,4 +1,11 @@
 # app.py
+import os
+import json
+import pickle
+import joblib
+import warnings
+from typing import Tuple, Dict
+
 import gradio as gr
 import torch
 import torch.nn as nn
@@ -7,15 +14,19 @@ import torchvision.models as tv_models
 import torchvision.transforms as T
 import numpy as np
 from PIL import Image
-import pickle, joblib, json, os, warnings
+
 warnings.filterwarnings("ignore")
 
+# Optional timm
 try:
     import timm
     HAS_TIMM = True
 except Exception:
     HAS_TIMM = False
 
+# ---------------------------
+# Defaults & metadata
+# ---------------------------
 DEFAULT_CLASSES = [
     "Ayrshire cattle","Brown Swiss cattle","Holstein Friesian cattle",
     "Jaffrabadi","Jersey cattle","Murrah","Red Dane cattle",
@@ -83,20 +94,19 @@ BREED_INFO = {
 IMAGENET_MEAN = [0.485, 0.456, 0.406]
 IMAGENET_STD  = [0.229, 0.224, 0.225]
 
-def strip_module_prefix(state_dict):
-    new_sd = {}
-    for k, v in state_dict.items():
-        new_sd[k[7:]] = v if k.startswith("module.") else v
-        if not k.startswith("module."):
-            new_sd[k] = v
+# ---------------------------
+# Helpers
+# ---------------------------
+def strip_module_prefix(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
     clean = {}
-    for k, v in new_sd.items():
+    for k, v in state_dict.items():
         if k.startswith("module."):
-            continue
-        clean[k] = v
+            clean[k[7:]] = v
+        else:
+            clean[k] = v
     return clean
 
-def file_to_path(file_obj):
+def file_to_path(file_obj) -> str:
     if isinstance(file_obj, str):
         return file_obj
     if hasattr(file_obj, "name"):
@@ -105,9 +115,12 @@ def file_to_path(file_obj):
         return file_obj.get("name") or file_obj.get("path") or file_obj.get("file")
     raise ValueError("Unsupported file input type")
 
-def make_head(in_dim, num_classes):
+def make_head(in_dim: int, num_classes: int) -> nn.Module:
     return nn.Sequential(nn.Dropout(0.2), nn.Linear(in_dim, num_classes))
 
+# ---------------------------
+# Classifier
+# ---------------------------
 class IndianBovineClassifier:
     def __init__(self):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -116,50 +129,51 @@ class IndianBovineClassifier:
         self.class_names = list(DEFAULT_CLASSES)
         self.num_classes = len(self.class_names)
         self.preprocess = T.Compose([
-            T.Resize((224,224)),
+            T.Resize((224, 224)),
             T.ToTensor(),
-            T.Normalize(IMAGENET_MEAN, IMAGENET_STD)
+            T.Normalize(IMAGENET_MEAN, IMAGENET_STD),
         ])
         self._try_autoload()
 
-    def _build_arch(self, arch: str, num_classes: int):
+    def _build_arch(self, arch: str, num_classes: int) -> nn.Module:
         a = (arch or "").strip()
         if a and HAS_TIMM:
             try:
                 m = timm.create_model(a, pretrained=False, num_classes=num_classes)
                 cfg = getattr(m, "default_cfg", None)
                 if cfg:
-                    size = cfg.get("input_size", (3,224,224))[-1]
+                    size = cfg.get("input_size", (3, 224, 224))[-1]
                     mean = list(cfg.get("mean", IMAGENET_MEAN))
-                    std  = list(cfg.get("std", IMAGENET_STD))
+                    std = list(cfg.get("std", IMAGENET_STD))
                     self.preprocess = T.Compose([
-                        T.Resize((size,size)),
+                        T.Resize((size, size)),
                         T.ToTensor(),
-                        T.Normalize(mean, std)
+                        T.Normalize(mean, std),
                     ])
                 return m
             except Exception:
                 pass
-        if a.lower() in {"resnet18","tv_resnet18"}:
+        if a.lower() in {"resnet18", "tv_resnet18"}:
             m = tv_models.resnet18(weights=None)
             m.fc = nn.Linear(m.fc.in_features, num_classes)
             return m
-        if a.lower() in {"efficientnet_b0","tv_efficientnet_b0"}:
+        if a.lower() in {"efficientnet_b0", "tv_efficientnet_b0"}:
             m = tv_models.efficientnet_b0(weights=None)
             in_dim = m.classifier[1].in_features
             m.classifier = make_head(in_dim, num_classes)
             return m
+        # fallback
         return self._simple_cnn(num_classes)
 
-    def _simple_cnn(self, nc: int):
+    def _simple_cnn(self, nc: int) -> nn.Module:
         class Simple(nn.Module):
             def __init__(self, out_dim):
                 super().__init__()
                 self.features = nn.Sequential(
-                    nn.Conv2d(3,64,3,padding=1), nn.ReLU(True), nn.MaxPool2d(2),
-                    nn.Conv2d(64,128,3,padding=1), nn.ReLU(True), nn.MaxPool2d(2),
-                    nn.Conv2d(128,256,3,padding=1), nn.ReLU(True),
-                    nn.AdaptiveAvgPool2d((1,1))
+                    nn.Conv2d(3, 64, 3, padding=1), nn.ReLU(True), nn.MaxPool2d(2),
+                    nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(True), nn.MaxPool2d(2),
+                    nn.Conv2d(128, 256, 3, padding=1), nn.ReLU(True),
+                    nn.AdaptiveAvgPool2d((1, 1)),
                 )
                 self.classifier = nn.Sequential(nn.Dropout(0.5), nn.Linear(256, out_dim))
             def forward(self, x):
@@ -169,6 +183,7 @@ class IndianBovineClassifier:
         return Simple(nc)
 
     def _try_autoload(self):
+        # Attempt common filenames; quietly fall back to demo if none
         candidates = [
             ("indian_bovine_breeds.pth","pytorch"),
             ("indian_bovine_model.pth","pytorch"),
@@ -176,12 +191,12 @@ class IndianBovineClassifier:
             ("indian_bovine_breeds.pkl","pickle"),
             ("model.pkl","pickle"),
             ("indian_bovine_breeds.joblib","joblib"),
-            ("model.joblib","joblib")
+            ("model.joblib","joblib"),
         ]
         for path, kind in candidates:
             if os.path.exists(path):
                 try:
-                    self._load_from_path(path, kind)
+                    self._load_from_path(path, kind=kind)
                     print(f"Loaded model: {path}")
                     return
                 except Exception as e:
@@ -189,15 +204,15 @@ class IndianBovineClassifier:
         self.model = self._simple_cnn(self.num_classes).to(self.device).eval()
         self.model_type = "demo"
 
-    def _maybe_set_classes_from_meta(self, meta: dict):
-        keys = ["classes","class_names","labels","breeds"]
+    def _maybe_set_classes_from_meta(self, meta: dict) -> bool:
+        keys = ["classes", "class_names", "labels", "breeds"]
         for k in keys:
             if k in meta and isinstance(meta[k], (list, tuple)) and len(meta[k]) > 1:
                 self.class_names = list(meta[k])
                 self.num_classes = len(self.class_names)
                 return True
         if "class_to_idx" in meta and isinstance(meta["class_to_idx"], dict):
-            inv = {v:k for k,v in meta["class_to_idx"].items()}
+            inv = {v: k for k, v in meta["class_to_idx"].items()}
             self.class_names = [inv[i] for i in range(len(inv))]
             self.num_classes = len(self.class_names)
             return True
@@ -210,11 +225,12 @@ class IndianBovineClassifier:
             nc = ckpt.get("num_classes", self.num_classes)
             state = ckpt.get("model_state_dict", ckpt.get("state_dict"))
             if state is None and all(isinstance(k, str) for k in ckpt.keys()):
-                state = ckpt
+                state = ckpt  # raw state dict
             if state is None:
                 raise ValueError("No state_dict in checkpoint.")
             state = strip_module_prefix(state)
             model = self._build_arch(arch or "efficientnet_b0", nc)
+            # ensure classifier head matches
             if hasattr(model, "classifier") and isinstance(model.classifier, nn.Sequential):
                 last = model.classifier[-1]
                 if isinstance(last, nn.Linear) and last.out_features != nc:
@@ -226,6 +242,7 @@ class IndianBovineClassifier:
             self.model = model.to(self.device).eval()
             self.model_type = f"pytorch:{arch or 'tv_efficientnet_b0'}"
         else:
+            # direct serialized torch.nn.Module
             self.model = ckpt.to(self.device).eval()
             self.model_type = "pytorch:serialized"
 
@@ -237,44 +254,67 @@ class IndianBovineClassifier:
             self.model = obj
             self.model_type = "sklearn"
         else:
-            raise ValueError("Unsupported object in file.")
+            raise ValueError("Unsupported object in file (expect torch module/state_dict or sklearn estimator).")
 
-    def _load_from_path(self, path, kind="auto"):
+    def _load_from_path(self, path: str, kind: str = "auto"):
         ext = os.path.splitext(path)[1].lower()
         if kind == "auto":
-            kind = "pytorch" if ext == ".pth" else ("pickle" if ext == ".pkl" else ("joblib" if ext == ".joblib" else "pytorch"))
-        if kind == "pytorch":
-            ckpt = torch.load(path, map_location=self.device)
-            self._load_pytorch_checkpoint(ckpt); return
-        if kind == "pickle":
+            if ext in {".pth"}:
+                kind = "pytorch"
+            elif ext in {".pkl"}:
+                kind = "pickle"
+            elif ext in {".joblib"}:
+                kind = "joblib"
+            else:
+                kind = "pytorch"
+
+        if kind in ("pytorch", "pickle"):
+            # Prefer torch.load first for torch checkpoints, even if extension is .pkl
             try:
                 ckpt = torch.load(path, map_location=self.device)
-                self._load_pytorch_checkpoint(ckpt); return
-            except Exception:
-                pass
-            with open(path, "rb") as f:
-                obj = pickle.load(f)
-            self._load_generic_object(obj); return
+                self._load_pytorch_checkpoint(ckpt)
+                return
+            except Exception as torch_err:
+                if kind == "pytorch":
+                    raise RuntimeError(f"PyTorch load failed: {torch_err}") from torch_err
+                # try sklearn-style pickle below
+
+            # sklearn pickle fallback
+            try:
+                with open(path, "rb") as f:
+                    obj = pickle.load(f)
+                self._load_generic_object(obj)
+                return
+            except pickle.UnpicklingError as pe:
+                # Likely a torch checkpoint mislabeled as .pkl
+                raise RuntimeError(
+                    "This .pkl appears to be a PyTorch checkpoint; load via torch.load or rename to .pth."
+                ) from pe
+
         if kind == "joblib":
             obj = joblib.load(path)
-            self._load_generic_object(obj); return
+            self._load_generic_object(obj)
+            return
+
         raise ValueError(f"Unknown model kind: {kind}")
 
-    def load_user_model(self, file_obj):
+    # public API for UI
+    def load_user_model(self, file_obj) -> str:
         path = file_to_path(file_obj)
         self._load_from_path(path, kind="auto")
         return f"✅ Loaded model: {os.path.basename(path)} | Type: {self.model_type} | Classes: {self.num_classes}"
 
-    def load_classes_json(self, file_obj):
+    def load_classes_json(self, file_obj) -> str:
         path = file_to_path(file_obj)
         with open(path, "r", encoding="utf-8") as f:
             names = json.load(f)
         if not isinstance(names, list) or len(names) < 2:
-            raise ValueError("classes.json must be a list of class names.")
+            raise ValueError("classes.json must be a list with 2 or more class names.")
         self.class_names = list(names)
         self.num_classes = len(names)
         return f"✅ Loaded {len(names)} class names from {os.path.basename(path)}"
 
+    # inference
     def preprocess_img(self, image: Image.Image):
         if image.mode != "RGB":
             image = image.convert("RGB")
@@ -282,18 +322,18 @@ class IndianBovineClassifier:
             x = self.preprocess(image).unsqueeze(0).to(self.device)
             return x
         else:
-            arr = np.array(image.resize((224,224))).astype(np.float32)/255.0
-            return arr.flatten().reshape(1,-1)
+            arr = np.array(image.resize((224, 224))).astype(np.float32) / 255.0
+            return arr.flatten().reshape(1, -1)
 
-    def predict(self, image: Image.Image):
+    def predict(self, image: Image.Image) -> Tuple[Dict[str, float], str]:
         if self.model is None:
-            return {"Error":"Model not loaded"}, "Unknown"
+            return {"Error": "Model not loaded"}, "Unknown"
         x = self.preprocess_img(image)
         if self.model_type.startswith("pytorch") or self.model_type == "demo":
             with torch.no_grad():
                 if self.model_type == "demo":
                     np.random.seed(hash(str(image.size)) % (2**32))
-                    probs = np.random.dirichlet(np.ones(self.num_classes)*3.0)
+                    probs = np.random.dirichlet(np.ones(self.num_classes) * 3.0)
                 else:
                     logits = self.model(x)
                     probs = F.softmax(logits, dim=1).cpu().numpy()[0]
@@ -301,11 +341,14 @@ class IndianBovineClassifier:
             probs = self.model.predict_proba(x)[0]
         else:
             np.random.seed(42)
-            probs = np.random.dirichlet(np.ones(self.num_classes)*2.0)
+            probs = np.random.dirichlet(np.ones(self.num_classes) * 2.0)
         top_idx = np.argsort(probs)[::-1][:3]
         results = {f"Top {i+1}: {self.class_names[idx]}": float(probs[idx]) for i, idx in enumerate(top_idx)}
         return results, self.class_names[top_idx[0]]
 
+# ---------------------------
+# UI callbacks
+# ---------------------------
 classifier = IndianBovineClassifier()
 
 def classify_image(image: Image.Image):
@@ -323,12 +366,13 @@ def classify_image(image: Image.Image):
             "Error occurred during classification",
             f"| Attribute | Value |\n|-----------|-------|\n| Status | Error: {msg} |",
         )
-    indicator = "🎲 DEMO - " if classifier.model_type == "demo" else f"🔥 {classifier.model_type} - "
+    indicator = "DEMO - " if classifier.model_type == "demo" else f"{classifier.model_type} - "
     md = f"{indicator}Classification Results:\n\n"
     for k, v in preds.items():
         md += f"- {k}: {v:.2%}\n"
     if classifier.model_type == "demo":
         md += "\nDemo mode: Upload a .pth/.pkl/.joblib model for real predictions."
+
     if top_breed in BREED_INFO:
         info = BREED_INFO[top_breed]
         desc = f"""
@@ -372,27 +416,32 @@ def upload_classes(file_obj):
     except Exception as e:
         return f"❌ Failed to load classes.json: {e}"
 
-# Minimal, responsive CSS (optional)
+# ---------------------------
+# Minimal, responsive CSS
+# ---------------------------
 CUSTOM_CSS = """
 .gradio-container { min-height: 100vh; }
 .header { text-align:center; padding: 1rem; }
 .header .title { font-size: 2em; font-weight: 700; }
 .footer { text-align:center; opacity:.75; padding:.75rem; }
 @media (max-width: 768px) {
-  .title { font-size: 1.5em !important; }
+  .title { font-size: 1.6em !important; }
 }
 """
 
+# ---------------------------
+# Interface
+# ---------------------------
 def create_interface():
-    with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft(), fill_width=True, title="🐄 Indian Bovine Classifier") as app:
+    with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft(), fill_width=True, title="Indian Bovine Classifier") as app:
         gr.HTML(f"""
             <div class="header">
-                <div class="title">🐄 Indian Bovine Breeds Classifier 🐃</div>
+                <div class="title">Indian Bovine Breeds Classifier</div>
                 <div>PyTorch runtime • {len(DEFAULT_CLASSES)} default classes • Device: {classifier.device}</div>
             </div>
         """)
 
-        # Collapsible sidebar for loaders
+        # Collapsible sidebar
         with gr.Sidebar():
             gr.Markdown("### Model loader")
             model_file = gr.File(label="Upload .pth / .pkl / .joblib", file_types=[".pth",".pkl",".joblib"], file_count="single")
@@ -406,22 +455,23 @@ def create_interface():
             classes_status = gr.Markdown("No external classes.json loaded.")
             load_classes_btn.click(upload_classes, inputs=[classes_file], outputs=[classes_status])
 
-        # Main responsive canvas
+        # Main canvas
         with gr.Row(equal_height=True):
             with gr.Column(scale=1, min_width=320, variant="panel"):
                 gr.Markdown("### Upload image")
                 image_input = gr.Image(type="pil", label="Cattle/Buffalo image")
-                classify_btn = gr.Button("🔍 Classify", variant="secondary")
+                classify_btn = gr.Button("Classify", variant="secondary")
             with gr.Column(scale=2, min_width=360, variant="panel"):
                 with gr.Tab("Results"):
-                    prediction_output = gr.Markdown(value="🔄 Upload an image to see classification.")
+                    prediction_output = gr.Markdown(value="Upload an image to see classification.")
                 with gr.Tab("Breed info"):
-                    breed_info_output = gr.Markdown(value="ℹ️ Breed info will appear here.")
+                    breed_info_output = gr.Markdown(value="Breed info will appear here.")
                 with gr.Tab("Stats"):
                     breed_stats_table = gr.Markdown(value="| Attribute | Value |\n|-----------|-------|\n| Status | Awaiting classification... |")
 
         gr.Markdown(f"""<div class="footer">Model type: {classifier.model_type} • PyTorch {torch.__version__}</div>""")
 
+        # Wiring
         classify_btn.click(classify_image, inputs=[image_input], outputs=[prediction_output, breed_info_output, breed_stats_table])
         image_input.change(classify_image, inputs=[image_input], outputs=[prediction_output, breed_info_output, breed_stats_table])
 
@@ -429,4 +479,7 @@ def create_interface():
 
 if __name__ == "__main__":
     app = create_interface()
-    app.launch(server_name="0.0.0.0", server_port=7860)
+    # Launch controls via env vars (optional)
+    share_flag = os.getenv("GRADIO_SHARE", "0").lower() in {"1", "true", "yes"}
+    ssr_flag = os.getenv("GRADIO_SSR_MODE", "true").lower() in {"1", "true", "yes"}
+    app.launch(server_name="0.0.0.0", server_port=7860, share=share_flag, ssr_mode=ssr_flag)