Blame for src/inference.py

import base64
from collections import Counter, defaultdict
from dataclasses import dataclass
from html import escape
from io import BytesIO
import os
from pathlib import Path
from typing import Any, Iterator, List, Optional, Sequence, Tuple
import warnings

import gradio as gr
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image, ImageDraw

from src.densenet import DenseNet121
from src.preprocessing import preprocess
from src.utils import download_google_drive_images, download_google_drive_folder_images, iter_image_paths, normalize_uploaded_paths

CASE_CONFIGS = {
    "Multiclass (4 Classes)": {
        "key": "multiclass",
        "labels": ["Kolitis Ulseratif", "Infeksi", "Crohn", "Tuberculosis"],
        "default_model_path": "models/multiclass.pth",
        "env_var": "MODEL_PATH_MULTICLASS",
    },
    "Crohn vs TB": {
        "key": "crohn_tb",
        "labels": ["Crohn", "Tuberculosis"],
        "default_model_path": "models/crohn_tb.pth",
        "env_var": "MODEL_PATH_CROHN_TB",
    },
    "Kolitis Ulseratif vs Infeksi": {
        "key": "ku_infeksi",
        "labels": ["Kolitis Ulseratif", "Infeksi"],
        "default_model_path": "models/ku_infeksi.pth",
        "env_var": "MODEL_PATH_KU_INFEKSI",
    },
}

DEFAULT_CASE_NAME = "Multiclass (4 Classes)"
CASE_OPTIONS = list(CASE_CONFIGS.keys())
DEFAULT_SAVE_DIR = "xdl_results"
DEFAULT_CUDA_SMOOTHGRAD_BATCH_SIZE = 32
GRADCAM_TARGET_LAYER_OPTIONS = (
    "denseblock3",
    "transition2",
    "transition1",
    "denseblock4",
    "transition3",
    "norm5_last",
)
DEFAULT_GRADCAM_TARGET_LAYER = "denseblock3"


def _detect_device() -> torch.device:
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", message="CUDA initialization:.*")
        try:
            has_cuda = torch.cuda.is_available()
        except Exception:
            has_cuda = False
    return torch.device("cuda" if has_cuda else "cpu")


DEVICE = _detect_device()
if DEVICE.type == "cuda":
    torch.backends.cudnn.benchmark = True
_case_state_cache: dict[str, dict] = {}

val_transform = preprocess(target_input_size=(3, 299, 299))


@dataclass(frozen=True)
class ClassifiedPrediction:
    path: Path
    pred_idx: int
    confidence: float


@dataclass(frozen=True)
class XdlVisualItem:
    filename: str
    predicted_label: str
    confidence: float
    original: np.ndarray
    gradcam: np.ndarray
    smoothgrad: np.ndarray


def _get_case_config(case_name: str) -> dict:
    return CASE_CONFIGS.get(case_name, CASE_CONFIGS[DEFAULT_CASE_NAME])


def _resolve_case_model_path(case_cfg: dict) -> str:
    return os.getenv(case_cfg["env_var"], case_cfg["default_model_path"])


def _smoothgrad_batch_size(device: torch.device) -> Optional[int]:
    if device.type != "cuda":
        return None

    raw_value = os.getenv("XDL_SMOOTHGRAD_BATCH_SIZE", str(DEFAULT_CUDA_SMOOTHGRAD_BATCH_SIZE))
    try:
        return max(1, int(raw_value))
    except ValueError:
        return DEFAULT_CUDA_SMOOTHGRAD_BATCH_SIZE


