Update app.py

app.py

@@ -1,6 +1,5 @@
-import os
 from pathlib import Path
-from typing import Optional, Tuple, Dict, List
+from typing import Optional, Dict, List
 
 import gradio as gr
 import torch
@@ -14,63 +13,28 @@ from torchvision.models import resnet18
 # Config
 # -----------------------------
 CIFAR10_CLASSES = [
-    "airplane","automobile","bird","cat","deer","dog","frog","horse","ship","truck"
+    "airplane", "automobile", "bird", "cat", "deer",
+    "dog", "frog", "horse", "ship", "truck"
 ]
 
-# CIFAR-10 normalization (standard)
 CIFAR10_MEAN = (0.4914, 0.4822, 0.4465)
 CIFAR10_STD  = (0.2470, 0.2435, 0.2616)
 
-EXAMPLES_DIR = Path("Examples")  # you uploaded to this folder
+EXAMPLES_DIR = Path("Examples")
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+# If you know the exact checkpoint name, lock it here:
+CKPT_PATH = Path("ast_cifar10_resnet18.pth")
+
 # -----------------------------
-# Utilities
+# Model helpers
 # -----------------------------
-def find_checkpoint(repo_root: Path) -> Optional[Path]:
-    """
-    Auto-find a checkpoint file in the repo.
-    Add your own filename here if you know it.
-    """
-    candidates = [
-        "model.pth", "model.pt",
-        "checkpoint.pth", "checkpoint.pt",
-        "best.pth", "best.pt",
-        "resnet18.pth", "resnet18.pt",
-        "weights.pth", "weights.pt",
-    ]
-
-    for name in candidates:
-        p = repo_root / name
-        if p.exists() and p.is_file():
-            return p
-
-    # Try pattern search
-    patterns = ["*.pth", "*.pt"]
-    for pat in patterns:
-        hits = sorted(repo_root.glob(pat))
-        # Prefer anything that looks like resnet/cifar/ast
-        preferred = [h for h in hits if any(k in h.name.lower() for k in ["resnet", "cifar", "ast", "sparse", "best"])]
-        if preferred:
-            return preferred[0]
-        if hits:
-            return hits[0]
-
-    return None
-
-
 def build_model(num_classes: int = 10) -> torch.nn.Module:
     m = resnet18(weights=None)
     m.fc = torch.nn.Linear(m.fc.in_features, num_classes)
     return m
 
-
 def load_weights(model: torch.nn.Module, ckpt_path: Path) -> None:
-    """
-    Loads common checkpoint formats:
-    - plain state_dict
-    - dict with 'state_dict' or 'model' keys
-    """
     ckpt = torch.load(ckpt_path, map_location="cpu")
 
     if isinstance(ckpt, dict):
@@ -79,26 +43,19 @@ def load_weights(model: torch.nn.Module, ckpt_path: Path) -> None:
         elif "model" in ckpt and isinstance(ckpt["model"], dict):
             state = ckpt["model"]
         else:
-            # might already be a state_dict-like dict
             state = ckpt
     else:
         raise ValueError(f"Unsupported checkpoint format: {type(ckpt)}")
 
-    # Strip possible 'module.' prefix if trained with DDP/DataParallel
-    new_state = {}
-    for k, v in state.items():
-        nk = k.replace("module.", "")
-        new_state[nk] = v
-
-    missing, unexpected = model.load_state_dict(new_state, strict=False)
-    # Strict=False to be robust; you can change to strict=True if you prefer.
+    # Remove "module." if saved from DDP
+    cleaned = {k.replace("module.", ""): v for k, v in state.items()}
+    missing, unexpected = model.load_state_dict(cleaned, strict=False)
     if missing or unexpected:
         print("[load_weights] Missing keys:", missing)
         print("[load_weights] Unexpected keys:", unexpected)
 
-
 # -----------------------------
