scripts/predict_single.py

"""
predict_single.py — GazeAlign single-image inference CLI + Python API

Usage (CLI):
    python scripts/predict_single.py \
        --image  examples/images/0c1c6a70-a96f5b27-5d042944-6b49b3b3-fd6a8293.jpg \
        --fixations examples/fixations/fixations.csv \
        --output outputs/output_mask.png \
        --preset cxr

Usage (Python API):
    from scripts.predict_single import GazeAlignPredictor

    predictor = GazeAlignPredictor.from_preset("cxr", presets_path="configs/presets.yaml")
    result = predictor.predict(
        "examples/images/0c1c6a70-a96f5b27-5d042944-6b49b3b3-fd6a8293.jpg",
        "examples/fixations/fixations.csv",
    )
    print(result.predicted_class)
    print(result.class_probs)
"""

from __future__ import annotations

import sys
import os

# Make the repo root importable so "import GazeAlign" works without pip install
_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _REPO_ROOT not in sys.path:
    sys.path.insert(0, _REPO_ROOT)

import argparse
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Optional

import numpy as np
import torch
import yaml


# ---------------------------------------------------------------------------
# Result container
# ---------------------------------------------------------------------------

@dataclass
class PredictionResult:
    """Holds all outputs from a single GazeAlign inference pass."""

    predicted_class: str
    """The top-1 class label string."""

    class_probs: Dict[str, float]
    """Softmax probability for every class, keyed by label string."""

    attention_mask: np.ndarray
    """Learned gaze-conditioned attention mask, shape [H, W], values in [0, 1]."""

    gaze_prior: np.ndarray
    """Raw fixation Gaussian-splat heatmap, shape [H, W], values in [0, 1]."""


# ---------------------------------------------------------------------------
# Predictor class
# ---------------------------------------------------------------------------