def _get_case_state(case_name: str) -> dict:
    case_cfg = _get_case_config(case_name)
    case_key = case_cfg["key"]
    if case_key in _case_state_cache:
        return _case_state_cache[case_key]

    labels = case_cfg["labels"]
    model_path = _resolve_case_model_path(case_cfg)
    model = DenseNet121(num_classes=len(labels)).to(DEVICE)

    model_error: Optional[str] = None
    try:
        model.load_state_dict(torch.load(model_path, map_location=DEVICE))
        model.eval()
    except Exception as exc:
        model_error = f"Model failed to load for case `{case_name}` from `{model_path}`: {exc}"

    state = {
        "case_name": case_name,
        "labels": labels,
        "model_path": model_path,
        "model": model,
        "model_error": model_error,
    }
    _case_state_cache[case_key] = state
    return state


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 iter_image_paths(folder)


def _normalize_uploaded_paths(uploaded_files: Optional[Sequence[str] | str]) -> List[Path]:
    return normalize_uploaded_paths(uploaded_files)


def _resolve_input_images(
    uploaded_files: Optional[Sequence[str] | str],
    folder_path: str,
    google_drive_link: str = "",
) -> Tuple[List[Path], str]:
    uploaded_paths = _normalize_uploaded_paths(uploaded_files)

    local_paths: List[Path] = []
    if folder_path:
        folder = Path(folder_path).expanduser().resolve()
        if not folder.exists() or not folder.is_dir():
            return [], f"Invalid folder: `{folder}`"
        local_paths = _iter_image_paths(folder)

    drive_paths, drive_error = download_google_drive_images(google_drive_link or "")
    if drive_error:
        return [], drive_error

    all_paths = sorted({p.resolve() for p in uploaded_paths + local_paths + drive_paths})
    if not all_paths:
        return [], "Upload a folder (or files), provide a valid folder path, or enter a public Google Drive link."

    return all_paths, ""


def _aggregate_classification(classified: List[ClassifiedPrediction], labels: List[str]) -> 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 _predict_top1(model: DenseNet121, image: Image.Image) -> Tuple[int, float, torch.Tensor]:
    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]
        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 _classify_image_paths(
    model: DenseNet121,
    labels: List[str],
    image_paths: List[Path],
    threshold: float,
) -> Tuple[List[ClassifiedPrediction], List[List[str]]]:
    classified: List[ClassifiedPrediction] = []
    rows: List[List[str]] = []

    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, _ = _predict_top1(model, 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}"])

    return classified, rows


def _resize_rgb(image: np.ndarray, side: int) -> np.ndarray:
    pil_image = Image.fromarray(image)
    if hasattr(Image, "Resampling"):
        resized = pil_image.resize((side, side), Image.Resampling.BILINEAR)
    else:
        resized = pil_image.resize((side, side), Image.BILINEAR)
    return np.array(resized)


def _build_compact_visual_row(
    original: np.ndarray,
    gradcam_img: np.ndarray,
    smoothgrad_img: np.ndarray,
    tile_size: int = 192,
) -> np.ndarray:
    original_small = _resize_rgb(original, tile_size)
    gradcam_small = _resize_rgb(gradcam_img, tile_size)
    smoothgrad_small = _resize_rgb(smoothgrad_img, tile_size)
    return np.concatenate([original_small, gradcam_small, smoothgrad_small], axis=1)


def _build_labeled_visual_row(
    original: np.ndarray,
    gradcam_img: np.ndarray,
    smoothgrad_img: np.ndarray,
    tile_size: int = 384,
) -> np.ndarray:
    label_height = 42
    tiles = [
        ("Original", _resize_rgb(original, tile_size)),
        ("GradCAM", _resize_rgb(gradcam_img, tile_size)),
        ("SmoothGrad", _resize_rgb(smoothgrad_img, tile_size)),
    ]
    canvas = Image.new("RGB", (tile_size * len(tiles), tile_size + label_height), "#ffffff")
    draw = ImageDraw.Draw(canvas)

    for index, (title, tile) in enumerate(tiles):
        x = index * tile_size
        draw.rectangle((x, 0, x + tile_size, label_height), fill="#0f172a")
        text_box = draw.textbbox((0, 0), title)
        text_w = text_box[2] - text_box[0]
        text_h = text_box[3] - text_box[1]
        draw.text(
            (x + (tile_size - text_w) / 2, (label_height - text_h) / 2 - 1),
            title,
            fill="#ffffff",
        )
        canvas.paste(Image.fromarray(_as_uint8_rgb(tile)), (x, label_height))

    return np.array(canvas)