-# Preprocess + Predict
+# Preprocess
 # -----------------------------
 preprocess = T.Compose([
     T.Resize((32, 32), interpolation=T.InterpolationMode.BILINEAR),
@@ -106,75 +63,68 @@ preprocess = T.Compose([
     T.Normalize(mean=CIFAR10_MEAN, std=CIFAR10_STD),
 ])
 
-def pil_to_model_tensor(img: Image.Image) -> torch.Tensor:
-    img = img.convert("RGB")
-    x = preprocess(img).unsqueeze(0)  # [1,3,32,32]
-    return x
+STATE: Dict[str, Optional[torch.nn.Module]] = {"model": None}
+
+def init():
+    if not CKPT_PATH.exists():
+        print(f"[init] Checkpoint not found: {CKPT_PATH}")
+        STATE["model"] = None
+        return
 
+    print(f"[init] Loading checkpoint: {CKPT_PATH}")
+    model = build_model(num_classes=len(CIFAR10_CLASSES))
+    load_weights(model, CKPT_PATH)
+    model.to(DEVICE).eval()
+    STATE["model"] = model
+
+def get_examples() -> List[List[str]]:
+    if not EXAMPLES_DIR.exists():
+        return []
+    imgs = sorted([p for p in EXAMPLES_DIR.iterdir() if p.suffix.lower() in [".png", ".jpg", ".jpeg"]])
+    return [[str(p)] for p in imgs]
+
+# -----------------------------
+# Predict
+# -----------------------------
 def predict(img: Image.Image):
     if img is None:
-        return None, None, None
+        return None, {}, [["", ""], ["", ""], ["", ""]], ""
 
     if STATE["model"] is None:
-        raise gr.Error("Model is not loaded. Check that your checkpoint exists in the Space repo.")
+        raise gr.Error("Model is not loaded. Ensure ast_cifar10_resnet18.pth exists in the repo root.")
 
-    # Show exactly what goes into the model (32x32)
+    # show the actual 32x32 that goes into model
     img32 = img.convert("RGB").resize((32, 32), resample=Image.BILINEAR)
 
-    x = pil_to_model_tensor(img).to(DEVICE)
+    x = preprocess(img.convert("RGB")).unsqueeze(0).to(DEVICE)  # [1,3,32,32]
     with torch.inference_mode():
         logits = STATE["model"](x)
         probs = F.softmax(logits, dim=1).squeeze(0)  # [10]
 
-    # Top-3
-    topk = torch.topk(probs, k=3)
-    top3 = [(CIFAR10_CLASSES[i], float(topk.values[j])) for j, i in enumerate(topk.indices.tolist())]
-
-    # Gradio Label expects dict label->confidence
+    # label dict for gr.Label
     label_dict = {cls: float(probs[i]) for i, cls in enumerate(CIFAR10_CLASSES)}
 
-    # Table for top-3
-    top3_table = [[name, f"{p*100:.2f}%"] for name, p in top3]
+    # top-3 table
+    topk = torch.topk(probs, k=3)
+    top3_rows = []
+    for j, idx in enumerate(topk.indices.tolist()):
+        top3_rows.append([CIFAR10_CLASSES[idx], f"{float(topk.values[j]) * 100:.2f}%"])
 
-    # Main prediction text
-    pred_name, pred_p = top3[0]
-    pred_text = f"**{pred_name}** ({pred_p*100:.2f}%)"
+    pred_name = CIFAR10_CLASSES[int(topk.indices[0])]
+    pred_conf = float(topk.values[0]) * 100.0
+    pred_text = f"**{pred_name}** ({pred_conf:.2f}%)"
 
-    return img32, label_dict, top3_table, pred_text
+    return img32, label_dict, top3_rows, pred_text
 
+def clear_all():
+    return None, None, {}, [["", ""], ["", ""], ["", ""]], ""
 
 # -----------------------------
-# App state
+# App
 # -----------------------------
-STATE: Dict[str, Optional[torch.nn.Module]] = {"model": None}
-
-def init():
-    repo_root = Path(".")
-    ckpt = find_checkpoint(repo_root)
-    if ckpt is None:
-        print("[init] No checkpoint found in repo root.")
-        STATE["model"] = None
-        return
-
-    print(f"[init] Loading checkpoint: {ckpt}")
-    model = build_model(num_classes=len(CIFAR10_CLASSES))
-    load_weights(model, ckpt)
-    model.to(DEVICE).eval()
-    STATE["model"] = model
-
-def get_examples() -> List[List[str]]:
-    if not EXAMPLES_DIR.exists():
-        return []
-    imgs = sorted([p for p in EXAMPLES_DIR.iterdir() if p.suffix.lower() in [".png", ".jpg", ".jpeg"]])
-    # Gradio expects list of lists, each inner list corresponds to inputs
-    return [[str(p)] for p in imgs]
-
 init()
 EXAMPLES = get_examples()
 
-# -----------------------------
-# UI
-# -----------------------------
 with gr.Blocks(title="AST CIFAR-10 Classifier") as demo:
     gr.Markdown(
         "# AST CIFAR-10 Classifier\n"
@@ -185,8 +135,6 @@ with gr.Blocks(title="AST CIFAR-10 Classifier") as demo:
     with gr.Row():
         with gr.Column(scale=1):
             img_in = gr.Image(type="pil", label="Upload CIFAR-like image")
-
-            # Show the exact 32×32 fed to model (useful for debugging)
             img_32 = gr.Image(type="pil", label="Model input (32×32)")
 
         with gr.Column(scale=1):
@@ -196,7 +144,7 @@ with gr.Blocks(title="AST CIFAR-10 Classifier") as demo:
                 headers=["class", "confidence"],
                 datatype=["str", "str"],
                 row_count=3,
-                col_count=(2, "fixed"),
+                column_count=2,   # <-- fixed (no deprecated col_count)
                 interactive=False,
                 label="Top-3"
             )
@@ -206,11 +154,14 @@ with gr.Blocks(title="AST CIFAR-10 Classifier") as demo:
         submit = gr.Button("Submit", variant="primary")
         clear = gr.Button("Clear")
 
+    # ✅ FIX: if cache_examples=True, you MUST provide fn and outputs
     if EXAMPLES:
         gr.Markdown("### Examples (from `Examples/` folder)")
         gr.Examples(
             examples=EXAMPLES,
             inputs=[img_in],
+            outputs=[img_32, pred_label, top3_table, pred_text],
+            fn=predict,
             cache_examples=True
         )
 
@@ -220,9 +171,11 @@ with gr.Blocks(title="AST CIFAR-10 Classifier") as demo:
         outputs=[img_32, pred_label, top3_table, pred_text]
     )
 
-    def _clear():
-        return None, None, None, ""
-    clear.click(fn=_clear, inputs=[], outputs=[img_in, img_32, top3_table, pred_text])
+    clear.click(
+        fn=clear_all,
+        inputs=[],
+        outputs=[img_in, img_32, pred_label, top3_table, pred_text]
+    )
 
 demo.queue()
 if __name__ == "__main__":