Add GradCAM target layer selection and UI defaults integration

config/ui_defaults.json

@@ -0,0 +1,18 @@
+{
+  "selected_case": "Multiclass (4 Classes)",
+  "confidence_threshold": 0.6,
+  "smoothgrad_samples": 50,
+  "smoothgrad_noise": 0.05,
+  "gradcam_target_layer": "denseblock3",
+  "save_xdl_results": false,
+  "save_xdl_dir": "xdl_results",
+  "_comment_xdl_target_layer": "UI dropdown reads this default. You can still override with XDL_TARGET_LAYER env var.",
+  "_supported_xdl_target_layer": {
+    "denseblock3": "Default. Usually less center-biased and more spatially varied.",
+    "transition2": "Good alternative; often broad and stable localization.",
+    "transition1": "Earlier layer; more detail but can be noisy.",
+    "denseblock4": "Late layer; stronger class semantics, can be center-heavy.",
+    "transition3": "Late transition; similar tradeoff to denseblock4.",
+    "norm5_last": "Original last layer behavior (legacy setting)."
+  }
+}

requirements.txt

@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv export --format requirements-txt
+#    uv export --format requirements-txt --output-file requirements.txt
 aiofiles==24.1.0 \
     --hash=sha256:22a075c9e5a3810f0c2e48f3008c94d68c65d763b9b03857924c99e57355166c \
     --hash=sha256:b4ec55f4195e3eb5d7abd1bf7e061763e864dd4954231fb8539a0ef8bb8260e5

src/inference.py

@@ -39,6 +39,15 @@ DEFAULT_CASE_NAME = "Multiclass (4 Classes)"
 CASE_OPTIONS = list(CASE_CONFIGS.keys())
 SUPPORTED_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"}
 DEFAULT_SAVE_DIR = "xdl_results"
+GRADCAM_TARGET_LAYER_OPTIONS = (
+    "denseblock3",
+    "transition2",
+    "transition1",
+    "denseblock4",
+    "transition3",
+    "norm5_last",
+)
+DEFAULT_GRADCAM_TARGET_LAYER = "denseblock3"
 
 
 def _detect_device() -> torch.device:
@@ -201,7 +210,8 @@ def _aggregate_classification(classified: List[ClassifiedPrediction], labels: Li
 
 
 def _predict_top1(model: DenseNet121, image: Image.Image) -> Tuple[int, float, torch.Tensor]:
-    input_tensor = val_transform(image).unsqueeze(0).to(DEVICE)
+    model_device = next(model.parameters()).device
+    input_tensor = val_transform(image).unsqueeze(0).to(model_device)
 
     with torch.no_grad():
         logits = model(input_tensor)[0]
@@ -368,6 +378,7 @@ def batch_predict_with_xdl(
     smoothgrad_noise: float,
     save_xdl_results: bool,
     save_xdl_dir: str,
+    gradcam_target_layer: str = DEFAULT_GRADCAM_TARGET_LAYER,
 ):
     last_output = None
     for payload in batch_predict_with_xdl_stream(
@@ -379,6 +390,7 @@ def batch_predict_with_xdl(
         smoothgrad_noise=smoothgrad_noise,
         save_xdl_results=save_xdl_results,
         save_xdl_dir=save_xdl_dir,
+        gradcam_target_layer=gradcam_target_layer,
     ):
         last_output = payload
 
@@ -396,6 +408,7 @@ def batch_predict_with_xdl_stream(
     smoothgrad_noise: float,
     save_xdl_results: bool,
     save_xdl_dir: str,
+    gradcam_target_layer: str = DEFAULT_GRADCAM_TARGET_LAYER,
 ) -> Iterator[Tuple[str, List[List[str]], List[Tuple[np.ndarray, str]]]]:
     case_state = _get_case_state(selected_case)
     model: DenseNet121 = case_state["model"]
@@ -408,9 +421,14 @@ def batch_predict_with_xdl_stream(
         yield _render_error_html(model_error), [], []
         return
 
+    model_device = next(model.parameters()).device
+
     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))
+    gradcam_target_layer = str(gradcam_target_layer or DEFAULT_GRADCAM_TARGET_LAYER).strip().lower()
+    if gradcam_target_layer not in GRADCAM_TARGET_LAYER_OPTIONS:
+        gradcam_target_layer = DEFAULT_GRADCAM_TARGET_LAYER
 
     image_paths, input_error = _resolve_input_images(uploaded_files, folder_path)
     if input_error:
@@ -445,7 +463,7 @@ def batch_predict_with_xdl_stream(
         final_class, mean_conf = _aggregate_classification(classified, labels)
         class_counter = Counter(item.pred_idx for item in classified)
         class_stats = ", ".join(f"{labels[idx]}: {count}" for idx, count in class_counter.items())
-        initial_xdl_status = "Processing overlays..."
+        initial_xdl_status = f"Processing overlays... (GradCAM layer: {gradcam_target_layer})"
     else:
         final_class = "N/A"
         mean_conf = None
@@ -455,7 +473,7 @@ def batch_predict_with_xdl_stream(
     summary_initial = _render_summary_html(
         case_name=case_name,
         model_path=model_path,
-        device_name=DEVICE.type,
+        device_name=model_device.type,
         processed=len(image_paths),
         classified=len(classified),
         threshold=threshold,
@@ -483,13 +501,17 @@ def batch_predict_with_xdl_stream(
     xdl_error_count = 0
 
     if xdl is not None:
-        target_layer = xdl["_get_target_layer"](model)
+        try:
+            target_layer = xdl["_get_target_layer"](model, layer_name=gradcam_target_layer)
+        except TypeError:
+            # Backward compatibility for older helper signature: _get_target_layer(model)
+            target_layer = xdl["_get_target_layer"](model)
         cam = xdl["GradCAM"](model=model, target_layers=[target_layer])
 
         for item in classified:
             try:
                 image = Image.open(item.path).convert("RGB")
-                input_tensor = val_transform(image).unsqueeze(0).to(DEVICE)
+                input_tensor = val_transform(image).unsqueeze(0).to(model_device)
 
                 base_img_float, base_img_uint8 = xdl["_preprocess_image"](input_tensor[0])
                 h, w = base_img_uint8.shape[:2]
@@ -506,6 +528,7 @@ def batch_predict_with_xdl_stream(
                     item.pred_idx,
                     n_samples=smoothgrad_samples,
                     noise_level=smoothgrad_noise,
+                    use_amp=(model_device.type == "cuda"),
                 )
                 _, smooth_heatmap = xdl["_process_smoothgrad_map"](
                     smooth_raw,
@@ -532,7 +555,7 @@ def batch_predict_with_xdl_stream(
     elif xdl_error_count:
         xdl_status = f"Completed with {xdl_error_count} overlay errors"
     else:
-        xdl_status = f"Completed ({len(gallery_items)} overlays)"
+        xdl_status = f"Completed ({len(gallery_items)} overlays, layer: {gradcam_target_layer})"
 
     if save_xdl_results:
         if save_error:
@@ -547,7 +570,7 @@ def batch_predict_with_xdl_stream(
     summary_final = _render_summary_html(
         case_name=case_name,
         model_path=model_path,
-        device_name=DEVICE.type,
+        device_name=model_device.type,
         processed=len(image_paths),
         classified=len(classified),
         threshold=threshold,

src/ui.py

@@ -1,59 +1,237 @@
+import json
+from pathlib import Path
+from typing import Any
+
 import gradio as gr
 
-from src.inference import CASE_OPTIONS, DEFAULT_CASE_NAME, DEVICE, batch_predict_with_xdl_stream
+from src.inference import (
+    CASE_OPTIONS,
+    DEFAULT_CASE_NAME,
+    DEFAULT_GRADCAM_TARGET_LAYER,
+    DEVICE,
+    GRADCAM_TARGET_LAYER_OPTIONS,
+    batch_predict_with_xdl_stream,
+)
+
+UI_DEFAULTS_PATH = Path(__file__).resolve().parent.parent / "config" / "ui_defaults.json"
+UI_DEFAULTS_FALLBACK = {
+    "selected_case": DEFAULT_CASE_NAME,
+    "confidence_threshold": 0.60,
+    "smoothgrad_samples": 50,
+    "smoothgrad_noise": 0.05,
+    "gradcam_target_layer": DEFAULT_GRADCAM_TARGET_LAYER,
+    "save_xdl_results": False,
+    "save_xdl_dir": "xdl_results",
+}
+
+GRADCAM_TARGET_LAYER_DROPDOWN_CHOICES = [
+    ("DenseBlock 3 (Default, balanced)", "denseblock3"),
+    ("Transition 2 (Broad, stable)", "transition2"),
+    ("Transition 1 (Earlier, detailed/noisier)", "transition1"),
+    ("DenseBlock 4 (Late, center-heavy)", "denseblock4"),
+    ("Transition 3 (Late, center-heavy)", "transition3"),
+    ("Norm5 Last (Legacy behavior)", "norm5_last"),
+]
+
+CUSTOM_CSS = """
+.app-shell {
+    max-width: 1120px;
+    margin: 0 auto;
+}
+.hero {
+    border: 1px solid #d1d5db;
+    background: linear-gradient(135deg, #f0fdfa 0%, #ecfeff 45%, #f8fafc 100%);
+    border-radius: 14px;
+    padding: 16px 18px;
+    margin-bottom: 12px;
+}
+.hero h1 {
+    margin: 0;
+    font-size: 24px;
+    color: #0f172a;
+}
+.hero p {
+    margin: 6px 0 0 0;
+    color: #334155;
+    font-size: 14px;
+}
+.panel {
+    border: 1px solid #e2e8f0;
+    border-radius: 12px;
+    background: #ffffff;
+    padding: 12px;
+}
+"""
+
+
+def _as_float(value: Any, fallback: float) -> float:
+    try:
+        return float(value)
+    except (TypeError, ValueError):
+        return float(fallback)
+
+
+def _as_int(value: Any, fallback: int) -> int:
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return int(fallback)
+
+
+def _as_bool(value: Any, fallback: bool) -> bool:
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        return value.strip().lower() in {"1", "true", "yes", "y", "on"}
+    if value is None:
+        return fallback
+    return bool(value)
+
+
+def _load_ui_defaults() -> dict[str, Any]:
+    defaults = dict(UI_DEFAULTS_FALLBACK)
+
+    try:
+        raw_text = UI_DEFAULTS_PATH.read_text(encoding="utf-8")
+        raw = json.loads(raw_text)
+        if isinstance(raw, dict):
+            for key in defaults:
+                if key in raw:
+                    defaults[key] = raw[key]
+    except Exception:
+        pass
+
+    selected_case = str(defaults.get("selected_case", DEFAULT_CASE_NAME))
+    defaults["selected_case"] = selected_case if selected_case in CASE_OPTIONS else DEFAULT_CASE_NAME
+    defaults["confidence_threshold"] = min(
+        1.0,
+        max(0.0, _as_float(defaults.get("confidence_threshold"), UI_DEFAULTS_FALLBACK["confidence_threshold"])),
+    )
+    defaults["smoothgrad_samples"] = max(
+        1,
+        _as_int(defaults.get("smoothgrad_samples"), UI_DEFAULTS_FALLBACK["smoothgrad_samples"]),
+    )
+    defaults["smoothgrad_noise"] = min(
+        1.0,
+        max(0.0, _as_float(defaults.get("smoothgrad_noise"), UI_DEFAULTS_FALLBACK["smoothgrad_noise"])),
+    )
+    gradcam_target_layer = str(defaults.get("gradcam_target_layer", DEFAULT_GRADCAM_TARGET_LAYER)).strip().lower()
+    defaults["gradcam_target_layer"] = (
+        gradcam_target_layer if gradcam_target_layer in GRADCAM_TARGET_LAYER_OPTIONS else DEFAULT_GRADCAM_TARGET_LAYER
+    )
+    defaults["save_xdl_results"] = _as_bool(defaults.get("save_xdl_results"), UI_DEFAULTS_FALLBACK["save_xdl_results"])
+    defaults["save_xdl_dir"] = str(defaults.get("save_xdl_dir") or UI_DEFAULTS_FALLBACK["save_xdl_dir"])
+    return defaults
+
+
+def _toggle_save_dir(enabled: bool):
+    is_enabled = bool(enabled)
+    return gr.update(visible=is_enabled, interactive=is_enabled)
 
 
 def build_demo() -> gr.Blocks:
-    with gr.Blocks(title="Medical Batch Classification + XDL") as demo:
-        gr.Markdown("## Medical Batch Classification + XDL")
-        gr.Markdown(
-            f"Detected device: **{DEVICE.type}**. "
-            "Upload a folder (preferred) or provide a local folder path."
-        )
+    ui_defaults = _load_ui_defaults()
 
-        selected_case = gr.Dropdown(
-            choices=CASE_OPTIONS,
-            value=DEFAULT_CASE_NAME,
-            label="Problem Case",
+    with gr.Blocks(title="XDL Colitis Demo") as demo:
+        gr.HTML(
+            f"""
+            <style>{CUSTOM_CSS}</style>
+            <div class="app-shell">
+              <div class="hero">
+                <h1>XDL Colitis Workbench</h1>
+                <p>Detected device: <b>{DEVICE.type}</b>. Upload a directory or enter a local folder path, then run batch inference.</p>
+              </div>
+            </div>
+            """
         )
 
-        with gr.Row():
-            upload_input = gr.File(
-                file_count="directory",
-                file_types=["image"],
-                type="filepath",
-                label="Upload Image Folder",
-            )
-            folder_path = gr.Textbox(
-                label="Local Folder Path (Optional)",
-                placeholder="/absolute/path/to/folder/with/images",
-            )
-
-        with gr.Row():
-            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")
-
-        with gr.Row():
-            save_xdl_results = gr.Checkbox(label="Save XDL Results Locally", value=False)
-            save_xdl_dir = gr.Textbox(
-                label="Save Folder",
-                value="xdl_results",
-                placeholder="xdl_results",
-            )
-
-        run_btn = gr.Button("Run Batch Inference")
-
-        summary_out = gr.HTML(label="Summary")
-        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="Compact XDL Results (Original | GradCAM | SmoothGrad)",
-            columns=2,
+        with gr.Row(elem_classes=["app-shell"]):
+            with gr.Column(scale=2, elem_classes=["panel"]):
+                gr.Markdown("### 1) Image Input")
+                selected_case = gr.Dropdown(
+                    choices=CASE_OPTIONS,
+                    value=ui_defaults["selected_case"],
+                    label="Problem Case",
+                    info="Choose the model group that matches your diagnosis scenario.",
+                )
+                upload_input = gr.File(
+                    file_count="directory",
+                    file_types=["image"],
+                    type="filepath",
+                    label="Upload Image Folder",
+                )
+                folder_path = gr.Textbox(
+                    label="Local Folder Path (Optional)",
+                    placeholder="/absolute/path/to/folder/with/images",
+                )
+
+            with gr.Column(scale=1, elem_classes=["panel"]):
+                gr.Markdown("### 2) Inference Settings")
+                threshold = gr.Number(
+                    value=ui_defaults["confidence_threshold"],
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    precision=2,
+                    label="Confidence Threshold",
+                    info="Range: 0.00 to 1.00",
+                )
+                smoothgrad_samples = gr.Number(
+                    value=ui_defaults["smoothgrad_samples"],
+                    minimum=1,
+                    maximum=1000,
+                    step=1,
+                    precision=0,
+                    label="SmoothGrad Samples",
+                    info="Higher values improve stability but increase runtime.",
+                )
+                smoothgrad_noise = gr.Number(
+                    value=ui_defaults["smoothgrad_noise"],
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    precision=2,
+                    label="SmoothGrad Noise Level",
+                    info="Typical range: 0.01 to 0.20",
+                )
+                gradcam_target_layer = gr.Dropdown(
+                    choices=GRADCAM_TARGET_LAYER_DROPDOWN_CHOICES,
+                    value=ui_defaults["gradcam_target_layer"],
+                    label="GradCAM Target Layer",
+                    info="Try `transition2` or `denseblock3` if CAM looks too centered.",
+                )
+                save_xdl_results = gr.Checkbox(
+                    label="Save XDL Results Locally",
+                    value=ui_defaults["save_xdl_results"],
+                )
+                save_xdl_dir = gr.Textbox(
+                    label="Save Folder",
+                    value=ui_defaults["save_xdl_dir"],
+                    placeholder="xdl_results",
+                    visible=bool(ui_defaults["save_xdl_results"]),
+                    interactive=bool(ui_defaults["save_xdl_results"]),
+                )
+                run_btn = gr.Button("Run Batch Inference", variant="primary")
+
+        with gr.Row(elem_classes=["app-shell"]):
+            with gr.Column(elem_classes=["panel"]):
+                gr.Markdown("### 3) Results")
+                summary_out = gr.HTML(label="Summary")
+                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="Compact XDL Results (Original | GradCAM | SmoothGrad)",
+                    columns=2,
+                )
+
+        save_xdl_results.change(
+            fn=_toggle_save_dir,
+            inputs=[save_xdl_results],
+            outputs=[save_xdl_dir],
         )
 
         run_btn.click(
@@ -67,6 +245,7 @@ def build_demo() -> gr.Blocks:
                 smoothgrad_noise,
                 save_xdl_results,
                 save_xdl_dir,
+                gradcam_target_layer,
             ],
             outputs=[summary_out, table_out, gallery_out],
         )

src/xdl.py

@@ -1,4 +1,5 @@
-from typing import Tuple
+import os
+from typing import Optional, Tuple
 
 import cv2
 import numpy as np
@@ -47,13 +48,50 @@ def _process_smoothgrad_map(
     return smoothgrad_map, heatmap
 
 
-def _get_target_layer(model: nn.Module):
-    """Return DenseNet feature layer used for GradCAM."""
+def _get_target_layer(model: nn.Module, layer_name: Optional[str] = None):
+    """
+    Return DenseNet feature layer used for GradCAM.
+
+    Default layer is `denseblock3` to reduce center-biased CAMs compared with
+    the original `norm5_last`.
+
+    You can override with:
+    - function arg: `layer_name`
+    - env var: `XDL_TARGET_LAYER`
+
+    Supported layer names:
+    `denseblock3`, `transition2`, `transition1`, `denseblock4`, `transition3`, `norm5_last`.
+    """
     if not isinstance(model, DenseNet121):
         raise TypeError(
             f"Unsupported model type for this demo: {type(model).__name__}. Expected DenseNet121."
         )
-    return model.densenet_model.features[-1]
+
+    requested = (layer_name or os.getenv("XDL_TARGET_LAYER") or "denseblock3").strip().lower()
+    aliases = {
+        "default": "denseblock3",
+        "last": "norm5_last",
+        "norm5": "norm5_last",
+    }
+    selected = aliases.get(requested, requested)
+
+    features = model.densenet_model.features
+    layer_map = {
+        "denseblock3": features.denseblock3,
+        "transition2": features.transition2,
+        "transition1": features.transition1,
+        "denseblock4": features.denseblock4,
+        "transition3": features.transition3,
+        "norm5_last": features[-1],
+    }
+
+    if selected not in layer_map:
+        supported = ", ".join(layer_map.keys())
+        raise ValueError(
+            f"Unsupported XDL target layer '{selected}'. Supported layers: {supported}"
+        )
+
+    return layer_map[selected]
 
 
 def smoothgrad(
@@ -62,27 +100,61 @@ def smoothgrad(
     target_class: int,
     n_samples: int = 100,
     noise_level: float = 0.05,
+    batch_size: Optional[int] = None,
+    use_amp: bool = False,
 ) -> np.ndarray:
-    """Compute SmoothGrad saliency map for one input tensor of shape (1, C, H, W)."""
+    """
+    Compute SmoothGrad saliency map for one input tensor of shape (1, C, H, W).
+
+    Notes:
+    - This implementation batches noisy samples to reduce per-step overhead.
+    - If `input_tensor` and `model` are on CUDA, computation runs on GPU.
+    """
+    if n_samples <= 0:
+        raise ValueError(f"n_samples must be > 0, got {n_samples}")
+
+    if input_tensor.ndim != 4 or input_tensor.shape[0] != 1:
+        raise ValueError(
+            f"input_tensor must have shape (1, C, H, W), got {tuple(input_tensor.shape)}"
+        )
+
+    if batch_size is None:
+        batch_size = min(16, n_samples)
+    if batch_size <= 0:
+        raise ValueError(f"batch_size must be > 0, got {batch_size}")
+
+    model_param = next(model.parameters(), None)
+    if model_param is not None and model_param.device != input_tensor.device:
+        raise ValueError(
+            f"Model device ({model_param.device}) and input device ({input_tensor.device}) must match."
+        )
+
     model.eval()
     accumulated_gradients = torch.zeros_like(input_tensor)
 
     input_range = torch.max(input_tensor) - torch.min(input_tensor)
     scaled_noise = noise_level * input_range
 
-    for _ in range(n_samples):
-        noise = torch.randn_like(input_tensor) * scaled_noise
-        noisy_input = (input_tensor + noise).clone().detach().requires_grad_(True)
-
-        output = model(noisy_input)
+    use_cuda_amp = use_amp and input_tensor.device.type == "cuda"
+    for start in range(0, n_samples, batch_size):
+        current_batch = min(batch_size, n_samples - start)
+        expanded_input = input_tensor.expand(current_batch, -1, -1, -1)
+        noise = torch.randn_like(expanded_input) * scaled_noise
+        noisy_input = (expanded_input + noise).clone().detach().requires_grad_(True)
+
+        if use_cuda_amp:
+            with torch.autocast(device_type="cuda", dtype=torch.float16):
+                output = model(noisy_input)
+        else:
+            output = model(noisy_input)
         if output.ndim == 1:
             output = output.unsqueeze(0)
 
-        model.zero_grad()
-        output[0, target_class].backward()
+        model.zero_grad(set_to_none=True)
+        output[:, target_class].sum().backward()
 
         if noisy_input.grad is not None:
-            accumulated_gradients += noisy_input.grad.data
+            accumulated_gradients += noisy_input.grad.data.sum(dim=0, keepdim=True)
 
     smooth = accumulated_gradients / n_samples
     smooth = torch.abs(smooth)