def _safe_stem(path: Path) -> str:
    stem = path.stem
    cleaned = "".join(ch if ch.isalnum() or ch in {"-", "_"} else "_" for ch in stem)
    return cleaned[:80] if cleaned else "image"


def _save_compact_panel(panel: np.ndarray, image_path: Path, save_dir: Path) -> Path:
    base = _safe_stem(image_path)
    out_path = save_dir / f"{base}_xdl.jpg"
    suffix = 1
    while out_path.exists():
        out_path = save_dir / f"{base}_xdl_{suffix}.jpg"
        suffix += 1
    Image.fromarray(panel).save(out_path, format="JPEG", quality=95)
    return out_path


def _as_uint8_rgb(image: np.ndarray) -> np.ndarray:
    arr = np.asarray(image)
    if arr.dtype != np.uint8:
        if np.issubdtype(arr.dtype, np.floating) and float(np.nanmax(arr)) <= 1.0:
            arr = arr * 255.0
        arr = np.clip(arr, 0, 255).astype(np.uint8)

    if arr.ndim == 2:
        arr = np.stack([arr, arr, arr], axis=-1)
    if arr.ndim == 3 and arr.shape[2] == 4:
        arr = arr[:, :, :3]
    return arr


def _image_to_data_uri(image: np.ndarray) -> str:
    buffer = BytesIO()
    Image.fromarray(_as_uint8_rgb(image)).save(buffer, format="PNG")
    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
    return f"data:image/png;base64,{encoded}"


def _render_xdl_panel(title: str, image: np.ndarray, filename: str, inspect_data_uri: str) -> str:
    title_html = escape(title)
    filename_html = escape(filename, quote=True)
    data_uri = _image_to_data_uri(image)
    return (
        "<div class='xdl-panel'>"
        f"<div class='xdl-panel-title'>{title_html}</div>"
        f"<a class='xdl-image-link' href='{inspect_data_uri}' target='_blank' rel='noopener noreferrer' "
        f"title='Open combined XDL panel for {filename_html} in a new tab'>"
        f"<img src='{data_uri}' alt='{title_html} for {filename_html}' />"
        "<span class='xdl-image-hint'>Open combined panel</span>"
        "</a>"
        "</div>"
    )


def _render_xdl_report_html(items: List[XdlVisualItem], status: str) -> str:
    status_html = escape(status)
    guide_html = (
        "<div class='xdl-guide'>"
        "<div class='xdl-guide-title'>How to Read the XDL Results</div>"
        "<div class='xdl-guide-grid'>"
        "<div><b>Original</b><span>Input image used by the classifier.</span></div>"
        "<div><b>GradCAM</b><span>Highlights regions most associated with the predicted class.</span></div>"
        "<div><b>SmoothGrad</b><span>Shows pixel-level sensitivity by averaging noisy gradients.</span></div>"
        "</div>"
        "</div>"
    )

    if not items:
        return (
            "<div class='xdl-report'>"
            f"{guide_html}"
            "<div class='xdl-empty'>"
            "<div class='xdl-empty-title'>No Classified XDL Cards Yet</div>"
            f"<div>{status_html}</div>"
            "</div>"
            "</div>"
        )

    cards = []
    total_items = len(items)
    for index, item in enumerate(items, start=1):
        filename = escape(item.filename)
        label = escape(item.predicted_label)
        confidence = float(np.clip(item.confidence * 100.0, 0.0, 100.0))
        inspect_data_uri = _image_to_data_uri(
            _build_labeled_visual_row(item.original, item.gradcam, item.smoothgrad)
        )
        cards.append(
            "<article class='xdl-card'>"
            "<div class='xdl-card-header'>"
            "<div>"
            f"<div class='xdl-index'>Image {index} of {total_items}</div>"
            f"<div class='xdl-filename'>{filename}</div>"
            "<div class='xdl-card-subtitle'>Classified image explanation</div>"
            "</div>"
            "<div class='xdl-meta'>"
            f"<span class='xdl-badge'>{label}</span>"
            f"<span class='xdl-confidence'>{confidence:.1f}% confidence</span>"
            "</div>"
            "</div>"
            "<div class='xdl-acceptance-note'>Accepted for XDL review because this image passed the selected confidence threshold.</div>"
            "<div class='xdl-panel-grid'>"
            f"{_render_xdl_panel('Original', item.original, item.filename, inspect_data_uri)}"
            f"{_render_xdl_panel('GradCAM', item.gradcam, item.filename, inspect_data_uri)}"
            f"{_render_xdl_panel('SmoothGrad', item.smoothgrad, item.filename, inspect_data_uri)}"
            "</div>"
            "</article>"
        )

    return (
        "<div class='xdl-report'>"
        f"{guide_html}"
        "<div class='xdl-report-heading'>"
        "<div>Classified Image XDL Cards</div>"
        f"<span>{len(items)} generated</span>"
        "</div>"
        "<div class='xdl-scroll'>"
        f"{''.join(cards)}"
        "</div>"
        f"<div class='xdl-status'>{status_html}</div>"
        "</div>"
    )


