Split app into launcher, UI, and inference modules

app.py

@@ -1,69 +1,7 @@
-import gradio as gr
-import torch
-import torch.nn.functional as F
-from torchvision import transforms
-from PIL import Image
+from src.ui import build_demo
 
-# ---------------------------------------------------------
-# 1. CONFIGURATION (Edit these!)
-# ---------------------------------------------------------
-MODEL_PATH = "models/multi_smoothing1.2_reducelr_epoch116.pth"       # Your model file name
-# IMPORTANT: These must match the order of folders/classes used during training!
-LABELS = ["uc", "infeksi", "crohn", "tb"]  # Example for your Colitis demo
 
-# ---------------------------------------------------------
-# 2. LOAD MODEL
-# ---------------------------------------------------------
-# We load to CPU since Hugging Face basic tier is CPU-only
-from src.densenet import DenseNet121
-model = DenseNet121(num_classes=len(LABELS))
-
-try:
-    model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu')))
-    model.eval()  # Set to evaluation mode
-except Exception as e:
-    print(f"Error loading model: {e}")
-    print("Make sure your model file is uploaded and matches MODEL_PATH.")
-    # Create a dummy model just so the app doesn't crash immediately for testing
-
-# ---------------------------------------------------------
-# 3. PREPROCESSING
-# ---------------------------------------------------------
-# Standard preprocessing. Adjust "Resize" if your model expects 299x299 (Inception) or others.
-from src.preprocessing import preprocess
-
-target_input_size = tuple([3, 299, 299])
-
-val_transform = preprocess(target_input_size=target_input_size)
-
-# ---------------------------------------------------------
-# 4. PREDICTION FUNCTION
-# ---------------------------------------------------------
-def predict(image):
-    if model is None:
-        return {"Error": 1.0}
-
-    input_tensor = val_transform(image).unsqueeze(0)
-
-    with torch.no_grad():
-        output = model(input_tensor)
-
-        probabilities = F.softmax(output[0], dim=0)
-
-    # 4. Map to labels
-    # Returns a dictionary: {"Class A": 0.9, "Class B": 0.1}
-    return {LABELS[i]: float(probabilities[i]) for i in range(len(LABELS))}
-
-# ---------------------------------------------------------
-# 5. GRADIO INTERFACE
-# ---------------------------------------------------------
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Label(num_top_classes=len(LABELS)),
-    title="Medical Image Classification Demo",
-    description="Upload a scan to classify it."
-)
+demo = build_demo()
 
 if __name__ == "__main__":
     demo.launch()

src/inference.py

