scripts/predict_single.py

"""
predict_single.py — Run GazeRefine on a single image + fixation CSV.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Command-line usage (matches the README exactly)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

    python scripts/predict_single.py \\
        --image   examples/image.png \\
        --fixations examples/fixations.csv \\
        --output  output_mask.png

Optional flags:
    --preset        colonoscopy | mri           (default: colonoscopy)
    --threshold     0.5                          binarization threshold
    --save_overlay                               also save a colour overlay PNG
    --device        cuda | cpu                   (auto-detected by default)

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Python API (matches the README exactly)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

    from scripts.predict_single import predict

    mask = predict(
        image_path="image.png",
        fixation_csv="fixations.csv",
    )
    # `mask` is a PIL Image of the binary segmentation mask.
    # Save it:
    mask.save("output_mask.png")

    # Extended API — also get overlays and raw arrays:
    result = predict(
        image_path="image.png",
        fixation_csv="fixations.csv",
        preset="mri",           # "colonoscopy" (default) or "mri"
        threshold=0.5,
        return_all=True,
    )
    result["mask"].save("mask.png")
    result["gaze_overlay"].save("gaze.png")
    result["mask_overlay"].save("overlay.png")

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Fixation CSV format
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

    x,y,duration
    340,221,180
    356,228,145
    368,244,205
    ...

    x, y     — fixation position in *raw pixel* coordinates of the input image.
               (These are automatically normalized by the image size internally.)
    duration — fixation duration in any consistent unit (milliseconds typical).
               The model only uses *relative* durations, so the unit does not matter.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path
from typing import Sequence

import numpy as np
import pandas as pd
import torch
from PIL import Image

# allow `python scripts/predict_single.py` from any working directory
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from gazecorrect import GazeRefine, overlay_heatmap, overlay_mask
from gazecorrect.constants import IMG_MEAN, IMG_STD, IMG_SIZE
from gazecorrect.gaze import load_fixation_csv
import torchvision.transforms as T


# ── per-modality hyperparameter presets ────────────────────────────────────
#    These match the exact settings used to produce the paper's Table 1 numbers.
PRESETS: dict[str, dict] = {
    "colonoscopy": dict(
        sigma=2.0,
        contrast_method="difference",
        max_iters=5,
        gaze_anchor_weight=0.5,
        knn_refine=True,
        knn_k=20,
        knn_temp=0.1,
    ),
    "mri": dict(
        sigma=1.5,
        contrast_method="difference",
        max_iters=1,
        gaze_anchor_weight=0.8,
        knn_refine=True,
        knn_k=3,
        knn_temp=0.1,
    ),
}

# one shared backbone name for both presets
DINO_NAME = "vit_large_patch16_dinov3.lvd1689m"

# module-level model cache: avoids reloading the backbone across repeated calls
# (useful when this module is imported by the Gradio Space or a notebook loop)
_MODEL_CACHE: dict[str, GazeRefine] = {}


def _get_model(preset: str, device: torch.device) -> GazeRefine:
    """Load (or return a cached) GazeRefine model for the given preset."""
    if preset not in _MODEL_CACHE:
        cfg = PRESETS[preset]
        _MODEL_CACHE[preset] = GazeRefine(dino_name=DINO_NAME, **cfg).to(device).eval()
    return _MODEL_CACHE[preset]


# ── main public function ────────────────────────────────────────────────────
import os
import pydicom
import numpy as np
from PIL import Image
@torch.no_grad()
def predict(
    image_path: "str | Path | Image.Image",
    fixation_csv: "str | Path",
    preset: str = "colonoscopy",
    threshold: float = 0.5,
    device: str | None = None,
    return_all: bool = False,
) -> "Image.Image | dict":
    """Run GazeRefine on one image and return the predicted segmentation mask.

    Parameters
    ----------
    image_path   : path to the input image (.jpg / .jpeg / .png) **or** an
                   already-loaded ``PIL.Image`` (used by the Gradio Space).
    fixation_csv : path to the fixation CSV (``x,y,duration`` columns,
                   pixel coordinates — see module docstring for the format).
    preset       : ``"colonoscopy"`` (default, Kvasir-SEG settings) or
                   ``"mri"`` (NCI-ISBI prostate-MRI settings).
    threshold    : binarization cutoff applied to the [0, 1] soft mask.
    device       : ``"cuda"`` / ``"cpu"`` — auto-detected when ``None``.
    return_all   : when ``True``, return a dict with the binary mask PIL Image
                   **plus** ``gaze_overlay``, ``mask_overlay``, and the raw
                   numpy arrays ``preds`` and ``gaze_heatmap``.
                   When ``False`` (default), return only the mask PIL Image.

    Returns
    -------
    ``PIL.Image`` of the binary mask, **or** a dict (see ``return_all``).
    """
    if preset not in PRESETS:
        raise ValueError(f"preset must be one of {list(PRESETS)}, got {preset!r}")

    _device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))

    # ── load the image ──
    ext = os.path.splitext(str(image_path))[1].lower()

    if isinstance(image_path, Image.Image):
        pil_image = image_path.convert("RGB")

    elif ext == ".dcm":
        import pydicom

        dcm = pydicom.dcmread(str(image_path))
        arr = dcm.pixel_array.astype(np.float32)

        # normalize properly (medical safe scaling)
        arr = arr - arr.min()
        arr = arr / (arr.max() + 1e-8)
        arr = (arr * 255).astype(np.uint8)

        pil_image = Image.fromarray(arr).convert("RGB")

    else:
        pil_image = Image.open(image_path).convert("RGB")    
    
    img_w, img_h = pil_image.size

    # ── load fixations and normalize pixel → [0, 1] ──
    from pathlib import Path
    fix_t = load_fixation_csv(
    str(fixation_csv),
    image_width=img_w,
    image_height=img_h,
    image_name = Path(image_path).stem
    )  # (N, 3) float tensor: x, y, duration
    fix_t = fix_t.to(_device)

    # ── preprocess the image for DINOv3 ──
    tf = T.Compose([
        T.Resize((IMG_SIZE, IMG_SIZE)),
        T.ToTensor(),
        T.Normalize(IMG_MEAN, IMG_STD),
    ])
    img_t = tf(pil_image).unsqueeze(0).to(_device)  # (1, 3, IMG_SIZE, IMG_SIZE)

    # ── run the model ──
    model = _get_model(preset, _device)
    out   = model(img_t, fix_t)

    # ── decode outputs ──
    soft_mask = out["preds"][0, 0].cpu().numpy()    # (H, W) float in [0, 1]
    gaze      = out["gaze_heatmap"][0].cpu().numpy()  # (h, w) float in [0, 1]
    bin_mask  = (soft_mask > threshold).astype(np.uint8) * 255

    mask_pil = Image.fromarray(bin_mask, mode="L")

    if not return_all:
        return mask_pil

    return dict(
        mask          = mask_pil,
        gaze_overlay  = overlay_heatmap(pil_image, gaze),
        mask_overlay  = overlay_mask(pil_image, (bin_mask / 255).astype(np.float32)),
        preds         = soft_mask,
        gaze_heatmap  = gaze,
    )


# ── CLI ─────────────────────────────────────────────────────────────────────

def _build_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(
        prog="predict_single.py",
        description="GazeRefine — zero-shot gaze-guided segmentation on a single image.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples
--------
  # colonoscopy polyp (default preset):
  python scripts/predict_single.py \\
      --image examples/images/kvasir_sample.jpg \\
      --fixations examples/fixations/kvasir_sample.csv \\
      --output output_mask.png

  # prostate MRI:
  python scripts/predict_single.py \\
      --image examples/images/prostate_sample.png \\
      --fixations examples/fixations/prostate_sample.csv \\
      --output output_mask.png \\
      --preset mri

  # save overlays too:
  python scripts/predict_single.py \\
      --image examples/images/kvasir_sample.jpg \\
      --fixations examples/fixations/kvasir_sample.csv \\
      --output output_mask.png \\
      --save_overlay
        """,
    )
    ap.add_argument("--image",       required=True,
                    help="Path to the input image (.jpg / .jpeg / .png).")
    ap.add_argument("--fixations",   required=True,
                    help="Path to the fixation CSV (x,y,duration — pixel coordinates).")
    ap.add_argument("--output",      required=True,
                    help="Where to save the predicted binary mask (.png).")
    ap.add_argument("--preset",      default="colonoscopy",
                    choices=list(PRESETS),
                    help="Hyperparameter preset: 'colonoscopy' (default) or 'mri'.")
    ap.add_argument("--threshold",   type=float, default=0.5,
                    help="Binarization threshold applied to the soft mask (default: 0.5).")
    ap.add_argument("--save_overlay", action="store_true",
                    help="Also save a colour overlay PNG next to --output.")
    ap.add_argument("--device",      default=None,
                    help="'cuda' or 'cpu' — auto-detected when not given.")
    return ap


def main():
    args = _build_parser().parse_args()

    print(f"[GazeRefine] image    : {args.image}")
    print(f"[GazeRefine] fixations: {args.fixations}")
    print(f"[GazeRefine] preset   : {args.preset}")
    print(f"[GazeRefine] threshold: {args.threshold}")

    result = predict(
        image_path   = args.image,
        fixation_csv = args.fixations,
        preset       = args.preset,
        threshold    = args.threshold,
        device       = args.device,
        return_all   = args.save_overlay,
    )

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

    if isinstance(result, dict):
        result["mask"].save(output_path)
        print(f"[GazeRefine] mask saved  → {output_path}")
        if args.save_overlay:
            overlay_path = output_path.with_stem(output_path.stem + "_overlay")
            result["mask_overlay"].save(overlay_path)
            gaze_path = output_path.with_stem(output_path.stem + "_gaze")
            result["gaze_overlay"].save(gaze_path)
            print(f"[GazeRefine] overlay     → {overlay_path}")
            print(f"[GazeRefine] gaze prior  → {gaze_path}")
    else:
        result.save(output_path)
        print(f"[GazeRefine] mask saved  → {output_path}")


if __name__ == "__main__":
    main()