def _render_summary_html(
    *,
    case_name: str,
    model_path: str,
    device_name: str,
    processed: int,
    classified: int,
    threshold: float,
    final_class: str,
    mean_confidence: Optional[float],
    distribution: str,
    xdl_status: str,
    save_status: str,
) -> str:
    if mean_confidence is None:
        mean_conf_text = "N/A"
        progress_percent = 0.0
    else:
        progress_percent = float(np.clip(mean_confidence * 100.0, 0.0, 100.0))
        mean_conf_text = f"{progress_percent:.1f}%"

    skipped = processed - classified
    metrics = [
        ("Final Class", final_class, "small"),
        ("Mean Confidence", mean_conf_text, ""),
        ("Processed", str(processed), ""),
        ("Classified", f"{classified} / {processed}", ""),
    ]
    metric_html = "".join(
        "<div class='metric-card'>"
        f"<div class='metric-label'>{escape(label)}</div>"
        f"<div class='metric-value {escape(size_class)}'>{escape(value)}</div>"
        "</div>"
        for label, value, size_class in metrics
    )

    details = [
        ("Case", case_name),
        ("Device", device_name),
        ("Confidence Threshold", f"{threshold:.2f}"),
        ("Skipped", str(skipped)),
        ("Class Distribution", distribution),
        ("XDL", xdl_status),
        ("Saved XDL", save_status),
        ("Model Path", model_path),
    ]
    detail_html = "".join(
        "<div class='summary-detail'>"
        f"<b>{escape(key)}:</b> {escape(value)}"
        "</div>"
        for key, value in details
    )

    return (
        "<div class='summary-card'>"
        "<div class='summary-title'>"
        "<h3>Batch Summary</h3>"
        f"<span>{escape(case_name)}</span>"
        "</div>"
        f"<div class='metric-grid'>{metric_html}</div>"
        f"<div class='summary-details'>{detail_html}</div>"
        "</div>"
    )


def _render_error_html(message: str) -> str:
    return (
        "<div style='border:1px solid #fca5a5;border-radius:10px;padding:12px;background:#fef2f2;color:#7f1d1d'>"
        "<div style='font-weight:600;margin-bottom:6px'>Input/Error</div>"
        f"<div>{escape(message)}</div>"
        "</div>"
    )


def _progress_percent(current: int, total: int) -> float:
    if total <= 0:
        return 0.0
    return float(np.clip((current / total) * 100.0, 0.0, 100.0))