class GazeAlignPredictor:
    """
    High-level wrapper for GazeAlign single-image inference.

    Instantiate via :meth:`from_preset` (recommended) or :meth:`__init__`
    when you already have a loaded model and class list.
    """

    def __init__(
        self,
        model: torch.nn.Module,
        classes: list[str],
        device: str = "cpu",
        img_size: int = 518,
    ):
        from GazeAlign.datasets import build_transform  # noqa: PLC0415

        self.model     = model.eval().to(device)
        self.classes   = classes
        self.device    = device
        self.img_size  = img_size
        # Same preprocessing used during training (resize + grayscale->3ch),
        # so single-image inference matches how the checkpoint was trained.
        self.transform = build_transform(img_size)

    # Expose the underlying sub-modules directly (used by the HF Space demo
    # and any code that wants to run the pipeline stage-by-stage).
    @property
    def image_encoder(self):
        return self.model.image_encoder

    @property
    def scanpath_encoder(self):
        return self.model.scanpath_encoder

    @property
    def mask_generator(self):
        return self.model.mask_generator

    @property
    def classifier(self):
        return self.model.classifier

    # ------------------------------------------------------------------
    # Construction helpers
    # ------------------------------------------------------------------

    @classmethod
    def from_preset(
        cls,
        preset_name: str,
        presets_path: str | Path = "configs/presets.yaml",
        device: Optional[str] = None,
    ) -> "GazeAlignPredictor":
        """
        Build a predictor from a named preset in *presets_path*.

        The YAML file must have a top-level key matching *preset_name*,
        with sub-keys ``checkpoint`` (path to .pth) and ``classes`` (list).

        Parameters
        ----------
        preset_name:
            Key in presets.yaml, e.g. ``"cxr"``.
        presets_path:
            Path to ``configs/presets.yaml``.
        device:
            ``"cuda"``, ``"cpu"``, or ``None`` (auto-detect).
        """
        presets_path = Path(presets_path)
        if not presets_path.exists():
            raise FileNotFoundError(f"Presets file not found: {presets_path}")

        with open(presets_path) as f:
            raw = yaml.safe_load(f)

        # Support both flat { cxr: {...} } and nested { presets: { cxr: {...} } }
        presets = raw.get("presets", raw)

        if preset_name not in presets:
            available = list(presets.keys())
            raise KeyError(
                f"Preset '{preset_name}' not found in {presets_path}. "
                f"Available presets: {available}"
            )

        cfg           = presets[preset_name]
        ckpt_path     = cfg["checkpoint"]
        classes       = cfg["classes"]
        img_size      = cfg.get("img_size", 518)
        grid_size     = cfg.get("grid_size", 37)
        backbone_name = cfg.get("backbone_name", "microsoft/rad-dino-maira-2")

        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"

        model = cls._load_model(
            ckpt_path,
            num_classes=len(classes),
            grid_size=grid_size,
            backbone_name=backbone_name,
            device=device,
        )
        return cls(model=model, classes=classes, device=device, img_size=img_size)

    @staticmethod
    def _load_model(
        ckpt_path: str | Path,
        num_classes: int,
        grid_size: int,
        backbone_name: str,
        device: str,
    ) -> torch.nn.Module:
        """Rebuild the five GazeAlign sub-modules and load a training checkpoint.

        Checkpoints written by ``scripts/train.py`` are wrapped dicts with
        one state-dict per sub-module (``image_encoder_state`` etc.), so we
        construct each module first and load them individually — mirroring
        ``scripts/run_eval.py``.
        """
        # Lazy imports so the module works even before GazeAlign is installed
        from GazeAlign.engine import GazeAlignModel  # noqa: PLC0415
        from GazeAlign.backbone import DINOv3Encoder  # noqa: PLC0415
        from GazeAlign.model import (  # noqa: PLC0415
            ScanpathTransformer,
            ViTMaskGenerator,
            classifier,
        )
        from GazeAlign.utils import (  # noqa: PLC0415
            load_gaze_checkpoint,
            load_submodule_state,
        )

        ckpt_path = Path(ckpt_path)
        if not ckpt_path.exists():
            raise FileNotFoundError(
                f"Checkpoint not found: {ckpt_path}\n"
                f"Download or train a model first (see scripts/train.py)."
            )

        image_encoder    = DINOv3Encoder(backbone_name=backbone_name)
        scanpath_encoder = ScanpathTransformer()
        mask_generator   = ViTMaskGenerator(num_patches=grid_size ** 2)
        clf              = classifier(num_classes=num_classes)

        # Normalizes both refactored and original research-script checkpoints.
        ckpt = load_gaze_checkpoint(ckpt_path, map_location=device)
        load_submodule_state(image_encoder, ckpt["image_encoder"])
        load_submodule_state(scanpath_encoder, ckpt["scanpath_encoder"])
        load_submodule_state(mask_generator, ckpt["mask_generator"])
        load_submodule_state(clf, ckpt["classifier"])
        if "epoch" in ckpt:
            print(f"[GazeAlign] Loaded checkpoint from epoch {ckpt['epoch']}")

        model = GazeAlignModel(
            image_encoder=image_encoder,
            scanpath_encoder=scanpath_encoder,
            mask_generator=mask_generator,
            classifier=clf,
            num_classes=num_classes,
            grid_size=grid_size,
        )
        return model

    # ------------------------------------------------------------------
    # Inference
    # ------------------------------------------------------------------

    def predict(
        self,
        image_path: str | Path,
        fixations_csv: str | Path,
        dicom_id: Optional[str] = None,
    ) -> PredictionResult:
        """
        Run a full GazeAlign forward pass on one image + its fixation CSV.

        Parameters
        ----------
        image_path:
            Path to the input image (JPEG, PNG, …).
        fixations_csv:
            Path to a MIMIC-style fixations CSV (columns ``DICOM_ID``,
            ``X_ORIGINAL``, ``Y_ORIGINAL``, ``Time (in secs)``).
        dicom_id:
            Which image's scanpath to use from the CSV. Defaults to the
            stem of *image_path*; if that id is absent and the CSV holds a
            single ``DICOM_ID``, that one is used automatically.

        Returns
        -------
        PredictionResult
        """
        from GazeAlign.constants import ID_COL  # noqa: PLC0415
        from GazeAlign.gaze import (  # noqa: PLC0415
            fixation_heatmap,
            get_scanpath,
            load_fixation_csv,
        )

        image_path    = Path(image_path)
        fixations_csv = Path(fixations_csv)

        if not image_path.exists():
            raise FileNotFoundError(f"Image not found: {image_path}")
        if not fixations_csv.exists():
            raise FileNotFoundError(f"Fixations CSV not found: {fixations_csv}")

        # 1. Load + preprocess image (also gives us original pixel dims for
        #    normalizing the fixation coordinates the way training did).
        img_tensor, orig_hw = self._load_image(image_path)
        orig_h, orig_w = orig_hw

        # 2. Resolve which scanpath (DICOM_ID) to use.
        df = load_fixation_csv(str(fixations_csv))
        ids = set(df[ID_COL].astype(str))
        if dicom_id is None:
            dicom_id = image_path.stem
        if dicom_id not in ids:
            if len(ids) == 1:
                dicom_id = next(iter(ids))
            else:
                raise ValueError(
                    f"DICOM_ID '{dicom_id}' not found in {fixations_csv}. "
                    f"Available ids: {sorted(ids)}. "
                    f"Pass --dicom_id to select one."
                )

        # 3. Build the normalized [T, 3] scanpath the model was trained on.
        scanpath = get_scanpath(df, dicom_id, img_height=orig_h, img_width=orig_w)
        if scanpath is None or scanpath.numel() == 0:
            raise ValueError(f"No fixations for DICOM_ID '{dicom_id}' in {fixations_csv}.")
        scanpath = scanpath[:200].to(self.device)  # match training truncation

        # 4. Raw gaze-prior heatmap (visualization only, not used by the model).
        fix_df = df[df[ID_COL].astype(str) == str(dicom_id)]
        gaze_prior = fixation_heatmap(
            fix_df, orig_h, orig_w, orig_h, orig_w
        )  # np.ndarray [H, W] in [0, 1]

        # 5. Forward pass through the sub-modules (single-image inference: no
        #    negatives / contrastive terms — just image + scanpath -> mask ->
        #    gaze-masked classification, mirroring engine.forward_batch).
        self.model.eval()
        with torch.no_grad():
            _, patch_tokens, _ = self.model.image_encoder(
                img_tensor.unsqueeze(0).to(self.device)
            )
            _, sp_emb, _ = self.model.scanpath_encoder([scanpath])
            patch_mask = torch.sigmoid(self.model.mask_generator(sp_emb))  # [1, g, g]

            weights_pos = patch_mask.view(1, -1, 1)
            feat_attended = (patch_tokens * weights_pos).mean(dim=1)
            logits = self.model.classifier(feat_attended)  # [1, num_classes]

        probs = torch.softmax(logits[0], dim=-1).cpu().numpy()
        top_idx = int(probs.argmax())

        attention_mask = patch_mask[0].cpu().numpy()  # [grid, grid]
        # Normalise to [0, 1] for display
        mn, mx = attention_mask.min(), attention_mask.max()
        if mx > mn:
            attention_mask = (attention_mask - mn) / (mx - mn)

        return PredictionResult(
            predicted_class=self.classes[top_idx],
            class_probs={cls: float(p) for cls, p in zip(self.classes, probs)},
            attention_mask=attention_mask,
            gaze_prior=gaze_prior,
        )

    # ------------------------------------------------------------------
    # Image loading
    # ------------------------------------------------------------------

    def _load_image(self, image_path: Path):
        """Return (preprocessed tensor [3, H, W], (orig_h, orig_w)).

        Uses the shared ``build_transform`` preprocessing (resize +
        grayscale-replicated-to-3-channel) so inference matches training.
        """
        from PIL import Image

        img = Image.open(image_path).convert("RGB")
        orig_hw = (img.height, img.width)
        return self.transform(img), orig_hw