@@ -0,0 +1,258 @@
+from collections import Counter, defaultdict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+from src.densenet import DenseNet121
+from src.preprocessing import preprocess
+
+MODEL_PATH = "models/multi_smoothing1.2_reducelr_epoch116.pth"
+LABELS = ["uc", "infeksi", "crohn", "tb"]
+NUM_CLASSES = len(LABELS)
+SUPPORTED_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"}
+
+model = DenseNet121(num_classes=NUM_CLASSES)
+_model_load_error: Optional[str] = None
+
+try:
+    model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device("cpu")))
+    model.eval()
+except Exception as exc:
+    _model_load_error = str(exc)
+
+val_transform = preprocess(target_input_size=(3, 299, 299))
+
+
+@dataclass(frozen=True)
+class ClassifiedPrediction:
+    path: Path
+    pred_idx: int
+    confidence: float
+
+
+def _model_error_message() -> Optional[str]:
+    if _model_load_error is None:
+        return None
+    return f"Model failed to load from `{MODEL_PATH}`: {_model_load_error}"
+
+
+def _load_xdl_modules():
+    """Lazy-load optional XDL dependencies."""
+    try:
+        from pytorch_grad_cam import GradCAM
+        from pytorch_grad_cam.utils.image import show_cam_on_image
+        from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+    except Exception as exc:
+        raise RuntimeError(f"GradCAM import failed: {exc}")
+
+    try:
+        import cv2
+    except Exception as exc:
+        raise RuntimeError(f"OpenCV import failed: {exc}")
+
+    try:
+        from src.xdl import (
+            _get_target_layer,
+            _preprocess_image,
+            _process_smoothgrad_map,
+            smoothgrad,
+        )
+    except Exception as exc:
+        raise RuntimeError(f"Failed to load XDL utilities from src/xdl.py: {exc}")
+
+    return {
+        "cv2": cv2,
+        "GradCAM": GradCAM,
+        "show_cam_on_image": show_cam_on_image,
+        "ClassifierOutputTarget": ClassifierOutputTarget,
+        "_get_target_layer": _get_target_layer,
+        "_preprocess_image": _preprocess_image,
+        "_process_smoothgrad_map": _process_smoothgrad_map,
+        "smoothgrad": smoothgrad,
+    }
+
+
+def _iter_image_paths(folder: Path) -> List[Path]:
+    return sorted(
+        p
+        for p in folder.rglob("*")
+        if p.is_file() and p.suffix.lower() in SUPPORTED_IMAGE_EXTENSIONS
+    )
+
+
+def _aggregate_classification(classified: List[ClassifiedPrediction]) -> Tuple[str, float]:
+    class_counter = Counter(item.pred_idx for item in classified)
+    top_count = max(class_counter.values())
+    tied_classes = [idx for idx, count in class_counter.items() if count == top_count]
+
+    if len(tied_classes) == 1:
+        final_idx = tied_classes[0]
+    else:
+        class_conf = defaultdict(list)
+        for item in classified:
+            class_conf[item.pred_idx].append(item.confidence)
+        final_idx = max(tied_classes, key=lambda idx: float(np.mean(class_conf[idx])))
+
+    final_conf = float(np.mean([item.confidence for item in classified if item.pred_idx == final_idx]))
+    return LABELS[final_idx], final_conf
+
+
+def _build_visual_row(
+    original: np.ndarray,
+    gradcam_img: np.ndarray,
+    smoothgrad_img: np.ndarray,
+) -> np.ndarray:
+    return np.concatenate([original, gradcam_img, smoothgrad_img], axis=1)
+
+
+def _predict_top1(image: Image.Image) -> Tuple[int, float, torch.Tensor]:
+    input_tensor = val_transform(image).unsqueeze(0)
+
+    with torch.no_grad():
+        logits = model(input_tensor)[0]
+        probs = F.softmax(logits, dim=0)
+
+    pred_idx = int(torch.argmax(probs).item())
+    confidence = float(probs[pred_idx].item())
+    return pred_idx, confidence, input_tensor
+
+
+def predict_single(image: Image.Image) -> Dict[str, float]:
+    model_error = _model_error_message()
+    if model_error:
+        return {model_error: 1.0}
+
+    input_tensor = val_transform(image).unsqueeze(0)
+
+    with torch.no_grad():
+        output = model(input_tensor)
+        probabilities = F.softmax(output[0], dim=0)
+
+    return {LABELS[i]: float(probabilities[i]) for i in range(NUM_CLASSES)}
+
+
+def batch_predict_with_xdl(
+    folder_path: str,
+    confidence_threshold: float,
+    smoothgrad_samples: int,
+    smoothgrad_noise: float,
+):
+    model_error = _model_error_message()
+    if model_error:
+        return model_error, [], []
+
+    if not folder_path:
+        return "Provide a folder path.", [], []
+
+    folder = Path(folder_path).expanduser().resolve()
+    if not folder.exists() or not folder.is_dir():
+        return f"Invalid folder: `{folder}`", [], []
+
+    threshold = float(np.clip(confidence_threshold, 0.0, 1.0))
+    smoothgrad_samples = int(max(1, smoothgrad_samples))
+    smoothgrad_noise = float(max(0.0, smoothgrad_noise))
+
+    image_paths = _iter_image_paths(folder)
+    if not image_paths:
+        return f"No supported image files found in `{folder}`.", [], []
+
+    xdl = None
+    xdl_error = ""
+    try:
+        xdl = _load_xdl_modules()
+    except RuntimeError as exc:
+        xdl_error = str(exc)
+
+    cam = None
+    if xdl:
+        target_layer = xdl["_get_target_layer"](model)
+        cam = xdl["GradCAM"](model=model, target_layers=[target_layer])
+
+    classified: List[ClassifiedPrediction] = []
+    rows: List[List[str]] = []
+    gallery_items = []
+
+    for img_path in image_paths:
+        try:
+            image = Image.open(img_path).convert("RGB")
+        except Exception as exc:
+            rows.append([img_path.name, "error", "-", str(exc)])
+            continue
+
+        pred_idx, confidence, input_tensor = _predict_top1(image)
+
+        if confidence < threshold:
+            rows.append([img_path.name, "below_threshold", LABELS[pred_idx], f"{confidence:.4f}"])
+            continue
+
+        prediction = ClassifiedPrediction(path=img_path, pred_idx=pred_idx, confidence=confidence)
+        classified.append(prediction)
+        rows.append([img_path.name, "classified", LABELS[pred_idx], f"{confidence:.4f}"])
+
+        if xdl and cam is not None:
+            try:
+                base_img_float, base_img_uint8 = xdl["_preprocess_image"](input_tensor[0])
+                h, w = base_img_uint8.shape[:2]
+
+                grayscale_cam = cam(
+                    input_tensor=input_tensor,
+                    targets=[xdl["ClassifierOutputTarget"](pred_idx)],
+                )[0, :]
+                gradcam_overlay = xdl["show_cam_on_image"](base_img_float, grayscale_cam, use_rgb=True)
+
+                smooth_raw = xdl["smoothgrad"](
+                    model,
+                    input_tensor,
+                    pred_idx,
+                    n_samples=smoothgrad_samples,
+                    noise_level=smoothgrad_noise,
+                )
+                _, smooth_heatmap = xdl["_process_smoothgrad_map"](
+                    smooth_raw,
+                    img_shape=(h, w),
+                    percentile=95,
+                    colormap="hot",
+                )
+                smooth_heatmap_rgb = xdl["cv2"].cvtColor(smooth_heatmap, xdl["cv2"].COLOR_BGR2RGB)
+                smooth_overlay = xdl["cv2"].addWeighted(base_img_uint8, 0.6, smooth_heatmap_rgb, 0.4, 0)
+
+                panel = _build_visual_row(base_img_uint8, gradcam_overlay, smooth_overlay)
+                caption = (
+                    f"{img_path.name} | pred={LABELS[pred_idx]} | conf={confidence:.4f} "
+                    "| left=Original middle=GradCAM right=SmoothGrad"
+                )
+                gallery_items.append((panel, caption))
+            except Exception as exc:
+                rows.append([img_path.name, "xdl_error", LABELS[pred_idx], str(exc)])
+
+    if not classified:
+        summary = (
+            f"Processed {len(image_paths)} images. "
+            f"0 images met threshold `{threshold:.2f}`."
+        )
+        if xdl_error:
+            summary += f"\n\nXDL status: {xdl_error}"
+        return summary, rows, []
+
+    final_class, mean_conf = _aggregate_classification(classified)
+    class_counter = Counter(item.pred_idx for item in classified)
+    class_stats = ", ".join(f"{LABELS[idx]}: {count}" for idx, count in class_counter.items())
+
+    summary = (
+        f"Processed {len(image_paths)} images. "
+        f"Classified: {len(classified)} (threshold >= {threshold:.2f}), "
+        f"Skipped: {len(image_paths) - len(classified)}.\n"
+        f"Final class (from per-image max class vote): **{final_class}**\n"
+        f"Mean confidence for final class: **{mean_conf:.4f}**\n"
+        f"Classified image distribution: {class_stats}"
+    )
+
+    if xdl_error:
+        summary += f"\n\nXDL status: {xdl_error}"
+
+    return summary, rows, gallery_items