def _render_progress_section(title: str, message: str, current: int = 0, total: int = 0) -> str:
    percent = _progress_percent(current, total)
    count_html = f"<span>{current} / {total}</span>" if total > 0 else ""
    return (
        "<div class='progress-card'>"
        "<div class='progress-card-head'>"
        f"<b>{escape(title)}</b>"
        f"{count_html}"
        "</div>"
        f"<div class='progress-message'>{escape(message)}</div>"
        "<div class='progress-track'>"
        f"<div class='progress-fill' style='width:{percent:.2f}%'></div>"
        "</div>"
        "</div>"
    )


def _render_processing_status_html(
    input_message: str = "Waiting for input...",
    input_current: int = 0,
    input_total: int = 0,
    xdl_message: str = "Waiting for classified images...",
    xdl_current: int = 0,
    xdl_total: int = 0,
) -> str:
    return (
        "<div class='processing-status'>"
        "<div class='processing-status-title'>Processing Status</div>"
        "<div class='processing-grid'>"
        f"{_render_progress_section('Input Preparation', input_message, input_current, input_total)}"
        f"{_render_progress_section('XDL Generation', xdl_message, xdl_current, xdl_total)}"
        "</div>"
        "</div>"
    )


def _loaded_images_update(image_paths: List[Path]) -> Any:
    return _loaded_image_names_update([path.name for path in image_paths])


def _loaded_image_names_update(image_names: List[str]) -> Any:
    names = list(dict.fromkeys(name for name in image_names if name))
    return gr.update(choices=names, value=names[0] if names else None, interactive=bool(names))


def _empty_loaded_images_update() -> Any:
    return gr.update(choices=[], value=None, interactive=False)


def batch_predict_with_xdl(
    uploaded_files,
    selected_case: str,
    folder_path: str,
    confidence_threshold: float,
    smoothgrad_samples: int,
    smoothgrad_noise: float,
    save_xdl_results: bool,
    save_xdl_dir: str,
    gradcam_target_layer: str = DEFAULT_GRADCAM_TARGET_LAYER,
    google_drive_link: str = "",
):
    last_output = None
    for payload in batch_predict_with_xdl_stream(
        uploaded_files=uploaded_files,
        selected_case=selected_case,
        folder_path=folder_path,
        confidence_threshold=confidence_threshold,
        smoothgrad_samples=smoothgrad_samples,
        smoothgrad_noise=smoothgrad_noise,
        save_xdl_results=save_xdl_results,
        save_xdl_dir=save_xdl_dir,
        gradcam_target_layer=gradcam_target_layer,
        google_drive_link=google_drive_link,
    ):
        last_output = payload

    if last_output is None:
        return _render_error_html("No output generated."), "", "", _empty_loaded_images_update(), []
    return last_output