# ---------------------------------------------------------------------------
# Output helpers
# ---------------------------------------------------------------------------

def _save_mask(mask: np.ndarray, path: Path) -> None:
    from PIL import Image
    arr = (mask * 255).clip(0, 255).astype(np.uint8)
    Image.fromarray(arr, mode="L").save(path)


def _save_overlay(
    image_path: Path,
    mask: np.ndarray,
    out_path: Path,
    alpha: float = 0.5,
    colormap: str = "jet",
) -> None:
    import matplotlib.cm as cm
    from PIL import Image

    img  = Image.open(image_path).convert("RGB")
    mask_resized = np.array(
        Image.fromarray((mask * 255).astype(np.uint8)).resize(
            (img.width, img.height), Image.BILINEAR
        )
    ) / 255.0

    cmap   = cm.get_cmap(colormap)
    heatmap = (cmap(mask_resized)[:, :, :3] * 255).astype(np.uint8)
    blended = Image.blend(img, Image.fromarray(heatmap), alpha=alpha)
    blended.save(out_path)


# ---------------------------------------------------------------------------
# CLI entry-point
# ---------------------------------------------------------------------------

def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="GazeAlign — single-image gaze-conditioned classification"
    )
    p.add_argument(
        "--image",
        default="examples/images/0c1c6a70-a96f5b27-5d042944-6b49b3b3-fd6a8293.jpg",
        help="Path to the input chest X-ray (or other medical image).",
    )
    p.add_argument(
        "--fixations",
        default="examples/fixations/fixations.csv",
        help=(
            "Path to a MIMIC-style fixations CSV with columns: "
            "DICOM_ID, X_ORIGINAL, Y_ORIGINAL, Time (in secs). "
            "The scanpath is selected by --dicom_id (default: stem of --image)."
        ),
    )
    p.add_argument(
        "--dicom_id",
        default=None,
        help=(
            "DICOM_ID whose scanpath to use from the fixations CSV. "
            "Defaults to the stem of --image; falls back to the sole id if "
            "the CSV contains only one."
        ),
    )
    p.add_argument(
        "--output",
        default="outputs/output_mask.png",
        help="Where to save the learned gaze-conditioned attention mask.",
    )
    p.add_argument(
        "--preset",
        default="cxr",
        help="Preset name in configs/presets.yaml (checkpoint + class list).",
    )
    p.add_argument(
        "--presets_path",
        default="configs/presets.yaml",
        help="Path to the presets YAML file.",
    )
    p.add_argument(
        "--device",
        default=None,
        help="'cuda' or 'cpu'. Defaults to cuda if available.",
    )
    p.add_argument(
        "--save_overlay",
        action="store_true",
        help=(
            "Also save a colour overlay (output_mask_overlay.png) "
            "and the raw gaze prior (output_mask_gaze_prior.png)."
        ),
    )
    return p.parse_args()