src/ui.py

@@ -0,0 +1,44 @@
+import gradio as gr
+
+from src.inference import NUM_CLASSES, batch_predict_with_xdl, predict_single
+
+
+def build_demo() -> gr.Blocks:
+    with gr.Blocks(title="Medical Image Classification + XDL") as demo:
+        gr.Markdown("## Medical Image Classification Demo")
+
+        with gr.Tab("Single Image"):
+            single_in = gr.Image(type="pil", label="Image")
+            single_out = gr.Label(num_top_classes=NUM_CLASSES, label="Prediction")
+            single_btn = gr.Button("Predict")
+            single_btn.click(fn=predict_single, inputs=single_in, outputs=single_out)
+
+        with gr.Tab("Folder Batch + XDL"):
+            folder_path = gr.Textbox(
+                label="Folder Path",
+                placeholder="/absolute/path/to/folder/with/images",
+            )
+            threshold = gr.Slider(0.0, 1.0, value=0.6, step=0.01, label="Confidence Threshold")
+            smoothgrad_samples = gr.Slider(10, 200, value=50, step=10, label="SmoothGrad Samples")
+            smoothgrad_noise = gr.Slider(0.01, 0.2, value=0.05, step=0.01, label="SmoothGrad Noise Level")
+            run_btn = gr.Button("Run Batch Inference")
+
+            summary_out = gr.Markdown()
+            table_out = gr.Dataframe(
+                headers=["filename", "status", "predicted_label", "confidence_or_error"],
+                datatype=["str", "str", "str", "str"],
+                interactive=False,
+                label="Per-image Results",
+            )
+            gallery_out = gr.Gallery(
+                label="XDL Overlays (Original | GradCAM | SmoothGrad)",
+                columns=1,
+            )
+
+            run_btn.click(
+                fn=batch_predict_with_xdl,
+                inputs=[folder_path, threshold, smoothgrad_samples, smoothgrad_noise],
+                outputs=[summary_out, table_out, gallery_out],
+            )
+
+    return demo