def batch_predict_with_xdl_stream(
    uploaded_files,
    selected_case: str,
    folder_path: str,
    confidence_threshold: float,
    smoothgrad_samples: int,
    smoothgrad_noise: float,
    save_xdl_results: bool,
    save_xdl_dir: str,
    gradcam_target_layer: str = DEFAULT_GRADCAM_TARGET_LAYER,
    google_drive_link: str = "",
) -> Iterator[Tuple[str, str, str, Any, List[List[str]]]]:
    case_state = _get_case_state(selected_case)
    model: DenseNet121 = case_state["model"]
    labels: List[str] = case_state["labels"]
    case_name: str = case_state["case_name"]
    model_path: str = case_state["model_path"]
    model_error: Optional[str] = case_state["model_error"]

    if model_error:
        yield _render_error_html(model_error), "", "", _empty_loaded_images_update(), []
        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

    uploaded_paths = _normalize_uploaded_paths(uploaded_files)
    local_paths: List[Path] = []
    if folder_path:
        folder = Path(folder_path).expanduser().resolve()
        if not folder.exists() or not folder.is_dir():
            yield _render_error_html(f"Invalid folder: `{folder}`"), "", "", _loaded_images_update(uploaded_paths), []
            return
        local_paths = _iter_image_paths(folder)

    base_paths = sorted({p.resolve() for p in uploaded_paths + local_paths})
    base_names = [path.name for path in base_paths]
    input_message = f"{len(base_paths)} uploaded images found." if base_paths else "Waiting for image input."
    yield (
        "",
        _render_processing_status_html(input_message, len(base_paths), len(base_paths)),
        _render_xdl_report_html([], "Waiting for classified images..."),
        _loaded_images_update(base_paths),
        [],
    )

    drive_paths: List[Path] = []
    if google_drive_link and google_drive_link.strip():
        for progress in download_google_drive_folder_images(google_drive_link):
            current_paths = sorted({p.resolve() for p in base_paths + progress.paths})
            current_names = base_names + (progress.filenames or [path.name for path in progress.paths])
            status_html = _render_processing_status_html(
                progress.message,
                progress.current,
                progress.total,
            )
            if progress.error:
                yield _render_error_html(progress.error), status_html, "", _loaded_image_names_update(current_names), []
                return
            yield (
                "",
                status_html,
                _render_xdl_report_html([], "Waiting for classified images..."),
                _loaded_image_names_update(current_names),
                [],
            )
            if progress.done:
                drive_paths = progress.paths

    image_paths = sorted({p.resolve() for p in base_paths + drive_paths})
    if not image_paths:
        message = "Upload a folder (or files) or enter a public Google Drive folder link."
        yield (
            _render_error_html(message),
            _render_processing_status_html("No supported images found.", 0, 0),
            "",
            _empty_loaded_images_update(),
            [],
        )
        return

    input_ready_status = f"Input ready. {len(image_paths)} supported images found."
    yield (
        "",
        _render_processing_status_html(input_ready_status, len(image_paths), len(image_paths), "Running classification..."),
        _render_xdl_report_html([], "Waiting for classification results..."),
        _loaded_images_update(image_paths),
        [],
    )

    save_dir_path: Optional[Path] = None
    save_error = ""
    saved_count = 0
    if save_xdl_results:
        raw_dir = save_xdl_dir.strip() if save_xdl_dir else DEFAULT_SAVE_DIR
        try:
            save_dir_path = Path(raw_dir).expanduser().resolve()
            save_dir_path.mkdir(parents=True, exist_ok=True)
        except Exception as exc:
            save_error = f"Save disabled: {exc}"
            save_dir_path = None

    classified, rows = _classify_image_paths(model, labels, image_paths, threshold)

    if save_xdl_results:
        if save_error:
            save_status = save_error
        elif save_dir_path is not None and not classified:
            save_status = f"No files saved (0 classified). Target: {save_dir_path}"
        else:
            save_status = "Pending..."
    else:
        save_status = "Disabled"

    if classified:
        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 = f"Processing overlays... (GradCAM layer: {gradcam_target_layer})"
    else:
        final_class = "N/A"
        mean_conf = None
        class_stats = "N/A"
        initial_xdl_status = "Skipped (no classified images)"

    summary_initial = _render_summary_html(
        case_name=case_name,
        model_path=model_path,
        device_name=model_device.type,
        processed=len(image_paths),
        classified=len(classified),
        threshold=threshold,
        final_class=final_class,
        mean_confidence=mean_conf,
        distribution=class_stats,
        xdl_status=initial_xdl_status,
        save_status=save_status,
    )

    # First UI update: classification results first while XDL cards are still being generated.
    yield (
        summary_initial,
        _render_processing_status_html(
            input_ready_status,
            len(image_paths),
            len(image_paths),
            initial_xdl_status,
            0,
            len(classified),
        ),
        _render_xdl_report_html([], initial_xdl_status),
        _loaded_images_update(image_paths),
        rows,
    )

    if not classified:
        return

    xdl = None
    xdl_error = ""
    try:
        xdl = _load_xdl_modules()
    except RuntimeError as exc:
        xdl_error = str(exc)

    visual_items: List[XdlVisualItem] = []
    xdl_error_count = 0
    xdl_processed = 0

    if xdl is not None:
        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:
            yield (
                summary_initial,
                _render_processing_status_html(
                    input_ready_status,
                    len(image_paths),
                    len(image_paths),
                    f"Generating overlays for {item.path.name}...",
                    xdl_processed,
                    len(classified),
                ),
                _render_xdl_report_html(
                    [],
                    f"Processing {xdl_processed} / {len(classified)} overlays. Gallery renders after all overlays finish.",
                ),
                _loaded_images_update(image_paths),
                rows,
            )
            try:
                image = Image.open(item.path).convert("RGB")
                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]

                grayscale_cam = cam(
                    input_tensor=input_tensor,
                    targets=[xdl["ClassifierOutputTarget"](item.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,
                    item.pred_idx,
                    n_samples=smoothgrad_samples,
                    noise_level=smoothgrad_noise,
                    batch_size=_smoothgrad_batch_size(model_device),
                    use_amp=(model_device.type == "cuda"),
                )
                _, smooth_heatmap = xdl["_process_smoothgrad_map"](
                    smooth_raw,
                    img_shape=(h, w),
                    percentile=95,
                    colormap="hot",
                )
                # Match old-original-xdl.py blending behavior exactly.
                smooth_overlay = xdl["cv2"].addWeighted(base_img_uint8, 0.5, smooth_heatmap, 0.5, 0)

                compact_panel = _build_compact_visual_row(base_img_uint8, gradcam_overlay, smooth_overlay)
                visual_items.append(
                    XdlVisualItem(
                        filename=item.path.name,
                        predicted_label=labels[item.pred_idx],
                        confidence=item.confidence,
                        original=base_img_uint8,
                        gradcam=gradcam_overlay,
                        smoothgrad=smooth_overlay,
                    )
                )

                if save_dir_path is not None:
                    _save_compact_panel(compact_panel, item.path, save_dir_path)
                    saved_count += 1
            except Exception as exc:
                xdl_error_count += 1
                rows.append([item.path.name, "xdl_error", labels[item.pred_idx], str(exc)])
            xdl_processed += 1
            yield (
                summary_initial,
                _render_processing_status_html(
                    input_ready_status,
                    len(image_paths),
                    len(image_paths),
                    f"XDL overlays complete: {xdl_processed} / {len(classified)}",
                    xdl_processed,
                    len(classified),
                ),
                _render_xdl_report_html(
                    [],
                    f"Generated {len(visual_items)} / {len(classified)} XDL cards. Gallery renders after all overlays finish.",
                ),
                _loaded_images_update(image_paths),
                rows,
            )

    if xdl_error:
        xdl_status = f"Disabled: {xdl_error}"
    elif xdl_error_count:
        xdl_status = f"Completed with {xdl_error_count} overlay errors"
    else:
        xdl_status = f"Completed ({len(visual_items)} overlays, layer: {gradcam_target_layer})"

    if save_xdl_results:
        if save_error:
            save_status = save_error
        elif save_dir_path is not None:
            save_status = f"Saved {saved_count} files to {save_dir_path}"
        else:
            save_status = "Disabled"
    else:
        save_status = "Disabled"

    summary_final = _render_summary_html(
        case_name=case_name,
        model_path=model_path,
        device_name=model_device.type,
        processed=len(image_paths),
        classified=len(classified),
        threshold=threshold,
        final_class=final_class,
        mean_confidence=mean_conf,
        distribution=class_stats,
        xdl_status=xdl_status,
        save_status=save_status,
    )
    yield (
        summary_final,
        _render_processing_status_html(
            input_ready_status,
            len(image_paths),
            len(image_paths),
            xdl_status,
            xdl_processed if xdl is not None else 0,
            len(classified),
        ),
        _render_xdl_report_html(visual_items, xdl_status),
        _loaded_images_update(image_paths),
        rows,
    )