def main() -> None:
    args = _parse_args()

    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    print(f"[GazeAlign] Loading preset '{args.preset}' from {args.presets_path} …")
    predictor = GazeAlignPredictor.from_preset(
        args.preset,
        presets_path=args.presets_path,
        device=args.device,
    )

    image_stem = Path(args.image).stem
    print(f"[GazeAlign] Running inference on '{image_stem}' …")
    result = predictor.predict(args.image, args.fixations, dicom_id=args.dicom_id)

    # --- Print results ---
    print(f"\n  Predicted class : {result.predicted_class}")
    print("  Class probabilities:")
    for cls, prob in sorted(result.class_probs.items(), key=lambda x: -x[1]):
        bar = "█" * int(prob * 40)
        print(f"    {cls:>12s}  {prob:.4f}  {bar}")

    # --- Save attention mask ---
    _save_mask(result.attention_mask, out_path)
    print(f"\n  Saved attention mask → {out_path}")

    if args.save_overlay:
        overlay_path    = out_path.with_name(out_path.stem + "_overlay" + out_path.suffix)
        gaze_prior_path = out_path.with_name(out_path.stem + "_gaze_prior" + out_path.suffix)

        _save_overlay(args.image, result.attention_mask, overlay_path)
        _save_mask(result.gaze_prior, gaze_prior_path)

        print(f"  Saved colour overlay  → {overlay_path}")
        print(f"  Saved gaze prior      → {gaze_prior_path}")

    print("\n[GazeAlign] Done.")


if __name__ == "__main__":
    main()