Create app.py

app.py

@@ -0,0 +1,135 @@
+import os, json, io
+import gradio as gr
+import numpy as np
+from PIL import Image
+import torch, torch.nn as nn
+from torchvision import models, transforms, datasets
+from cam_utils import grad_cam
+
+# -------- Config --------
+MODEL_NAME = os.environ.get("MODEL_NAME", "efficientnet_b0")
+NUM_CLASSES = int(os.environ.get("NUM_CLASSES", "2"))
+IMAGE_SIZE = int(os.environ.get("IMAGE_SIZE", "224"))
+WEIGHTS_PATH = os.environ.get("WEIGHTS_PATH", "checkpoints/best.pt")  # you will upload this file
+CLASS_NAMES = ["Parasitized", "Uninfected"] if NUM_CLASSES == 2 else [str(i) for i in range(NUM_CLASSES)]
+
+# -------- Model loading --------
+def build_model(name: str, num_classes: int):
+    name = name.lower()
+    if name == "efficientnet_b0":
+        m = models.efficientnet_b0(weights=None)
+        m.classifier[1] = nn.Linear(m.classifier[1].in_features, num_classes)
+        return m
+    elif name == "resnet50":
+        m = models.resnet50(weights=None)
+        m.fc = nn.Linear(m.fc.in_features, num_classes)
+        return m
+    else:
+        raise ValueError(f"Unsupported model_name: {name}")
+
+_device = "cuda" if torch.cuda.is_available() else "cpu"
+_model = build_model(MODEL_NAME, NUM_CLASSES).to(_device)
+
+if not os.path.exists(WEIGHTS_PATH):
+    raise FileNotFoundError(
+        f"Missing weights at {WEIGHTS_PATH}. "
+        "Upload your trained file (e.g., checkpoints/best.pt) to the Space repo."
+    )
+_state = torch.load(WEIGHTS_PATH, map_location=_device)
+_model.load_state_dict(_state)
+_model.eval()
+
+# -------- Inference helpers --------
+_pre = transforms.Compose([
+    transforms.Resize(int(IMAGE_SIZE*1.15)),
+    transforms.CenterCrop(IMAGE_SIZE),
+    transforms.ToTensor(),
+])
+
+def predict(image: Image.Image, show_cam: bool):
+    if image is None:
+        return None, None, None
+    img = image.convert("RGB")
+    x = _pre(img).unsqueeze(0).to(_device)
+    with torch.no_grad():
+        logits = _model(x).cpu().numpy().squeeze()
+    probs = np.exp(logits - logits.max()); probs = probs / probs.sum()
+    pred_idx = int(np.argmax(probs))
+    label = CLASS_NAMES[pred_idx]
+    probs_dict = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}
+
+    overlay = None
+    if show_cam:
+        cam = grad_cam(_model, img, img_size=IMAGE_SIZE, device=_device)
+        overlay = Image.fromarray((cam["overlay"]*255).astype("uint8"))
+
+    return label, probs_dict, overlay
+
+# -------- Optional: quick validation on an uploaded ZIP of a val folder --------
+# Structure expected: root_dir/ClassA/*.png, root_dir/ClassB/*.png, ...
+def validate(zip_or_folder):
+    import tempfile, zipfile, shutil
+    if zip_or_folder is None:
+        return "Upload a .zip of your validation set."
+
+    tmp = tempfile.mkdtemp()
+    path = zip_or_folder.name
+    if path.endswith(".zip"):
+        with zipfile.ZipFile(zip_or_folder.name, 'r') as zf:
+            zf.extractall(tmp)
+        root = tmp
+    else:
+        root = path  # folder
+
+    ds = datasets.ImageFolder(root, transform=_pre)
+    dl = torch.utils.data.DataLoader(ds, batch_size=64, shuffle=False, num_workers=2)
+    ys, ps = [], []
+    with torch.no_grad():
+        for xb, yb in dl:
+            preds = _model(xb.to(_device)).argmax(1).cpu().numpy()
+            ys.extend(yb.numpy()); ps.extend(preds)
+    import sklearn.metrics as sk
+    rep = sk.classification_report(ys, ps, target_names=ds.classes, output_dict=True)
+    cm  = sk.confusion_matrix(ys, ps)
+
+    # render small CM image
+    import matplotlib.pyplot as plt
+    import seaborn as sns
+    fig, ax = plt.subplots(figsize=(4.5,4))
+    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
+                xticklabels=ds.classes, yticklabels=ds.classes, ax=ax)
+    ax.set_xlabel("Predicted"); ax.set_ylabel("True"); fig.tight_layout()
+    buf = io.BytesIO(); fig.savefig(buf, format="png", dpi=160); buf.seek(0)
+    cm_img = buf
+
+    return json.dumps(rep, indent=2), cm_img
+
+# -------- Gradio UI --------
+with gr.Blocks(title="Malaria Diagnostic Assistant") as demo:
+    gr.Markdown("# 🩸 Malaria Diagnostic Assistant")
+    gr.Markdown("Prototype — energy-efficient triage with human-in-the-loop (Adaptive Sparse Training)")
+
+    with gr.Tab("🔍 Inference"):
+        with gr.Row():
+            with gr.Column():
+                img_in = gr.Image(type="pil", label="Upload blood smear image")
+                show_cam = gr.Checkbox(value=True, label="Show Grad-CAM")
+                btn = gr.Button("Predict", variant="primary")
+            with gr.Column():
+                label_out = gr.Label(num_top_classes=2, label="Prediction & Probabilities")
+                cam_out = gr.Image(type="pil", label="Grad-CAM overlay")
+        btn.click(fn=predict, inputs=[img_in, show_cam], outputs=[label_out, label_out, cam_out])
+
+    with gr.Tab("✅ Validation (optional)"):
+        gr.Markdown("Upload a **.zip** containing a folder with class subfolders (e.g., `Parasitized/`, `Uninfected/`).")
+        val_zip = gr.File(label="Validation ZIP", file_types=[".zip"])
+        run_eval = gr.Button("Compute report + confusion matrix")
+        rep_out = gr.Textbox(label="classification_report (JSON)")
+        cm_img = gr.Image(type="filepath", label="Confusion Matrix")
+        run_eval.click(fn=validate, inputs=[val_zip], outputs=[rep_out, cm_img])
+
+    gr.Markdown("---")
+    gr.Markdown("Built with EfficientNet-B0 + Adaptive Sparse Training (AST) — not a diagnostic device.")
+
+if __name__ == "__main__":
+    demo.launch()