""" 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 gazerefine import GazeRefine, overlay_heatmap, overlay_mask from gazerefine.constants import IMG_MEAN, IMG_STD, IMG_SIZE from gazerefine.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()