Upload 41 files

.gitattributes

@@ -1,35 +1,35 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,75 @@
----
-title: GazeCorrect
-emoji: 👁
-colorFrom: green
-colorTo: yellow
-sdk: gradio
-sdk_version: 6.24.0
-python_version: '3.12'
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: GazeRefine
+emoji: 👁️
+colorFrom: blue
+colorTo: red
+sdk: gradio
+sdk_version: 4.44.1
+app_file: app.py
+pinned: false
+license: mit
+short_description: Zero-shot, training-free gaze-guided medical segmentation
+---
+
+# GazeRefine — Expert Gaze as a Test-Time Prompt
+
+Interactive demo for **GazeRefine**, a training-free, zero-shot framework
+that turns expert eye-gaze into an inference-time prompt for medical image
+segmentation. Frozen DINOv3 patch features + gaze-anchored
+foreground/background prototypes + recurrent contrastive cleaning + kNN
+affinity propagation — no masks, no clicks-as-boxes, no fine-tuning, no
+adapters, no prompt encoder.
+
+## How to use
+1. Upload a colonoscopy or grayscale-MRI-style image.
+2. Click on the image 1–5 times where a clinician's gaze would land on the
+   structure of interest (a polyp, the prostate, ...). Each click adds a
+   numbered fixation marker; the slider controls that fixation's relative
+   duration/weight before your next click.
+3. Pick a hyperparameter preset (tuned per-modality, see the paper).
+4. Press **Run GazeRefine** to get the gaze-prior overlay and the predicted
+   segmentation mask.
+5. Optionally open **Correct the gaze-attended region**. The predicted mask is
+   feathered, Gaussian noise is added only inside that region, and the noised
+   image is sent to RoentGen-v2 image-to-image inference. The generated result
+   is composited back only inside the mask, retaining the original background.
+
+## Synthetic chest X-ray correction
+
+Install the additional inference dependencies:
+
+```bash
+pip install -U diffusers transformers accelerate
+```
+
+RoentGen-v2 is a gated model. Before using the correction button, sign in to
+Hugging Face and accept the model's access conditions. The app uses
+`StableDiffusionImg2ImgPipeline` (rather than the model-card's text-to-image
+`DiffusionPipeline`) so the masked/noised image can be used as the starting
+image. The model and its outputs are for research/education only, never for
+clinical diagnosis.
+
+### Deploying on Hugging Face Spaces
+
+The Space must have GPU hardware selected. In **Settings → Secrets**, add an
+`HF_TOKEN` secret created by the Hugging Face account that accepted access to
+`stanfordmimi/RoentGen-v2`. The token is read server-side and is never shown in
+the web interface. The first correction request downloads the model, so it can
+take several minutes; the interface displays the attention-weighted Gaussian
+noise preview immediately if RoentGen cannot start.
+
+For a **ZeroGPU** Space, this application already decorates the segmentation
+and correction callbacks with `@spaces.GPU`. Keep `app.py` unchanged at the
+top level; removing those decorators causes the ZeroGPU startup error
+“No @spaces.GPU function detected”.
+
+## Notes
+- Inference uses a frozen `vit_large_patch16_dinov3.lvd1689m` backbone from
+  `timm`. First run will download the checkpoint.
+- CPU inference works but is slow; a GPU Space is recommended for a smooth
+  demo.
+- This Space is for research/demonstration only — it is **not** a clinical
+  diagnostic tool.
+
+Full code, configs, and the unified evaluation pipeline:
+[GitHub repository](https://github.com/<your-org>/gazerefine).

app.py

@@ -0,0 +1,860 @@
+"""
+app.py — GazeRefine interactive demo (Hugging Face Space).
+
+Upload flow
+-----------
+* Standard images (.jpg / .png / etc.) — drag-and-drop or click on the main
+  gr.Image widget.  The same widget also accepts clicks to place fixations, so
+  there is now only ONE image panel instead of two.
+* DICOM files (.dcm) — use the separate "Upload DICOM" file picker.  The file
+  is decoded with pydicom and converted to an RGB PIL image before being handed
+  to the same fixation / run pipeline.
+* Fixation file (.csv / .xlsx / .xls) — use the separate "Upload fixation
+  file" picker.  This can contain fixations for one or many images (e.g. an
+  eye-tracker export with one row per fixation).  After upload, three
+  dropdowns let the user pick which column is the image-name/ID column and
+  which columns hold X / Y (and, optionally, duration).  Rows are matched to
+  the currently loaded image by filename; X/Y values are auto-detected as
+  either normalised [0,1] or raw pixel coordinates.
+"""
+from __future__ import annotations
+
+import sys
+import types
+import tempfile
+import csv
+import os
+from pathlib import Path
+
+# ZeroGPU must see this import and decorator before any module imports torch.
+# The fallback keeps the same app runnable in a regular CPU/GPU Space and
+# during local development, where Hugging Face's `spaces` package is absent.
+try:
+    import spaces
+except ImportError:
+    class _SpacesFallback:
+        @staticmethod
+        def GPU(*args, **kwargs):
+            if args and callable(args[0]) and len(args) == 1 and not kwargs:
+                return args[0]
+            return lambda function: function
+    spaces = _SpacesFallback()
+
+# ── 1. audioop shim (Python 3.13 removed audioop; pydub needs it) ─────────────
+if sys.version_info >= (3, 13):
+    for _mod in ("audioop", "pyaudioop"):
+        if _mod not in sys.modules:
+            sys.modules[_mod] = types.ModuleType(_mod)
+
+# ── 2. Patch starlette Jinja2Templates.TemplateResponse ──────────────────────
+import starlette.templating as _st
+
+_orig_TR = _st.Jinja2Templates.TemplateResponse
+
+def _compat_TR(self, *args, **kwargs):
+    if args and isinstance(args[0], str) and len(args) >= 2 and isinstance(args[1], dict):
+        name       = args[0]
+        context    = args[1]
+        status_code = args[2] if len(args) > 2 else kwargs.get("status_code", 200)
+        headers    = kwargs.get("headers")
+        media_type = kwargs.get("media_type")
+        background = kwargs.get("background")
+        template   = self.get_template(name)
+        return _st._TemplateResponse(
+            template, context,
+            status_code=status_code,
+            headers=headers,
+            media_type=media_type,
+            background=background,
+        )
+    return _orig_TR(self, *args, **kwargs)
+
+_st.Jinja2Templates.TemplateResponse = _compat_TR  # type: ignore[method-assign]
+
+import gradio as gr
+
+# ── 3. gradio_client schema shim ──────────────────────────────────────────────
+try:
+    import gradio_client.utils as _gcu
+    _orig_inner = _gcu._json_schema_to_python_type
+
+    def _safe_inner(schema, defs=None):
+        if not isinstance(schema, dict):
+            return "Any"
+        if not isinstance(schema.get("additionalProperties"), dict):
+            schema = {k: v for k, v in schema.items() if k != "additionalProperties"}
+        return _orig_inner(schema, defs)
+
+    _gcu._json_schema_to_python_type = _safe_inner
+except Exception:
+    pass
+
+# ── 4. huggingface_hub HfFolder shim ─────────────────────────────────────────
+try:
+    from huggingface_hub import HfFolder  # noqa: F401
+except ImportError:
+    import huggingface_hub as _hfh
+    class _FakeHfFolder:
+        @staticmethod
+        def get_token(): return None
+    _hfh.HfFolder = _FakeHfFolder  # type: ignore[attr-defined]
+    sys.modules["huggingface_hub"].HfFolder = _FakeHfFolder  # type: ignore[assignment]
+
+import numpy as np
+from PIL import Image, ImageDraw
+
+# ── 5. Path setup ─────────────────────────────────────────────────────────────
+_here = Path(__file__).resolve().parent
+for _candidate in [_here] + list(_here.parents):
+    _s = str(_candidate)
+    if _s not in sys.path:
+        sys.path.insert(0, _s)
+
+import scripts.predict_single as _predict_module  # noqa: E402
+from scripts.predict_single import predict         # noqa: E402
+from gazecorrect.correction import (               # noqa: E402
+    apply_correction_only_to_attention,
+    correct_with_roentgen,
+    make_attention_noise,
+)
+
+# ── Monkey-patch load_fixation_csv ────────────────────────────────────────────
+# Our single-image temp CSV has x,y,duration in raw pixel coordinates with no
+# image_name column.  The original loader expects a dataset CSV and returns an
+# empty tensor when that column is absent.
+# This patch detects the missing column, reads the CSV directly, normalises
+# pixel → [0,1], and adds the batch dimension the model requires: (N,3)→(1,N,3).
+
+import pandas as _pd
+import torch as _torch
+
+try:
+    from gazecorrect.gaze import load_fixation_csv as _orig_load_fixation_csv
+except Exception:
+    _orig_load_fixation_csv = None
+
+def _patched_load_fixation_csv(csv_path, image_width=1, image_height=1, image_name=None):
+    df = _pd.read_csv(csv_path)
+    print(f"[PATCH] load_fixation_csv — columns: {list(df.columns)}, rows: {len(df)}")
+
+    if "image_name" in df.columns and _orig_load_fixation_csv is not None:
+        print("[PATCH] image_name column present — using original loader")
+        return _orig_load_fixation_csv(csv_path, image_width=image_width,
+                                       image_height=image_height, image_name=image_name)
+
+    x   = df["x"].values.astype(float)
+    y   = df["y"].values.astype(float)
+    dur = df["duration"].values.astype(float)
+
+    x_n   = x   / max(float(image_width),  1.0)
+    y_n   = y   / max(float(image_height), 1.0)
+    dur_n = dur / (dur.max() + 1e-8)
+
+    # model expects (B, N, 3) — add batch dim
+    fixations = _torch.tensor(
+        list(zip(x_n, y_n, dur_n)), dtype=_torch.float32
+    ).unsqueeze(0)  # (N, 3) → (1, N, 3)
+
+    print(f"[PATCH] tensor shape: {tuple(fixations.shape)}")
+    print(f"[PATCH] fixations (x_norm, y_norm, dur_norm):\n{fixations[0]}")
+    return fixations
+
+_predict_module.load_fixation_csv = _patched_load_fixation_csv
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Helpers
+# ─────────────────────────────────────────────────────────────────────────────
+
+PRESETS = {
+    "Colonoscopy / polyp (Kvasir-SEG settings)": "colonoscopy",
+    "Grayscale MRI / CT (prostate-MRI settings)": "mri",
+}
+POINT_COLORS = ["#ff3b30", "#ff9500", "#ffcc00", "#34c759", "#5ac8fa", "#007aff", "#af52de"]
+
+_NO_COL = "— none —"
+
+
+def dcm_to_pil(dcm_path: str) -> Image.Image:
+    """Load a DICOM file and return an RGB PIL image."""
+    import pydicom
+    dcm = pydicom.dcmread(dcm_path)
+    arr = dcm.pixel_array.astype(np.float32)
+    arr = arr - arr.min()
+    arr = arr / (arr.max() + 1e-8)
+    arr = (arr * 255).astype(np.uint8)
+    # Handle multi-frame / greyscale / RGB DICOM
+    if arr.ndim == 2:
+        return Image.fromarray(arr, mode="L").convert("RGB")
+    if arr.ndim == 3 and arr.shape[0] in (1, 3, 4):
+        # (C, H, W) → (H, W, C)
+        arr = arr.transpose(1, 2, 0)
+    return Image.fromarray(arr).convert("RGB")
+
+
+def draw_points(image: Image.Image, points: list) -> Image.Image:
+    """Overlay fixation circles on a copy of `image`.
+
+    `points`: list of (x_px, y_px, duration) in original-image pixel coords.
+    """
+    if image is None:
+        return None
+    vis = image.convert("RGB").copy()
+    draw = ImageDraw.Draw(vis)
+    w, h = vis.size
+    r = max(6, min(w, h) // 80)
+    for i, (x_px, y_px, dur) in enumerate(points):
+        color = POINT_COLORS[i % len(POINT_COLORS)]
+        rad = r * (0.6 + 0.8 * dur)
+        draw.ellipse(
+            [x_px - rad, y_px - rad, x_px + rad, y_px + rad],
+            outline=color, width=3,
+        )
+        draw.text((x_px + rad + 2, y_px - rad), str(i + 1), fill=color)
+    return vis
+
+
+def read_table(path: str) -> "_pd.DataFrame":
+    """Load a .csv / .xlsx / .xls fixation file into a DataFrame."""
+    ext = Path(path).suffix.lower()
+    if ext in (".xlsx", ".xls"):
+        return _pd.read_excel(path)
+    # Sniff delimiter for csv/tsv/txt — eye-tracker exports are sometimes
+    # tab-separated even with a .csv extension.
+    return _pd.read_csv(path, sep=None, engine="python")
+
+
+def normalize_xy(x_vals: np.ndarray, y_vals: np.ndarray, img_w: int, img_h: int):
+    """Convert X/Y column values to pixel coords for the given image size.
+
+    Values already in [0, 1] (inclusive, with a little slack for rounding)
+    are treated as normalised; anything else is assumed to already be raw
+    pixel coordinates and is left as-is (but clamped to the image bounds).
+    """
+    looks_normalized = (
+        np.nanmax(x_vals) <= 1.05 and np.nanmax(y_vals) <= 1.05
+        and np.nanmin(x_vals) >= -0.05 and np.nanmin(y_vals) >= -0.05
+    )
+    if looks_normalized:
+        x_px = np.clip(x_vals, 0, 1) * img_w
+        y_px = np.clip(y_vals, 0, 1) * img_h
+    else:
+        x_px = np.clip(x_vals, 0, img_w)
+        y_px = np.clip(y_vals, 0, img_h)
+    return x_px, y_px
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Event handlers
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+_UPLOAD_LABEL    = "Drop / click to load  .jpg  .png  .bmp  .tif  .tiff  .webp  .dcm"
+_FIXATION_LABEL  = "Click to place fixations"
+_FIXFILE_LABEL   = "Upload fixation file (.csv / .xlsx / .xls) — optional"
+
+
+def _resolve_path(file_obj):
+    """Extract a filesystem path from whatever gr.File passes."""
+    if isinstance(file_obj, str):
+        return file_obj
+    if isinstance(file_obj, dict):
+        return file_obj.get("name") or file_obj.get("path") or file_obj.get("tmp_path") or ""
+    if hasattr(file_obj, "name"):
+        return file_obj.name
+    return ""
+
+
+def on_file_upload(file_obj):
+    """Load any image or DICOM and switch the panel to fixation-click mode."""
+    _no_change = (None, [], gr.update(), gr.update(), gr.update(), gr.update(), None)
+
+    if file_obj is None:
+        return _no_change
+
+    image_name = ""
+
+    # gr.Image gives PIL/numpy; gr.File gives a path
+    if isinstance(file_obj, Image.Image):
+        pil = file_obj.convert("RGB")
+    elif isinstance(file_obj, np.ndarray):
+        pil = Image.fromarray(file_obj).convert("RGB")
+    else:
+        path = _resolve_path(file_obj)
+        if not path:
+            gr.Warning("Could not resolve file path.")
+            return _no_change
+        image_name = Path(path).name
+        ext = Path(path).suffix.lower()
+        try:
+            pil = dcm_to_pil(path) if ext == ".dcm" else Image.open(path).convert("RGB")
+        except Exception as e:
+            gr.Warning(f"Could not load file: {e}")
+            return _no_change
+
+    print(f"[DEBUG] on_file_upload — size={pil.size}  name={image_name!r}")
+    # Switch: hide upload zone, show image panel + delete button
+    return (
+        pil,                                    # orig_image_state
+        [],                                     # points_state
+        image_name,                              # image_name_state
+        gr.update(visible=False),               # upload_zone  → hide
+        gr.update(value=pil, visible=True,
+                  label=_FIXATION_LABEL),       # image_panel  → show with image
+        gr.update(visible=True),                # delete_btn   → show
+        None,                                   # attention_mask_state → reset
+    )
+
+
+def on_select(orig_image: Image.Image, points: list, duration: float, evt: gr.SelectData):
+    """Record a fixation click in original-image pixel coords."""
+    if orig_image is None:
+        gr.Warning("Upload an image first.")
+        return points, gr.update(), None
+    x_px, y_px = float(evt.index[0]), float(evt.index[1])
+    new_points = points + [(x_px, y_px, duration)]
+    print(f"[DEBUG] fixation #{len(new_points)}: x={x_px:.1f}  y={y_px:.1f}  dur={duration}")
+    return new_points, draw_points(orig_image, new_points), None
+
+
+def on_clear(orig_image):
+    """Remove all fixations but keep the current image."""
+    if orig_image is None:
+        return [], gr.update(), None
+    return [], gr.update(value=orig_image), None
+
+
+def on_delete():
+    """Delete the current image and return to upload mode."""
+    return (
+        None,                                   # orig_image_state
+        [],                                     # points_state
+        "",                                      # image_name_state
+        gr.update(value=None, visible=True),    # upload_zone  → show (reset)
+        gr.update(value=None, visible=False),   # image_panel  → hide
+        gr.update(visible=False),               # delete_btn   → hide
+        None,                                   # attention_mask_state → reset
+    )
+
+
+# ── Fixation-file upload → column mapping ────────────────────────────────────
+
+def on_fixfile_upload(file_obj):
+    """Load the fixation table and populate the column-mapping dropdowns."""
+    _hide = (
+        None, gr.update(visible=False),
+        gr.update(choices=[], value=None), gr.update(choices=[], value=None),
+        gr.update(choices=[], value=None), gr.update(choices=[], value=None),
+        gr.update(visible=False),
+    )
+    if file_obj is None:
+        return _hide
+
+    path = _resolve_path(file_obj)
+    if not path:
+        gr.Warning("Could not resolve fixation file path.")
+        return _hide
+
+    try:
+        df = read_table(path)
+    except Exception as e:
+        gr.Warning(f"Could not read fixation file: {e}")
+        return _hide
+
+    if df.empty or len(df.columns) == 0:
+        gr.Warning("Fixation file appears to be empty.")
+        return _hide
+
+    cols = [str(c) for c in df.columns]
+    print(f"[DEBUG] fixation file loaded — columns: {cols}, rows: {len(df)}")
+
+    def _guess(*keywords, fallback=None):
+        for c in cols:
+            cl = c.lower()
+            if any(k in cl for k in keywords):
+                return c
+        return fallback if fallback is not None else cols[0]
+
+    guess_id  = _guess("image", "id", "name", "file", fallback=cols[0])
+    # exact / boundary-aware matches first (avoids "fix_index" matching "x"),
+    # then fall back to a bare trailing "x" / "y".
+    guess_x = _guess("fix_x", "pos_x", "gaze_x", fallback=None)
+    if guess_x is None:
+        guess_x = next((c for c in cols if c.lower().rstrip("_") .endswith("x")
+                         and "index" not in c.lower()), cols[0])
+    guess_y = _guess("fix_y", "pos_y", "gaze_y", fallback=None)
+    if guess_y is None:
+        guess_y = next((c for c in cols if c.lower().rstrip("_").endswith("y")
+                         and "index" not in c.lower()), cols[0])
+
+    dur_choices = [_NO_COL] + cols
+    guess_dur = _guess("duration", "dur", fallback=_NO_COL)
+
+    return (
+        df.to_json(),                                          # fixfile_df_state (serialized)
+        gr.update(visible=True),                                # mapping_row → show
+        gr.update(choices=cols, value=guess_id),                # id_col_dd
+        gr.update(choices=cols, value=guess_x),                 # x_col_dd
+        gr.update(choices=cols, value=guess_y),                 # y_col_dd
+        gr.update(choices=dur_choices, value=guess_dur),         # dur_col_dd
+        gr.update(visible=True),                                # apply_fix_btn → show
+    )
+
+
+def on_apply_fixfile(fixfile_json, id_col, x_col, y_col, dur_col,
+                      orig_image, image_name):
+    """Match rows to the currently loaded image (by filename) and load
+    them as fixation points, replacing whatever points are currently set.
+
+    If no rows match the loaded image's filename, nothing is loaded — the
+    existing points (if any) are left untouched, and the user is warned so
+    they can check the ID column / image filename instead of silently
+    getting fixations for the wrong image."""
+    if orig_image is None:
+        gr.Warning("Load an image first, then apply the fixation file.")
+        return gr.update(), gr.update(), None
+    if not fixfile_json:
+        gr.Warning("Upload a fixation file first.")
+        return gr.update(), gr.update(), None
+    if not id_col or not x_col or not y_col:
+        gr.Warning("Pick the ID, X and Y columns first.")
+        return gr.update(), gr.update(), None
+    if not image_name:
+        gr.Warning(
+            "Couldn't determine the loaded image's filename (this can "
+            "happen if the image was pasted/dropped without a filename). "
+            "Re-upload the image as a file and try again."
+        )
+        return gr.update(), gr.update(), None
+
+    df = _pd.read_json(fixfile_json)
+
+    # Match by exact filename first, then by stem-without-extension, so the
+    # mapping still works if the fixation file's IMAGE column omits the
+    # extension or uses a different one than the uploaded image.
+    mask = df[id_col].astype(str) == image_name
+    if not mask.any():
+        stem_no_ext = Path(image_name).stem
+        mask = df[id_col].astype(str).apply(lambda v: Path(str(v)).stem) == stem_no_ext
+
+    sub = df[mask]
+
+    if sub.empty:
+        gr.Warning(
+            f"No rows in the fixation file match the loaded image "
+            f"('{image_name}'). Nothing was loaded — check that the ID "
+            f"column values match the image filename."
+        )
+        return gr.update(), gr.update(), None
+
+    w, h = orig_image.size
+    x_vals = sub[x_col].astype(float).to_numpy()
+    y_vals = sub[y_col].astype(float).to_numpy()
+    x_px, y_px = normalize_xy(x_vals, y_vals, w, h)
+
+    if dur_col and dur_col != _NO_COL and dur_col in sub.columns:
+        dur_raw = sub[dur_col].astype(float).to_numpy()
+        dmax = float(np.nanmax(dur_raw)) if len(dur_raw) else 1.0
+        dur_n = dur_raw / (dmax + 1e-8)
+    else:
+        dur_n = np.full(len(sub), 1.0)
+
+    new_points = [
+        (float(xp), float(yp), float(d))
+        for xp, yp, d in zip(x_px, y_px, dur_n)
+    ]
+    print(f"[DEBUG] loaded {len(new_points)} fixations from file for image '{image_name}'")
+
+    return new_points, draw_points(orig_image, new_points), None
+
+
+@spaces.GPU(duration=120)
+def run(orig_image: Image.Image, points: list, preset_name: str, threshold: float):
+    import traceback, uuid
+    print(f"[DEBUG] run — points={len(points)}  preset={preset_name}")
+
+    if orig_image is None:
+        gr.Warning("Upload an image first.")
+        return None, None, None, None, "Upload an image first."
+    if not points:
+        gr.Warning("Click on the image at least once to place a fixation (or load a fixation file).")
+        return None, None, None, None, "Add at least one gaze fixation, then run segmentation."
+
+    preset_key = PRESETS[preset_name]
+    w, h = orig_image.size
+
+    shared_stem       = f"gazerefine_{uuid.uuid4().hex}"
+    tmp_img_path      = os.path.join(tempfile.gettempdir(), f"{shared_stem}.png")
+    fixation_csv_path = os.path.join(tempfile.gettempdir(), f"{shared_stem}.csv")
+
+    orig_image.convert("RGB").save(tmp_img_path)
+
+    with open(fixation_csv_path, "w", newline="") as f:
+        writer = csv.writer(f)
+        writer.writerow(["x", "y", "duration"])
+        for x_px, y_px, dur in points:
+            writer.writerow([x_px, y_px, dur])
+
+    print(f"[DEBUG] image {w}x{h}  |  {len(points)} fixations  |  preset={preset_key}  thr={threshold}")
+
+    with open(fixation_csv_path) as f:
+        print(f"[DEBUG] CSV:\n{f.read()}")
+
+    try:
+        out = predict(
+            image_path=tmp_img_path,
+            fixation_csv=fixation_csv_path,
+            preset=preset_key,
+            threshold=threshold,
+            return_all=True,
+        )
+    except Exception as e:
+        print(f"[ERROR] predict() raised: {e}")
+        traceback.print_exc()
+        message = f"Segmentation failed: {e}"
+        gr.Warning(message)
+        return None, None, None, None, message
+    finally:
+        for p in (tmp_img_path, fixation_csv_path):
+            try:
+                os.unlink(p)
+            except OSError:
+                pass
+
+    mask_arr = np.array(out["mask"])
+    print(f"[DEBUG] mask non-zero: {(mask_arr > 0).sum()} / {mask_arr.size}")
+    return (
+        out["gaze_overlay"], out["mask_overlay"], out["mask"], out["mask"],
+        "Segmentation completed. You can now apply attention-weighted Gaussian noise and correction.",
+    )
+
+
+@spaces.GPU(duration=180)
+def correct_attention_region(
+    orig_image: Image.Image,
+    attention_mask: Image.Image,
+    noise_strength: float,
+    mask_feather: float,
+    diffusion_strength: float,
+    steps: int,
+    prompt: str,
+    seed: float,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Noise the DINOv3/gaze-attended region, then restore it with RoentGen."""
+    if orig_image is None or attention_mask is None:
+        message = "Run GazeRefine first so a gaze-guided attention mask is available."
+        gr.Warning(message)
+        return None, None, message
+    seed_value = int(seed) if seed >= 0 else None
+    try:
+        progress(0.05, desc="Applying Gaussian noise to the attended region")
+        noised, feathered_mask = make_attention_noise(
+            orig_image, attention_mask, noise_strength, mask_feather, seed_value
+        )
+    except Exception as exc:
+        message = f"Could not create the attention-weighted noise image: {exc}"
+        print(f"[ERROR] {message}")
+        gr.Warning(message)
+        return None, None, message
+
+    try:
+        progress(0.15, desc="Loading RoentGen-v2 (first run downloads model weights)")
+        generated = correct_with_roentgen(
+            noised,
+            prompt=prompt.strip() or "chest x-ray",
+            strength=diffusion_strength,
+            steps=int(steps),
+            seed=seed_value,
+        )
+        progress(0.95, desc="Compositing corrected attention region")
+        corrected = apply_correction_only_to_attention(orig_image, generated, feathered_mask)
+        return noised, corrected, "Completed: RoentGen-v2 corrected the gaze-attended region."
+    except Exception as exc:
+        message = (
+            "Gaussian-noise preview created, but RoentGen-v2 could not run. "
+            "In your Hugging Face Space, add a Settings → Secrets value named "
+            "HF_TOKEN, accept access for stanfordmimi/RoentGen-v2 with that account, "
+            f"and use GPU hardware. Error: {exc}"
+        )
+        print(f"[ERROR] diffusion correction failed: {exc}")
+        gr.Warning(message)
+        return noised, None, message
+
+
+@spaces.GPU(duration=300)
+def segment_noise_and_regenerate(
+    orig_image: Image.Image,
+    points: list,
+    preset_name: str,
+    threshold: float,
+    disease_description: str,
+    noise_strength: float,
+    mask_feather: float,
+    diffusion_strength: float,
+    steps: int,
+    seed: float,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """One-click workflow: gaze segmentation -> masked noise -> regeneration."""
+    if orig_image is None:
+        return None, None, None, None, None, None, "Upload an image first."
+    if not points:
+        return None, None, None, None, None, None, "Add gaze clicks or load a gaze CSV first."
+
+    import traceback
+    import uuid
+
+    shared_stem = f"gazerefine_{uuid.uuid4().hex}"
+    tmp_img_path = os.path.join(tempfile.gettempdir(), f"{shared_stem}.png")
+    fixation_csv_path = os.path.join(tempfile.gettempdir(), f"{shared_stem}.csv")
+    try:
+        progress(0.03, desc="Preparing gaze input")
+        orig_image.convert("RGB").save(tmp_img_path)
+        with open(fixation_csv_path, "w", newline="") as file:
+            writer = csv.writer(file)
+            writer.writerow(["x", "y", "duration"])
+            writer.writerows(points)
+
+        progress(0.12, desc="DINOv3 gaze-guided segmentation")
+        out = predict(
+            image_path=tmp_img_path,
+            fixation_csv=fixation_csv_path,
+            preset=PRESETS[preset_name],
+            threshold=threshold,
+            return_all=True,
+        )
+    except Exception as exc:
+        traceback.print_exc()
+        message = f"Segmentation failed: {exc}"
+        gr.Warning(message)
+        return None, None, None, None, None, None, message
+    finally:
+        for path in (tmp_img_path, fixation_csv_path):
+            try:
+                os.unlink(path)
+            except OSError:
+                pass
+
+    noised = None
+    try:
+        seed_value = int(seed) if seed >= 0 else None
+        progress(0.45, desc="Adding attention-weighted Gaussian noise")
+        noised, feathered_mask = make_attention_noise(
+            orig_image, out["mask"], noise_strength, mask_feather, seed_value
+        )
+        progress(0.55, desc="RoentGen-v2 regeneration")
+        generated = correct_with_roentgen(
+            noised,
+            prompt=disease_description.strip() or "Normal chest radiograph.",
+            strength=diffusion_strength,
+            steps=int(steps),
+            seed=seed_value,
+        )
+        corrected = apply_correction_only_to_attention(orig_image, generated, feathered_mask)
+        return (
+            out["gaze_overlay"], out["mask_overlay"], out["mask"], out["mask"],
+            noised, corrected,
+            "Completed: gaze mask, attention-weighted noise, and regenerated image are ready.",
+        )
+    except Exception as exc:
+        message = (
+            "The gaze mask and noise image were created, but RoentGen-v2 is not authorized. "
+            "Accept access to the model and set HF_TOKEN. Error: " + str(exc)
+        )
+        gr.Warning(message)
+        return out["gaze_overlay"], out["mask_overlay"], out["mask"], out["mask"], noised, None, message
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# UI
+# ─────────────────────────────────────────────────────────────────────────────
+
+with gr.Blocks(title="GazeRefine — gaze-guided zero-shot segmentation") as demo:
+    gr.Markdown(
+        """
+        # 👁️ GazeRefine — Expert Gaze as a Test-Time Prompt
+        Training-free, zero-shot medical image segmentation.
+        Upload an image or DICOM, click to place fixations (or load a fixation
+        file), then hit **Run**.
+        """
+    )
+
+    orig_image_state = gr.State(None)
+    points_state      = gr.State([])
+    image_name_state  = gr.State("")     # filename of the currently loaded image
+    fixfile_df_state  = gr.State(None)   # serialized DataFrame (to_json) of the uploaded fixation file
+    attention_mask_state = gr.State(None)  # binary DINOv3/gaze segmentation output
+
+    with gr.Row():
+        # ── Left column ───────────────────────────────────────────────────────
+        with gr.Column(scale=1):
+
+            # ── Upload zone (visible when no image loaded) ────────────────────
+            upload_zone = gr.File(
+                label=_UPLOAD_LABEL,
+                file_types=[".jpg", ".jpeg", ".png", ".bmp",
+                            ".tif", ".tiff", ".webp", ".gif", ".dcm"],
+                file_count="single",
+                visible=True,
+                elem_id="upload_zone",
+            )
+
+            # ── Image panel (hidden until image loaded; click to fixate) ──────
+            image_panel = gr.Image(
+                type="pil",
+                label=_FIXATION_LABEL,
+                height=430,
+                interactive=False,          # no toolbar → .select fires on click
+                show_download_button=False,
+                visible=False,
+                elem_id="image_panel",
+            )
+
+            # ── Delete button (hidden until image loaded) ─────────────────────
+            delete_btn = gr.Button("🗑 Delete image — load another", visible=False, variant="secondary")
+
+            # ── Fixation file upload (optional alternative to manual clicks) ──
+            with gr.Accordion("📄 Load fixations from file", open=False):
+                fixfile_upload = gr.File(
+                    label=_FIXFILE_LABEL,
+                    file_types=[".csv", ".xlsx", ".xls", ".tsv", ".txt"],
+                    file_count="single",
+                    elem_id="fixfile_upload",
+                )
+                with gr.Row(visible=False) as mapping_row:
+                    id_col_dd  = gr.Dropdown(label="Image / ID column", choices=[])
+                    x_col_dd   = gr.Dropdown(label="X column", choices=[])
+                    y_col_dd   = gr.Dropdown(label="Y column", choices=[])
+                    dur_col_dd = gr.Dropdown(label="Duration column (optional)", choices=[])
+                apply_fix_btn = gr.Button(
+                    "📥 Load fixations for current image", visible=False,
+                )
+
+            # ── Controls ──────────────────────────────────────────────────────
+            with gr.Row():
+                duration_slider = gr.Slider(
+                    0.1, 1.0, value=1.0, step=0.1,
+                    label="Fixation duration weight",
+                )
+                clear_btn = gr.Button("✖ Clear fixations")
+
+            preset = gr.Radio(
+                list(PRESETS.keys()), value=list(PRESETS.keys())[0],
+                label="Preset",
+            )
+            threshold = gr.Slider(
+                0.1, 0.9, value=0.5, step=0.05,
+                label="Mask threshold",
+            )
+            run_btn = gr.Button("▶ Run GazeRefine", variant="primary")
+            segmentation_status = gr.Textbox(
+                label="Segmentation status", value="Upload an image and add gaze points.",
+                interactive=False, lines=2,
+            )
+
+            with gr.Accordion("🩻 Correct the gaze-attended region (RoentGen-v2)", open=False):
+                gr.Markdown(
+                    "After segmentation, Gaussian noise is applied only inside the "
+                    "gaze-guided mask. RoentGen-v2 receives that image for image-to-image "
+                    "restoration; unmasked pixels are kept from the source image. Research use only."
+                )
+                correction_prompt = gr.Textbox(
+                    value="Normal chest radiograph.",
+                    label="Disease / radiology description for regeneration",
+                    placeholder="Example: Right lower-lobe opacity. No pleural effusion.",
+                )
+                with gr.Row():
+                    noise_strength = gr.Slider(0.0, 1.0, value=0.35, step=0.05, label="Gaussian noise")
+                    mask_feather = gr.Slider(0, 20, value=4, step=1, label="Mask feather (pixels)")
+                with gr.Row():
+                    diffusion_strength = gr.Slider(0.05, 1.0, value=0.35, step=0.05, label="Diffusion strength")
+                    diffusion_steps = gr.Slider(10, 75, value=30, step=1, label="Diffusion steps")
+                    diffusion_seed = gr.Number(value=42, precision=0, label="Seed (-1 = random)")
+                correct_btn = gr.Button("✨ Noise attended region and correct", variant="secondary")
+                one_click_btn = gr.Button(
+                    "▶ Segment → noise attended region → regenerate", variant="primary"
+                )
+                correction_status = gr.Textbox(
+                    label="Correction status", value="Run GazeRefine, then start correction.",
+                    interactive=False, lines=3,
+                )
+
+        # ── Right column: outputs ─────────────────────────────────────────────
+        with gr.Column(scale=1):
+            gaze_out = gr.Image(label="Gaze prior", height=260)
+            with gr.Row():
+                mask_overlay_out = gr.Image(label="Mask overlay", height=260)
+                mask_only_out    = gr.Image(label="Binary mask",  height=260)
+            with gr.Row():
+                noised_out = gr.Image(label="Gaussian-noised attended region", height=260)
+                corrected_out = gr.Image(label="RoentGen corrected image", height=260)
+
+    # ── Event wiring ──────────────────────────────────────────────────────────
+
+    _upload_outputs = [orig_image_state, points_state, image_name_state,
+                        upload_zone, image_panel, delete_btn, attention_mask_state]
+
+    upload_zone.upload(on_file_upload, inputs=[upload_zone], outputs=_upload_outputs)
+    upload_zone.change(on_file_upload, inputs=[upload_zone], outputs=_upload_outputs)
+
+    image_panel.select(
+        on_select,
+        inputs=[orig_image_state, points_state, duration_slider],
+        outputs=[points_state, image_panel, attention_mask_state],
+    )
+
+    clear_btn.click(
+        on_clear,
+        inputs=[orig_image_state],
+        outputs=[points_state, image_panel, attention_mask_state],
+    )
+
+    delete_btn.click(
+        on_delete,
+        outputs=[orig_image_state, points_state, image_name_state,
+                 upload_zone, image_panel, delete_btn, attention_mask_state],
+    )
+
+    _fixfile_outputs = [fixfile_df_state, mapping_row, id_col_dd, x_col_dd, y_col_dd, dur_col_dd, apply_fix_btn]
+
+    fixfile_upload.upload(on_fixfile_upload, inputs=[fixfile_upload], outputs=_fixfile_outputs)
+    fixfile_upload.change(on_fixfile_upload, inputs=[fixfile_upload], outputs=_fixfile_outputs)
+
+    apply_fix_btn.click(
+        on_apply_fixfile,
+        inputs=[fixfile_df_state, id_col_dd, x_col_dd, y_col_dd, dur_col_dd,
+                orig_image_state, image_name_state],
+        outputs=[points_state, image_panel, attention_mask_state],
+    )
+
+    run_btn.click(
+        run,
+        inputs=[orig_image_state, points_state, preset, threshold],
+        outputs=[gaze_out, mask_overlay_out, mask_only_out, attention_mask_state,
+                 segmentation_status],
+    )
+
+    correct_btn.click(
+        correct_attention_region,
+        inputs=[orig_image_state, attention_mask_state, noise_strength, mask_feather,
+                diffusion_strength, diffusion_steps, correction_prompt, diffusion_seed],
+        outputs=[noised_out, corrected_out, correction_status],
+    )
+
+    one_click_btn.click(
+        segment_noise_and_regenerate,
+        inputs=[orig_image_state, points_state, preset, threshold, correction_prompt,
+                noise_strength, mask_feather, diffusion_strength, diffusion_steps,
+                diffusion_seed],
+        outputs=[gaze_out, mask_overlay_out, mask_only_out, attention_mask_state,
+                 noised_out, corrected_out, correction_status],
+    )
+
+    gr.Markdown(
+        "Segmentation: frozen DINOv3 + gaze-anchored prototypes + recurrent "
+        "foreground/background refinement. Correction: attention-local Gaussian noise "
+        "followed by optional RoentGen-v2 image-to-image restoration."
+    )
+
+demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)

gazecorrect/__init__.py

@@ -0,0 +1,32 @@
+"""
+GazeRefine
+==========
+Expert gaze as a test-time prompt for training-free medical image segmentation.
+
+Quick start
+-----------
+    from scripts.predict_single import predict
+    mask = predict(image_path="image.png", fixation_csv="fixations.csv")
+    mask.save("output_mask.png")
+"""
+
+from .backbone  import FrozenDINOv3
+from .gaze      import get_scanpath, load_fixation_csv, generate_gaze_heatmap
+from .model     import GazeRefine, knn_affinity_refinement
+from .metrics   import compute_metrics
+from .visualize import save_prediction, overlay_heatmap, overlay_mask
+
+__all__ = [
+    "FrozenDINOv3",
+    "get_scanpath",
+    "load_fixation_csv",
+    "generate_gaze_heatmap",
+    "GazeRefine",
+    "knn_affinity_refinement",
+    "compute_metrics",
+    "save_prediction",
+    "overlay_heatmap",
+    "overlay_mask",
+]
+
+__version__ = "0.1.0"

gazecorrect/backbone.py

@@ -0,0 +1,78 @@
+"""
+backbone.py — Frozen DINOv3 feature extractor.
+
+GazeRefine never updates the backbone. We load a pretrained DINOv3 ViT via
+`timm`, freeze every parameter, and pull out raw patch tokens from one or
+more transformer blocks using forward hooks. No adapters, no projections,
+no fine-tuning.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class FrozenDINOv3(nn.Module):
+    """Completely frozen DINOv3 ViT backbone that exposes raw patch tokens.
+
+    Parameters
+    ----------
+    model_name : str
+        Any DINOv3 variant available in `timm` (e.g.
+        ``"vit_large_patch16_dinov3.lvd1689m"``,
+        ``"vit_base_patch16_dinov3.lvd1689m"``). Larger backbones generally
+        give cleaner semantic separation but cost more memory/compute.
+    extract_mode : {"last", "all"}
+        - ``"last"``: use only the final block's patch tokens (fast, the
+          default used in our reported results).
+        - ``"all"``: pool patch tokens from 4 evenly-spaced blocks
+          (0, n/4, 3n/4, n-1) and average the resulting similarity maps in
+          ``GazeRefine``. This sometimes helps on harder modalities at the
+          cost of ~4x compute.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "vit_large_patch16_dinov3.lvd1689m",
+        extract_mode: str = "last",
+    ):
+        super().__init__()
+        import timm  # local import: keeps `timm` optional for users who only read the code
+
+        print(f"[GazeRefine] Loading frozen backbone: {model_name} (extract_mode={extract_mode})")
+        bb = timm.create_model(model_name, pretrained=True, num_classes=0)
+        for p in bb.parameters():
+            p.requires_grad_(False)
+        bb.eval()
+
+        self.backbone = bb
+        self.embed_dim = bb.embed_dim
+        self.num_blocks = len(bb.blocks)
+
+        if extract_mode == "last":
+            self.levels = [-1]
+        elif extract_mode == "all":
+            self.levels = [0, self.num_blocks // 4, (self.num_blocks * 3) // 4, self.num_blocks - 1]
+        else:
+            raise ValueError(f"extract_mode must be 'last' or 'all', got {extract_mode!r}")
+
+        print(f"[GazeRefine] Hooking transformer blocks at levels: {self.levels}")
+        self._feats: dict[int, torch.Tensor] = {}
+        for lvl in self.levels:
+            bb.blocks[lvl].register_forward_hook(self._make_hook(lvl))
+
+    def _make_hook(self, lvl: int):
+        def _hook_fn(module, inp, out):
+            # DINOv3 token layout: [CLS, register_1..register_k, patch_1..patch_N]
+            n_prefix = 1 + getattr(self.backbone, "num_register_tokens", 4)
+            self._feats[lvl] = out[:, n_prefix:, :]  # (B, N, D) — patch tokens only
+        return _hook_fn
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
+        """Run the frozen backbone and return a list of (B, N, D) patch-token tensors,
+        one per hooked level, in the order given by ``self.levels``."""
+        self._feats.clear()
+        self.backbone(x)
+        return [self._feats[lvl] for lvl in self.levels]

gazecorrect/constants.py

@@ -0,0 +1,13 @@
+"""constants.py — shared sizing and normalization constants.
+
+IMG_SIZE is fixed to an exact multiple of the ViT patch size so the patch
+grid divides evenly with no rounding/cropping artifacts.
+"""
+
+PATCH_SIZE = 16
+H_PATCH = 37 * 2                      # 74 patches per side
+IMG_SIZE = H_PATCH * PATCH_SIZE       # 1184 px  (use 518 = 37*14 if you switch to a /14 ViT)
+N_PATCHES = H_PATCH ** 2
+
+IMG_MEAN = [0.485, 0.456, 0.406]
+IMG_STD = [0.229, 0.224, 0.225]

gazecorrect/correction.py

@@ -0,0 +1,112 @@
+"""Gaze-mask-guided image correction with an image-to-image diffusion model.
+
+The segmentation mask is deliberately used only to localise corruption: pixels
+outside the attended region are preserved in the source image before the
+diffusion model is invoked.  This is a research workflow, not a diagnostic or
+clinical image reconstruction method.
+"""
+from __future__ import annotations
+
+from functools import lru_cache
+import os
+
+import numpy as np
+import torch
+from PIL import Image, ImageFilter
+
+
+ROENTGEN_MODEL = "stanfordmimi/RoentGen-v2"
+
+
+def make_attention_noise(
+    image: Image.Image,
+    mask: Image.Image | np.ndarray,
+    noise_strength: float = 0.35,
+    blur_radius: float = 4.0,
+    seed: int | None = None,
+) -> tuple[Image.Image, Image.Image]:
+    """Return an image with Gaussian noise inside a feathered attention mask.
+
+    ``noise_strength`` is the standard deviation on the 0--255 image scale.
+    The returned mask is the feathered mask actually used for compositing.
+    """
+    source = np.asarray(image.convert("RGB"), dtype=np.float32)
+    if isinstance(mask, Image.Image):
+        mask_image = mask.convert("L")
+    else:
+        values = np.asarray(mask, dtype=np.float32)
+        if values.max(initial=0.0) <= 1.0:
+            values = values * 255.0
+        mask_image = Image.fromarray(np.clip(values, 0, 255).astype(np.uint8), "L")
+    mask_image = mask_image.resize(image.size, Image.Resampling.BILINEAR)
+    if blur_radius > 0:
+        mask_image = mask_image.filter(ImageFilter.GaussianBlur(radius=blur_radius))
+
+    alpha = np.asarray(mask_image, dtype=np.float32)[..., None] / 255.0
+    rng = np.random.default_rng(seed)
+    noise = rng.normal(0.0, 255.0 * float(noise_strength), size=source.shape)
+    noised = np.clip(source + alpha * noise, 0, 255).astype(np.uint8)
+    return Image.fromarray(noised, "RGB"), mask_image
+
+
+def _device_and_dtype() -> tuple[str, torch.dtype]:
+    if torch.cuda.is_available():
+        return "cuda", torch.float16
+    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+        return "mps", torch.float32
+    return "cpu", torch.float32
+
+
+@lru_cache(maxsize=1)
+def _load_roentgen_pipeline():
+    """Load lazily so segmentation remains usable without diffusers installed."""
+    try:
+        from diffusers import StableDiffusionImg2ImgPipeline
+    except ImportError as exc:
+        raise RuntimeError(
+            "Diffusion correction requires `diffusers`, `transformers`, and "
+            "`accelerate`. Install the project's updated requirements first."
+        ) from exc
+
+    device, dtype = _device_and_dtype()
+    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACEHUB_API_TOKEN")
+    # RoentGen-v2 is a Stable Diffusion 2.1 fine-tune.  Its published example
+    # shows text-to-image with DiffusionPipeline; this image-to-image variant
+    # loads the same weights but accepts the noised attention image as `image`.
+    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+        ROENTGEN_MODEL, torch_dtype=dtype, token=token
+    )
+    return pipe.to(device)
+
+
+def correct_with_roentgen(
+    noised_image: Image.Image,
+    prompt: str = "chest x-ray",
+    strength: float = 0.35,
+    guidance_scale: float = 3.5,
+    steps: int = 30,
+    seed: int | None = None,
+) -> Image.Image:
+    """Use RoentGen-v2 image-to-image inference to restore a noised region."""
+    if not 0.0 < strength <= 1.0:
+        raise ValueError("Diffusion strength must be in (0, 1].")
+    pipe = _load_roentgen_pipeline()
+    device, _ = _device_and_dtype()
+    generator = None if seed is None else torch.Generator(device=device).manual_seed(int(seed))
+    result = pipe(
+        prompt=prompt,
+        image=noised_image.convert("RGB"),
+        strength=float(strength),
+        guidance_scale=float(guidance_scale),
+        num_inference_steps=int(steps),
+        generator=generator,
+    )
+    return result.images[0].convert("RGB")
+
+
+def apply_correction_only_to_attention(
+    source_image: Image.Image, generated_image: Image.Image, feathered_mask: Image.Image
+) -> Image.Image:
+    """Keep un-attended anatomy exactly from the input image after diffusion."""
+    generated = generated_image.convert("RGB").resize(source_image.size, Image.Resampling.LANCZOS)
+    return Image.composite(generated, source_image.convert("RGB"), feathered_mask)

gazecorrect/datasets.py

@@ -0,0 +1,154 @@
+"""
+datasets.py — Image/mask/fixation loaders for the two reported benchmarks,
+plus a small base class so adding a third dataset is just a few lines.
+
+Expected folder layout (same for both datasets):
+
+    <root>/
+        images/    one image per case (.jpg/.png for Kvasir, .dcm for prostate MRI)
+        masks/     matching binary mask, same basename, .png
+    <fixation_csv>
+        one row per fixation, see gazerefine.gaze for the expected columns.
+        The IMAGE column must match an images/ filename (Kvasir) or the
+        DICOM basename + ".jpg" (prostate — fixation collection was run on
+        JPEG-rendered slices while the model reads the original DICOM).
+"""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+import pandas as pd
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+import torchvision.transforms as T
+
+from .constants import IMG_SIZE, IMG_MEAN, IMG_STD
+from .gaze import get_scanpath, IMG_COL
+
+
+class _BaseGazeDataset(Dataset):
+    """Shared image/mask transform + fixation-grouping logic.
+
+    Subclasses only need to implement ``_load_image(name)`` and provide the
+    set of valid image ids that have a matching mask and fixation entries.
+    """
+
+    def __init__(self, root: str, fixation_csv: str, img_size: int = IMG_SIZE):
+        self.root = root
+        self.img_dir = os.path.join(root, "images")
+        self.mask_dir = os.path.join(root, "masks")
+        self.img_size = img_size
+
+        self.df = pd.read_csv(fixation_csv)
+        self.df.columns = self.df.columns.str.strip()
+        self.fix_df = self.df.groupby(IMG_COL)
+        self.max_len = int(self.df.groupby(IMG_COL).size().max())
+        print(f"[Dataset] max scanpath length = {self.max_len}")
+
+        self.image_ids: list[str] = []  # set by subclass __init__
+
+        self.img_tf = T.Compose([
+            T.Resize((img_size, img_size)),
+            T.ToTensor(),
+            T.Normalize(IMG_MEAN, IMG_STD),
+        ])
+        self.mask_tf = T.Compose([
+            T.Resize((img_size, img_size), interpolation=T.InterpolationMode.NEAREST),
+            T.ToTensor(),
+        ])
+
+    def __len__(self):
+        return len(self.image_ids)
+
+    def _load_image(self, name: str) -> Image.Image:
+        raise NotImplementedError
+
+    def _mask_path(self, name: str) -> str:
+        raise NotImplementedError
+
+    def _fixation_key(self, name: str) -> str:
+        """CSV IMAGE-column key for this sample. Override if it differs from
+        the on-disk basename (e.g. prostate MRI uses .jpg keys for .dcm files)."""
+        return name
+
+    def __getitem__(self, idx: int):
+        name = self.image_ids[idx]
+
+        image = self.img_tf(self._load_image(name))
+
+        mask = self.mask_tf(Image.open(self._mask_path(name)).convert("L"))
+        mask = (mask > 0.5).float()
+
+        fix_rows = self.fix_df.get_group(self._fixation_key(name))
+        fixation = get_scanpath(fix_rows, self.max_len)
+
+        return {"image": image, "fixation": fixation, "mask": mask, "name": name}
+
+
+class KvasirSEGDataset(_BaseGazeDataset):
+    """Kvasir-SEG colonoscopy polyp segmentation. images/masks share filenames
+    (e.g. ``cju0qkwl35piu0993l0dewei2.jpg`` in both folders)."""
+
+    def __init__(self, root: str, fixation_csv: str, img_size: int = IMG_SIZE):
+        super().__init__(root, fixation_csv, img_size)
+        img_files = set(os.listdir(self.img_dir))
+        mask_files = set(os.listdir(self.mask_dir))
+        csv_imgs = set(self.df[IMG_COL].unique())
+        self.image_ids = sorted(img_files & mask_files & csv_imgs)
+        if not self.image_ids:
+            raise RuntimeError("No overlap between images/, masks/ and the fixation CSV.")
+        print(f"[Dataset] Kvasir-SEG valid samples = {len(self.image_ids)}")
+
+    def _load_image(self, name: str) -> Image.Image:
+        return Image.open(os.path.join(self.img_dir, name)).convert("RGB")
+
+    def _mask_path(self, name: str) -> str:
+        return os.path.join(self.mask_dir, name)
+
+
+class ProstateMRIDataset(_BaseGazeDataset):
+    """NCI-ISBI prostate MRI. images/ holds DICOM (.dcm), masks/ holds PNG,
+    and the fixation CSV references each case as ``<basename>.jpg`` (the
+    format the eye-tracking session was actually rendered/displayed in)."""
+
+    def __init__(self, root: str, fixation_csv: str, img_size: int = IMG_SIZE):
+        super().__init__(root, fixation_csv, img_size)
+
+        dcm_basenames = {os.path.splitext(f)[0] for f in os.listdir(self.img_dir) if f.endswith(".dcm")}
+        png_basenames = {os.path.splitext(f)[0] for f in os.listdir(self.mask_dir) if f.endswith(".png")}
+        csv_basenames = {os.path.splitext(f)[0] for f in self.df[IMG_COL].unique()}
+
+        self.image_ids = sorted(dcm_basenames & png_basenames & csv_basenames)
+        if not self.image_ids:
+            raise RuntimeError("No overlap between images/ (.dcm), masks/ (.png) and the fixation CSV.")
+        print(f"[Dataset] Prostate MRI valid samples = {len(self.image_ids)}")
+
+    def _load_image(self, name: str) -> Image.Image:
+        import pydicom  # local import: optional dependency, only needed for DICOM datasets
+
+        dcm_path = os.path.join(self.img_dir, f"{name}.dcm")
+        dicom = pydicom.dcmread(dcm_path)
+        arr = dicom.pixel_array.astype(np.float32)
+
+        if arr.max() > arr.min():
+            arr = (arr - arr.min()) / (arr.max() - arr.min())
+        else:
+            arr = np.zeros_like(arr)
+
+        rgb = np.stack([arr, arr, arr], axis=-1)  # grayscale -> 3-channel, DINOv3 expects RGB
+        return Image.fromarray((rgb * 255).astype(np.uint8))
+
+    def _mask_path(self, name: str) -> str:
+        return os.path.join(self.mask_dir, f"{name}.png")
+
+    def _fixation_key(self, name: str) -> str:
+        return f"{name}.jpg"
+
+
+DATASET_REGISTRY = {
+    "kvasir": KvasirSEGDataset,
+    "prostate_mri": ProstateMRIDataset,
+}

gazecorrect/gaze.py

@@ -0,0 +1,191 @@
+"""
+gaze.py — Turn an eye-tracking scanpath into a spatial prior.
+
+Implements the gaze heatmap from the paper:
+
+    H(u, v) = Σ_m  d̃_m · exp( -((u-xm)² + (v-ym)²) / 2σ² )
+
+i.e. a duration-weighted sum of 2D Gaussians, one per fixation, on the
+DINOv3 patch grid.  The result is min-max normalized and used as the soft
+foreground prior W_fg; (1-H) gives the background complement W_bg.
+
+───────────────────────────────────────────────────────────────────────────
+Simple CSV format (what predict_single.py / the README document):
+
+    x,y,duration
+    340,221,180
+    356,228,145
+    ...
+
+    x, y     — fixation position in *raw pixel* coordinates of the original image.
+               load_fixation_csv() divides these by the image width/height to
+               produce the normalized [0,1] coords the model needs.
+    duration — fixation duration in any consistent unit (ms typical).
+
+───────────────────────────────────────────────────────────────────────────
+EyeLink/Tobii batch format (used by the full-dataset loaders in datasets.py):
+
+    IMAGE                  image filename the fixation belongs to
+    CURRENT_FIX_INDEX      fixation order within the trial
+    CURRENT_FIX_X          fixation x, already normalized to [0, 1]
+    CURRENT_FIX_Y          fixation y, already normalized to [0, 1]
+    CURRENT_FIX_DURATION   fixation duration (ms or any consistent unit)
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+import torch
+from PIL import Image
+
+# ── EyeLink-style column names (used by datasets.py / batch evaluation) ──
+TIME_COL    = "CURRENT_FIX_DURATION"
+FIX_IDX_COL = "CURRENT_FIX_INDEX"
+FIX_X_COL   = "CURRENT_FIX_X"
+FIX_Y_COL   = "CURRENT_FIX_Y"
+IMG_COL     = "IMAGE"
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+#  Simple CSV loader  (README-documented format)
+# ═══════════════════════════════════════════════════════════════════════════
+
+def load_fixation_csv(
+    csv_path: str,
+    image_width: int = None,
+    image_height: int = None,
+    image_name: str = None,
+) -> torch.Tensor:
+
+    df = pd.read_csv(csv_path)
+    df.columns = df.columns.str.strip()
+
+    # --------------------------------------------------
+    # Match DICOM image IDs to fixation CSV image names
+    # Example:
+    #   Prostate3T-01-0002_16.dcm
+    # -> Prostate3T-01-0002_16.jpg
+    # --------------------------------------------------
+    if "IMAGE" in df.columns and image_name is not None:
+
+        target_name = f"{image_name}.jpg"
+
+        df = df[df["IMAGE"] == target_name]
+
+    if len(df) == 0:
+        raise ValueError(
+            f"No fixations found for image: {image_name}"
+        )
+
+    # EyeLink/Tobii format
+    if (
+        "CURRENT_FIX_X" in df.columns
+        and "CURRENT_FIX_Y" in df.columns
+        and "CURRENT_FIX_DURATION" in df.columns
+    ):
+        x = df["CURRENT_FIX_X"].to_numpy(np.float32)
+        y = df["CURRENT_FIX_Y"].to_numpy(np.float32)
+        t = df["CURRENT_FIX_DURATION"].to_numpy(np.float32)
+
+    # README format
+    elif "x" in df.columns and "y" in df.columns:
+        x = df["x"].to_numpy(np.float32) / image_width
+        y = df["y"].to_numpy(np.float32) / image_height
+        t = df["duration"].to_numpy(np.float32)
+
+    else:
+        raise ValueError(f"Unsupported CSV format: {df.columns}")
+
+    # duration normalization
+    t_min, t_max = t.min(), t.max()
+    t = (t - t_min) / (t_max - t_min + 1e-8)
+
+    seq = np.stack([x, y, t], axis=1).astype(np.float32)
+
+    return torch.tensor(
+        seq,
+        dtype=torch.float32
+    ).unsqueeze(0)
+
+# ═══════════════════════════════════════════════════════════════════════════
+#  Batch-evaluation helpers  (used by datasets.py)
+# ═══════════════════════════════════════════════════════════════════════════
+
+def get_scanpath(df: pd.DataFrame, max_len: int, time_col: str = TIME_COL) -> torch.Tensor:
+    """Build a zero-padded ``(max_len, 3)`` scanpath tensor from a
+    fixation-report DataFrame already filtered to a single image.
+
+    Columns: ``[CURRENT_FIX_X, CURRENT_FIX_Y, CURRENT_FIX_DURATION]``
+    already normalized to [0,1] as produced by the EyeLink/Tobii export
+    pipeline (datasets.py normalizes these before calling this function).
+    """
+    if len(df) == 0:
+        return torch.zeros(max_len, 3)
+
+    df = df.sort_values(FIX_IDX_COL)
+    x = df[FIX_X_COL].to_numpy(np.float32)
+    y = df[FIX_Y_COL].to_numpy(np.float32)
+    t = df[time_col].to_numpy(np.float32)
+
+    t_min, t_max = t.min(), t.max()
+    t = (t - t_min) / (t_max - t_min + 1e-8)
+
+    seq = np.stack([x, y, t], axis=1)
+    if len(seq) > max_len:
+        seq = seq[:max_len]
+    if len(seq) < max_len:
+        seq = np.concatenate([seq, np.zeros((max_len - len(seq), 3), np.float32)], axis=0)
+
+    return torch.tensor(seq, dtype=torch.float32)
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+#  Heatmap generation
+# ═══════════════════════════════════════════════════════════════════════════
+
+def generate_gaze_heatmap(
+    fixation: torch.Tensor, h_patch: int, sigma: float = 2.0
+) -> torch.Tensor:
+    """Convert a batch of scanpaths into duration-weighted Gaussian heatmaps.
+
+    Parameters
+    ----------
+    fixation : (B, L, 3) tensor  [x_norm, y_norm, t_norm], zero-padded.
+    h_patch  : patch-grid side length (74 for IMG_SIZE=1184 with patch_size=16).
+    sigma    : Gaussian spread in patch units. Larger → softer, more diffuse
+               foreground prior; smaller → tighter, closer to exactly the
+               fixated patches.
+
+    Returns
+    -------
+    (B, h_patch, h_patch) float32 tensor, each sample min-max normalized
+    to [0, 1] independently.
+    """
+    B, L, _ = fixation.shape
+    device = fixation.device
+
+    y_coords = torch.arange(h_patch, device=device, dtype=torch.float32)
+    x_coords = torch.arange(h_patch, device=device, dtype=torch.float32)
+    grid_y, grid_x = torch.meshgrid(y_coords, x_coords, indexing="ij")
+
+    heatmaps = []
+    for b in range(B):
+        valid = fixation[b].abs().sum(dim=-1) > 1e-8   # filter zero-padding rows
+        valid_fix = fixation[b][valid]
+
+        h_b = torch.zeros(h_patch, h_patch, device=device)
+        if len(valid_fix) > 0:
+            for k in range(valid_fix.shape[0]):
+                x_k, y_k, t_k = valid_fix[k, 0], valid_fix[k, 1], valid_fix[k, 2]
+                gx = x_k * (h_patch - 1)
+                gy = y_k * (h_patch - 1)
+                dist_sq = (grid_x - gx) ** 2 + (grid_y - gy) ** 2
+                h_b += t_k * torch.exp(-dist_sq / (2 * sigma ** 2))
+
+            h_max = h_b.max()
+            if h_max > 1e-8:
+                h_b = h_b / h_max
+        heatmaps.append(h_b)
+
+    return torch.stack(heatmaps, dim=0)

gazecorrect/metrics.py

@@ -0,0 +1,27 @@
+"""metrics.py — Dice / IoU against a binarized prediction."""
+
+from __future__ import annotations
+
+import torch
+
+
+@torch.no_grad()
+def compute_metrics(preds: torch.Tensor, mask: torch.Tensor, thr: float = 0.5) -> dict:
+    """
+    preds, mask : (B, 1, H, W). ``preds`` is the continuous [0, 1] GazeRefine
+    output, ``mask`` is the ground-truth binary mask.
+    """
+    pred_bin = (preds > thr).float()
+    tp = (pred_bin * mask).sum((1, 2, 3))
+    fp = (pred_bin * (1 - mask)).sum((1, 2, 3))
+    fn = ((1 - pred_bin) * mask).sum((1, 2, 3))
+
+    dice_per = 2 * tp / (2 * tp + fp + fn + 1e-8)
+    iou_per = tp / (tp + fp + fn + 1e-8)
+
+    return {
+        "dice": dice_per.mean().item(),
+        "iou": iou_per.mean().item(),
+        "dice_per": dice_per,
+        "iou_per": iou_per,
+    }

gazecorrect/model.py

@@ -0,0 +1,242 @@
+"""
+model.py — The GazeRefine zero-shot segmentation model.
+
+This is a 1:1 refactor of the two task-specific scripts (Kvasir-SEG polyp /
+NCI-ISBI prostate MRI) into a single, dataset-agnostic module. The numerics
+are unchanged from the original experiments — only structure, naming and
+comments were cleaned up. Every per-dataset difference (sigma, kNN k,
+gaze-anchor weight, max_iters, ...) is now a constructor / config argument
+instead of a hardcoded default, see ``configs/*.yaml``.
+
+Maps onto the paper (Section 2) as follows:
+
+    gaze.generate_gaze_heatmap   -> H(u, v)                          (Eq. gaze prior)
+    W_fg, W_bg                   -> W_fg^(0), W_bg^(0)                (gaze-derived weights)
+    F_proto_init, B_proto_init   -> F^(0), B^(0)                      (initial prototypes)
+    contrast_method="difference" -> S_i^(t) = max(0, s_fg - alpha*s_bg)  (contrastive cleaning)
+    knn_affinity_refinement      -> S̄_i^(t)                           (kNN affinity propagation)
+    gaze_anchor_weight           -> lambda                            (anchor blending strength)
+    convergence check            -> ||F^(t+1) - F^(t)|| < eps
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .backbone import FrozenDINOv3
+from .gaze import generate_gaze_heatmap
+from .constants import IMG_SIZE, PATCH_SIZE
+
+EPS = 1e-12
+
+
+def knn_affinity_refinement(
+    Pv: torch.Tensor, S: torch.Tensor, k: int = 5, temperature: float = 0.05
+) -> torch.Tensor:
+    """Propagate a per-patch score map ``S`` across its k nearest neighbors in
+    frozen DINOv3 feature space (patch-to-patch affinity), encouraging
+    coherent, object-level responses instead of isolated high-confidence
+    patches. Corresponds to S̄ in the paper.
+
+    Pv : (B, N, D) raw (unnormalized) patch embeddings.
+    S  : (B, N)    current per-patch score to be smoothed.
+    """
+    B, N, D = Pv.shape
+    Pv_norm = F.normalize(Pv, p=2, dim=-1)
+
+    sim_matrix = torch.bmm(Pv_norm, Pv_norm.transpose(1, 2))            # (B, N, N)
+    topk_vals, topk_indices = torch.topk(sim_matrix, k=k, dim=-1)        # (B, N, k)
+    weights = F.softmax(topk_vals / temperature, dim=-1)                 # (B, N, k)
+
+    S_expanded = S.unsqueeze(1).expand(B, N, N)
+    S_neighbors = torch.gather(S_expanded, dim=2, index=topk_indices)    # (B, N, k)
+
+    return torch.sum(weights * S_neighbors, dim=-1)                     # (B, N)
+
+
+class GazeRefine(nn.Module):
+    """Zero-shot, training-free, gaze-guided segmentation model.
+
+    The model has exactly one set of *learned* weights: the frozen,
+    pretrained DINOv3 backbone. Everything else — prototype construction,
+    contrastive cleaning, kNN propagation, gaze anchoring — is a closed-form
+    operation re-run from scratch on every image at inference time.
+
+    Parameters
+    ----------
+    dino_name : timm DINOv3 checkpoint name.
+    img_size, patch_size : input resolution / ViT patch size. Must satisfy
+        ``img_size % patch_size == 0``.
+    sigma : Gaussian spread (in patch units) for the gaze heatmap.
+    extract_mode : ``"last"`` (1 block) or ``"all"`` (4 blocks, averaged).
+    contrast_method : ``"difference"`` (paper default, contrastive cleaning),
+        ``"softmax"`` (foreground/background softmax ratio), or
+        ``"original"`` (plain foreground cosine similarity, no background
+        suppression — kept for the ablation in Table 2).
+    temperature : softmax temperature, only used when contrast_method="softmax".
+    max_iters : maximum recurrent refinement iterations (T in the paper).
+    knn_refine : whether to apply kNN affinity propagation each iteration.
+    knn_k, knn_temp : kNN neighborhood size / softmax temperature.
+    gaze_anchor_weight : lambda — how strongly each iteration's prototypes
+        are pulled back toward the original gaze-only prototypes. Higher =
+        trust the raw fixations more; lower = let the model drift further
+        from the initial gaze region.
+    """
+
+    def __init__(
+        self,
+        dino_name: str = "vit_large_patch16_dinov3.lvd1689m",
+        img_size: int = IMG_SIZE,
+        patch_size: int = PATCH_SIZE,
+        sigma: float = 2.0,
+        extract_mode: str = "last",
+        contrast_method: str = "difference",
+        temperature: float = 0.05,
+        max_iters: int = 10,
+        knn_refine: bool = True,
+        knn_k: int = 5,
+        knn_temp: float = 0.1,
+        gaze_anchor_weight: float = 0.6,
+    ):
+        super().__init__()
+        assert img_size % patch_size == 0, "img_size must be a multiple of patch_size"
+        self.img_size = img_size
+        self.h_patch = img_size // patch_size
+        self.sigma = sigma
+        self.contrast_method = contrast_method
+        self.temperature = temperature
+        self.max_iters = max_iters
+        self.knn_refine = knn_refine
+        self.knn_k = knn_k
+        self.knn_temp = knn_temp
+        self.gaze_anchor_weight = gaze_anchor_weight
+
+        self.visual_enc = FrozenDINOv3(dino_name, extract_mode=extract_mode)
+
+    # ------------------------------------------------------------------ #
+    def _gaze_weights(self, fixation: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Scanpath -> heatmap -> normalized foreground/background spatial weights."""
+        gaze_heatmap = generate_gaze_heatmap(fixation, h_patch=self.h_patch, sigma=self.sigma)
+        H_flat = gaze_heatmap.view(gaze_heatmap.size(0), -1)
+
+        W_fg = H_flat / (H_flat.sum(dim=-1, keepdim=True) + EPS)
+        bg_w = 1.0 - H_flat
+        W_bg = bg_w / (bg_w.sum(dim=-1, keepdim=True) + EPS)
+        return gaze_heatmap, W_fg, W_bg
+
+    def _score(self, Pv_norm: torch.Tensor, F_proto_norm: torch.Tensor, B_proto_norm: torch.Tensor) -> torch.Tensor:
+        """Foreground/background contrastive scoring for one level, one iteration."""
+        if self.contrast_method == "original":
+            S = torch.bmm(Pv_norm, F_proto_norm.unsqueeze(-1)).squeeze(-1)
+            return torch.clamp(S, min=0.0)
+
+        sim_fg = torch.bmm(Pv_norm, F_proto_norm.unsqueeze(-1)).squeeze(-1)
+        sim_bg = torch.bmm(Pv_norm, B_proto_norm.unsqueeze(-1)).squeeze(-1)
+
+        if self.contrast_method == "difference":
+            return torch.clamp(sim_fg - sim_bg, min=0.0)
+        elif self.contrast_method == "softmax":
+            stacked = torch.stack([sim_fg, sim_bg], dim=-1) / self.temperature
+            probs = F.softmax(stacked, dim=-1)
+            return probs[:, :, 0]
+        raise ValueError(f"Unknown contrast_method: {self.contrast_method!r}")
+
+    def _refine_level(self, Pv: torch.Tensor, W_fg: torch.Tensor, W_bg: torch.Tensor) -> torch.Tensor:
+        """Run the full recurrent gaze-anchored refinement loop for one DINOv3 level.
+        Returns the final (B, N) per-patch foreground score map for that level."""
+        B = Pv.size(0)
+        device = Pv.device
+        Pv_norm = F.normalize(Pv, p=2, dim=-1)
+
+        # Initial gaze-only prototypes — F^(0), B^(0)
+        F_proto_init = torch.sum(W_fg.unsqueeze(-1) * Pv, dim=1)
+        B_proto_init = torch.sum(W_bg.unsqueeze(-1) * Pv, dim=1)
+
+        W_fg_curr, W_bg_curr = W_fg.clone(), W_bg.clone()
+        best_S = W_fg.clone()
+        S_prev = W_fg.clone()
+        active = torch.ones(B, dtype=torch.bool, device=device)
+
+        for _ in range(self.max_iters):
+            W_fg_norm = W_fg_curr / (W_fg_curr.sum(dim=-1, keepdim=True) + EPS)
+            W_bg_norm = W_bg_curr / (W_bg_curr.sum(dim=-1, keepdim=True) + EPS)
+
+            F_proto_curr = torch.sum(W_fg_norm.unsqueeze(-1) * Pv, dim=1)
+            B_proto_curr = torch.sum(W_bg_norm.unsqueeze(-1) * Pv, dim=1)
+
+            # Anchor toward the initial gaze-only prototypes — lambda blending
+            lam = self.gaze_anchor_weight
+            F_proto = (1 - lam) * F_proto_curr + lam * F_proto_init
+            B_proto = (1 - lam) * B_proto_curr + lam * B_proto_init
+            F_proto_norm = F.normalize(F_proto, p=2, dim=-1)
+            B_proto_norm = F.normalize(B_proto, p=2, dim=-1)
+
+            S_iter = self._score(Pv_norm, F_proto_norm, B_proto_norm)
+            if self.knn_refine:
+                S_iter = knn_affinity_refinement(Pv, S_iter, k=self.knn_k, temperature=self.knn_temp)
+
+            S_new = S_prev.clone()
+            S_new[active] = S_iter[active]
+
+            # Collapse prevention: if an active sample's map flattened to ~0,
+            # roll it back and freeze it instead of letting it degenerate.
+            collapsed = S_new.max(dim=-1)[0] < 1e-4
+            freeze = active & collapsed
+            S_new[freeze] = S_prev[freeze]
+            active = active & ~collapsed
+
+            best_S = S_new.clone()
+            if not active.any():
+                break
+
+            # Turn the refined score into the next iteration's spatial weights
+            S_sig = torch.sigmoid(S_new)
+            W_fg_next = S_sig / (S_sig.sum(dim=-1, keepdim=True) + EPS)
+            bg_w = 1.0 - S_sig
+            bg_w = torch.clamp(bg_w - bg_w.min(dim=-1, keepdim=True).values, min=0.0)
+            W_bg_next = bg_w / (bg_w.sum(dim=-1, keepdim=True) + EPS)
+
+            F_proto_next = torch.sum(W_fg_next.unsqueeze(-1) * Pv, dim=1)
+            F_proto_next_norm = F.normalize(F_proto_next, p=2, dim=-1)
+
+            # Convergence: foreground prototype stopped moving
+            diff = torch.abs(F_proto_next_norm - F_proto_norm).mean(dim=-1)
+            if torch.all(diff < 1e-6):
+                break
+
+            S_prev = S_new.clone()
+            W_fg_curr[active] = W_fg_next[active]
+            W_bg_curr[active] = W_bg_next[active]
+
+        return best_S
+
+    # ------------------------------------------------------------------ #
+    def forward(self, image: torch.Tensor, fixation: torch.Tensor) -> dict:
+        """
+        image : (B, 3, H, W) normalized RGB tensor (ImageNet mean/std).
+        fixation : (B, L, 3) zero-padded scanpath, see ``gazerefine.gaze.get_scanpath``.
+
+        Returns a dict with:
+            preds        (B, 1, H, W) final mask in [0, 1] — threshold at 0.5 for a binary mask
+            gaze_heatmap (B, h, w) the raw gaze prior, useful for visualization
+            final_map    (B, N) the un-upsampled patch-level score map
+        """
+        B = image.size(0)
+        Pv_list = self.visual_enc(image)
+        gaze_heatmap, W_fg, W_bg = self._gaze_weights(fixation)
+
+        all_maps = [self._refine_level(Pv, W_fg, W_bg) for Pv in Pv_list]
+        final_map = torch.stack(all_maps, dim=1).mean(dim=1)  # average across levels
+
+        grid = final_map.view(B, 1, self.h_patch, self.h_patch)
+        upsampled = F.interpolate(grid, size=(self.img_size, self.img_size), mode="bilinear", align_corners=False)
+
+        # per-sample min-max normalization -> ready-to-threshold mask
+        flat = upsampled.view(B, -1)
+        m_min = flat.min(dim=-1, keepdim=True).values.view(B, 1, 1, 1)
+        m_max = flat.max(dim=-1, keepdim=True).values.view(B, 1, 1, 1)
+        preds = (upsampled - m_min) / (m_max - m_min + 1e-8)
+
+        return dict(preds=preds, gaze_heatmap=gaze_heatmap, final_map=final_map)

gazecorrect/visualize.py

@@ -0,0 +1,86 @@
+"""visualize.py — Plotting / overlay helpers.
+
+``save_prediction`` reproduces the 4-panel (image / gaze / GT / prediction)
+grid used during batch evaluation. ``overlay_heatmap`` and ``overlay_mask``
+are lighter-weight, matplotlib-free PIL helpers meant for the Gradio demo
+and the example notebook, where you typically want a single composited
+image rather than a saved subplot file.
+"""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+import torch
+from PIL import Image
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import matplotlib.cm as cm
+
+from .constants import IMG_MEAN, IMG_STD
+
+
+def _denormalize(img: torch.Tensor) -> np.ndarray:
+    """(3, H, W) normalized tensor -> (H, W, 3) uint8-range float array in [0, 1]."""
+    mean = torch.tensor(IMG_MEAN, device=img.device).view(3, 1, 1)
+    std = torch.tensor(IMG_STD, device=img.device).view(3, 1, 1)
+    return (img * std + mean).clamp(0, 1).permute(1, 2, 0).cpu().numpy()
+
+
+@torch.no_grad()
+def save_prediction(
+    img: torch.Tensor,
+    gaze_heatmap: torch.Tensor,
+    gt: torch.Tensor,
+    pred: torch.Tensor,
+    name: str,
+    out_dir: str = "predictions",
+) -> None:
+    """Save a 4-panel [image | gaze heatmap | GT mask | predicted mask] figure."""
+    os.makedirs(out_dir, exist_ok=True)
+
+    img_np = _denormalize(img)
+    gaze_np = gaze_heatmap.squeeze().cpu().numpy()
+    gt_np = gt.squeeze().cpu().numpy()
+    pred_np = pred.squeeze().cpu().numpy()
+
+    fig, ax = plt.subplots(1, 4, figsize=(20, 5))
+    ax[0].imshow(img_np); ax[0].set_title("Image")
+    ax[1].imshow(gaze_np, cmap="jet"); ax[1].set_title("Gaze Heatmap")
+    ax[2].imshow(gt_np, cmap="gray"); ax[2].set_title("GT mask")
+    ax[3].imshow(pred_np, cmap="gray"); ax[3].set_title("Predicted mask")
+    for a in ax:
+        a.axis("off")
+    plt.tight_layout()
+    plt.savefig(os.path.join(out_dir, name), dpi=100, bbox_inches="tight")
+    plt.close()
+
+
+def overlay_heatmap(image: Image.Image, heatmap: np.ndarray, alpha: float = 0.45) -> Image.Image:
+    """Composite a [0, 1] heatmap (any H'xW', will be resized) onto a PIL image
+    using a jet colormap. Used to show the gaze prior over the original image."""
+    heatmap = np.asarray(heatmap, dtype=np.float32)
+    heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min() + 1e-8)
+
+    heat_img = Image.fromarray((cm.jet(heatmap)[:, :, :3] * 255).astype(np.uint8))
+    heat_img = heat_img.resize(image.size, resample=Image.BILINEAR)
+
+    base = image.convert("RGB")
+    return Image.blend(base, heat_img, alpha=alpha)
+
+
+def overlay_mask(
+    image: Image.Image, mask: np.ndarray, color: tuple[int, int, int] = (255, 60, 60), alpha: float = 0.45
+) -> Image.Image:
+    """Composite a binary/soft [0, 1] mask onto a PIL image as a solid color wash."""
+    mask = np.asarray(mask, dtype=np.float32)
+    mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)
+    mask_img = Image.fromarray((mask * 255).astype(np.uint8)).resize(image.size, resample=Image.NEAREST)
+
+    base = image.convert("RGB")
+    color_layer = Image.new("RGB", base.size, color)
+    composited = Image.composite(color_layer, base, mask_img.point(lambda p: int(p * alpha)))
+    return composited

notebooks/GazeRefine_Demo.ipynb

@@ -0,0 +1,306 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# GazeRefine \u2014 Interactive Walkthrough\n",
+    "\n",
+    "This notebook walks through the **full GazeRefine pipeline** step by step:\n",
+    "\n",
+    "1. Build a gaze heatmap from a scanpath\n",
+    "2. Construct gaze-anchored foreground/background prototypes in frozen DINOv3 feature space\n",
+    "3. Run the recurrent refinement loop (contrastive cleaning + kNN affinity propagation)\n",
+    "4. Visualize the final mask\n",
+    "\n",
+    "It ships with a **synthetic toy example** (a blob on a noisy background, with a few clicked\n",
+    "fixation points) so you can see every intermediate step without needing any dataset.\n",
+    "Swap in a real image + fixation CSV from `examples/` (see `examples/README.md`) to run it\n",
+    "on actual data \u2014 the API is identical.\n",
+    "\n",
+    "> Section 4 needs `torch` + `timm` (and a DINOv3 checkpoint download, so internet access)\n",
+    "> to call the real backbone. Sections 1\u20133 only need `numpy`/`matplotlib` and run anywhere,\n",
+    "> using a tiny **mock feature extractor** so you can still see the prototype/refinement math."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys, os\n",
+    "from pathlib import Path\n",
+    "\n",
+    "REPO_ROOT = Path.cwd().parent if Path.cwd().name == \"notebooks\" else Path.cwd()\n",
+    "sys.path.insert(0, str(REPO_ROOT))\n",
+    "\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "np.random.seed(0)\n",
+    "%matplotlib inline"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. A synthetic example\n",
+    "\n",
+    "We draw a soft circular \"lesion\" on a noisy background and pick a few fixation points near\n",
+    "its center \u2014 mimicking a clinician glancing at the structure of interest a handful of times,\n",
+    "with the first, longest fixation landing closest to the center."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "IMG_SIZE = 224  # small, just for a fast visual demo (the real pipeline defaults to 1184 / DINOv3 patch grid)\n",
+    "\n",
+    "def make_toy_image(size=IMG_SIZE, cx=0.55, cy=0.45, r=0.12):\n",
+    "    yy, xx = np.mgrid[0:size, 0:size] / size\n",
+    "    dist = np.sqrt((xx - cx) ** 2 + (yy - cy) ** 2)\n",
+    "    blob = np.clip(1 - dist / r, 0, 1) ** 2\n",
+    "    background = 0.25 + 0.05 * np.random.randn(size, size)\n",
+    "    img = np.clip(background + 0.6 * blob, 0, 1)\n",
+    "    gt_mask = (dist < r).astype(np.float32)\n",
+    "    return img, gt_mask\n",
+    "\n",
+    "toy_img, toy_mask = make_toy_image()\n",
+    "\n",
+    "# a few fixations clustered near the blob, decaying duration, with a touch of natural jitter\n\n",
+    "fixations = [\n",
+    "    (0.55, 0.45, 1.0),\n",
+    "    (0.57, 0.47, 0.7),\n",
+    "    (0.53, 0.44, 0.5),\n",
+    "    (0.50, 0.50, 0.3),\n",
+    "]\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 2, figsize=(8, 4))\n",
+    "ax[0].imshow(toy_img, cmap=\"gray\"); ax[0].set_title(\"Synthetic image\")\n",
+    "for x, y, d in fixations:\n",
+    "    ax[0].scatter(x * IMG_SIZE, y * IMG_SIZE, s=200 * d, c=\"red\", alpha=0.6, edgecolors=\"white\")\n",
+    "ax[1].imshow(toy_mask, cmap=\"gray\"); ax[1].set_title(\"Ground-truth mask (synthetic)\")\n",
+    "for a in ax: a.axis(\"off\")\n",
+    "plt.tight_layout(); plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 2. Scanpath \u2192 gaze heatmap\n",
+    "\n",
+    "`gazerefine.gaze.generate_gaze_heatmap` places a duration-weighted Gaussian at each fixation\n",
+    "and min-max normalizes the result. This becomes the soft foreground prior $W_{fg}^{(0)}$;\n",
+    "$1 - H$ becomes the background prior $W_{bg}^{(0)}$."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from gazerefine.gaze import generate_gaze_heatmap\n",
+    "\n",
+    "H_PATCH = 28  # toy patch grid (real pipeline uses 74x74, see gazerefine/constants.py)\n",
+    "fixation_t = torch.tensor([fixations], dtype=torch.float32)  # (1, L, 3)\n",
+    "\n",
+    "gaze_heatmap = generate_gaze_heatmap(fixation_t, h_patch=H_PATCH, sigma=2.0)[0].numpy()\n",
+    "\n",
+    "plt.figure(figsize=(4, 4))\n",
+    "plt.imshow(gaze_heatmap, cmap=\"jet\")\n",
+    "plt.title(\"Duration-weighted gaze heatmap (patch grid)\")\n",
+    "plt.axis(\"off\"); plt.colorbar(fraction=0.046); plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 3. Prototype construction & refinement \u2014 mechanics demo\n",
+    "\n",
+    "To see the **math** of gaze-anchored prototypes, contrastive cleaning, and kNN affinity\n",
+    "propagation without downloading a backbone, we substitute a tiny mock \"feature extractor\":\n",
+    "patches near the blob get embeddings clustered around one random direction, background\n",
+    "patches cluster around another, with noise \u2014 a cartoon stand-in for \"DINOv3 already groups\n",
+    "semantically similar patches together\", which is the real property GazeRefine exploits.\n",
+    "\n",
+    "**This section is for intuition only.** Section 4 below runs the real frozen DINOv3 backbone."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from gazerefine.model import knn_affinity_refinement\n\n",
+    "\n",
+    "D = 64  # toy embedding dim\n\n",
+    "rng = np.random.default_rng(0)\n\n",
+    "fg_direction = rng.normal(size=D); fg_direction /= np.linalg.norm(fg_direction)\n\n",
+    "bg_direction = rng.normal(size=D); bg_direction /= np.linalg.norm(bg_direction)\n\n",
+    "\n\n",
+    "# resize the toy GT mask down to the patch grid to decide which mock cluster each patch belongs to\n\n",
+    "import torch.nn.functional as Fnn\n\n",
+    "mask_small = Fnn.interpolate(torch.tensor(toy_mask)[None, None], size=(H_PATCH, H_PATCH), mode=\"nearest\")[0, 0].numpy()\n\n",
+    "patch_is_fg = mask_small.flatten() > 0.5\n\n",
+    "\n\n",
+    "N = H_PATCH * H_PATCH\n\n",
+    "Pv = np.stack([\n\n",
+    "    (fg_direction if patch_is_fg[i] else bg_direction) + 0.3 * rng.normal(size=D)\n\n",
+    "    for i in range(N)\n\n",
+    "])\n\n",
+    "Pv = torch.tensor(Pv, dtype=torch.float32)[None]  # (1, N, D)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "H_flat = torch.tensor(gaze_heatmap.flatten(), dtype=torch.float32)[None]  # (1, N)\n\n",
+    "eps = 1e-12\n\n",
+    "W_fg = H_flat / (H_flat.sum(-1, keepdim=True) + eps)\n\n",
+    "W_bg = (1 - H_flat) / (1 - H_flat).sum(-1, keepdim=True)\n\n",
+    "\n\n",
+    "Pv_norm = Fnn.normalize(Pv, p=2, dim=-1)\n\n",
+    "F_proto = Fnn.normalize((W_fg.unsqueeze(-1) * Pv).sum(1), p=2, dim=-1)\n\n",
+    "B_proto = Fnn.normalize((W_bg.unsqueeze(-1) * Pv).sum(1), p=2, dim=-1)\n\n",
+    "\n\n",
+    "sim_fg = torch.bmm(Pv_norm, F_proto.unsqueeze(-1)).squeeze(-1)\n\n",
+    "sim_bg = torch.bmm(Pv_norm, B_proto.unsqueeze(-1)).squeeze(-1)\n\n",
+    "S_clean = torch.clamp(sim_fg - sim_bg, min=0.0)              # contrastive cleaning\n\n",
+    "S_knn = knn_affinity_refinement(Pv, S_clean, k=8, temperature=0.1)  # kNN affinity propagation\n\n",
+    "\n\n",
+    "def to_grid(s):\n\n",
+    "    g = s[0].reshape(H_PATCH, H_PATCH).detach().numpy()\n\n",
+    "    return (g - g.min()) / (g.max() - g.min() + 1e-8)\n\n",
+    "\n\n",
+    "fig, ax = plt.subplots(1, 4, figsize=(16, 4))\n\n",
+    "ax[0].imshow(gaze_heatmap, cmap=\"jet\"); ax[0].set_title(\"Gaze prior H\")\n\n",
+    "ax[1].imshow(to_grid(sim_fg - sim_bg), cmap=\"viridis\"); ax[1].set_title(\"sim_fg \u2212 sim_bg (pre-clean)\")\n\n",
+    "ax[2].imshow(to_grid(S_clean), cmap=\"viridis\"); ax[2].set_title(\"After contrastive cleaning\")\n\n",
+    "ax[3].imshow(to_grid(S_knn), cmap=\"viridis\"); ax[3].set_title(\"After kNN propagation\")\n\n",
+    "for a in ax: a.axis(\"off\")\n\n",
+    "plt.tight_layout(); plt.show()\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Notice how the gaze prior alone is a soft, blurry blob centered on the clicked points,\n\n",
+    "but after contrastive cleaning + kNN propagation the response sharpens to the full extent\n\n",
+    "of the (mock-)semantically coherent foreground region \u2014 this is the localization gain\n\n",
+    "described in the paper, beyond the initially fixated patches."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. The real pipeline \u2014 frozen DINOv3 + `GazeRefine`\n",
+    "\n",
+    "This cell calls the actual model end to end. It needs `torch`, `timm`, and internet access\n",
+    "the first time (to download the DINOv3 checkpoint). If those aren't available in your current\n",
+    "environment, this cell will print instructions instead of failing silently \u2014 run it in a\n",
+    "GPU-enabled environment with the repo's `requirements.txt` installed."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "try:\n\n",
+    "    from gazerefine import GazeRefine\n\n",
+    "    from gazerefine.constants import IMG_SIZE\n\n",
+    "    import torchvision.transforms as T\n\n",
+    "\n\n",
+    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n",
+    "\n\n",
+    "    model = GazeRefine(\n\n",
+    "        dino_name=\"vit_large_patch16_dinov3.lvd1689m\",\n\n",
+    "        sigma=2.0, contrast_method=\"difference\", max_iters=5,\n\n",
+    "        gaze_anchor_weight=0.5, knn_k=20,\n\n",
+    "    ).to(device).eval()\n\n",
+    "\n\n",
+    "    img_3ch = np.stack([toy_img] * 3, axis=-1)\n\n",
+    "    tf = T.Compose([T.ToPILImage(), T.Resize((IMG_SIZE, IMG_SIZE)), T.ToTensor(),\n\n",
+    "                     T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])\n\n",
+    "    img_t = tf((img_3ch * 255).astype(np.uint8)).unsqueeze(0).to(device)\n\n",
+    "\n\n",
+    "    fix_t = torch.tensor([fixations], dtype=torch.float32).to(device)\n\n",
+    "\n\n",
+    "    with torch.no_grad():\n\n",
+    "        out = model(img_t, fix_t)\n\n",
+    "\n\n",
+    "    pred = out[\"preds\"][0, 0].cpu().numpy()\n\n",
+    "    fig, ax = plt.subplots(1, 3, figsize=(12, 4))\n\n",
+    "    ax[0].imshow(toy_img, cmap=\"gray\"); ax[0].set_title(\"Input\")\n\n",
+    "    ax[1].imshow(out[\"gaze_heatmap\"][0].cpu().numpy(), cmap=\"jet\"); ax[1].set_title(\"Gaze prior\")\n\n",
+    "    ax[2].imshow(pred > 0.5, cmap=\"gray\"); ax[2].set_title(\"GazeRefine mask\")\n\n",
+    "    for a in ax: a.axis(\"off\")\n\n",
+    "    plt.tight_layout(); plt.show()\n\n",
+    "\n\n",
+    "except ModuleNotFoundError as e:\n\n",
+    "    print(f\"[skipped] missing dependency: {e}\")\n\n",
+    "    print(\"Install requirements.txt (torch, timm, ...) and re-run this cell to see real DINOv3 output.\")\n\n",
+    "except Exception as e:\n\n",
+    "    print(f\"[skipped] could not run the real backbone in this environment: {e}\")\n\n",
+    "    print(\"This is most likely a missing internet connection for the first-time checkpoint download.\")\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 5. Try it on a real image + fixation CSV\n",
+    "\n",
+    "Once you've added real assets under `examples/` (see `examples/README.md`), point\n",
+    "`scripts.predict_single.predict` at them directly:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# from scripts.predict_single import predict\n\n",
+    "#\n\n",
+    "# out = predict(\n\n",
+    "#     image=\"../examples/images/kvasir_sample_01.jpg\",\n\n",
+    "#     fixation_csv=\"../examples/fixations/kvasir_sample_01_fixations.csv\",\n\n",
+    "#     image_key=\"kvasir_sample_01.jpg\",\n\n",
+    "#     sigma=2.0, contrast_method=\"difference\", max_iters=5,\n\n",
+    "#     gaze_anchor_weight=0.5, knn_k=20,\n\n",
+    "# )\n\n",
+    "# out[\"mask_overlay\"]\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "pygments_lexer": "ipython3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

notebooks/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 GazeRefine Authors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

notebooks/configs/kvasir.yaml

@@ -0,0 +1,27 @@
+# GazeRefine config — Kvasir-SEG (colonoscopy polyp segmentation)
+# These are the exact hyperparameters used to obtain the 89.49% Dice
+# reported in Table 1 of the paper.
+
+dataset: kvasir
+root: /path/to/Kvasir-SEG          # expects <root>/images, <root>/masks
+fixation_csv: /path/to/kvasir_fixation.csv
+
+dino_name: vit_large_patch16_dinov3.lvd1689m
+extract_mode: last                 # "last" | "all"
+
+sigma: 2.0                         # gaze Gaussian spread, in patch units
+contrast_method: difference        # "difference" | "softmax" | "original"
+temperature: 0.05                  # only used by contrast_method=softmax
+max_iters: 5                       # T, recurrent refinement iterations
+gaze_anchor_weight: 0.5            # lambda, anchoring strength
+
+knn_refine: true
+knn_k: 20
+knn_temp: 0.1
+
+threshold: 0.5                     # binarization threshold
+
+batch_size: 4
+num_workers: 4
+save_num: 10
+pred_dir: ./predictions/kvasir

notebooks/configs/prostate_mri.yaml

@@ -0,0 +1,27 @@
+# GazeRefine config — NCI-ISBI Prostate MRI segmentation
+# These are the exact hyperparameters used to obtain the 75.44% Dice
+# reported in Table 1 of the paper.
+
+dataset: prostate_mri
+root: /path/to/Pros_MRI            # expects <root>/images (.dcm), <root>/masks (.png)
+fixation_csv: /path/to/nci-isbi_fixation.csv
+
+dino_name: vit_large_patch16_dinov3.lvd1689m
+extract_mode: last                 # "last" | "all"
+
+sigma: 1.5                         # gaze Gaussian spread, in patch units
+contrast_method: difference        # "difference" | "softmax" | "original"
+temperature: 0.05                  # only used by contrast_method=softmax
+max_iters: 1                       # T, recurrent refinement iterations
+gaze_anchor_weight: 0.8            # lambda, anchoring strength (trust gaze more — low-contrast MRI)
+
+knn_refine: true
+knn_k: 3
+knn_temp: 0.1
+
+threshold: 0.5                     # binarization threshold
+
+batch_size: 4
+num_workers: 4
+save_num: 10
+pred_dir: ./predictions/prostate_mri

notebooks/gazerefine/__init__.py

@@ -0,0 +1,33 @@
+"""
+GazeRefine
+==========
+Expert gaze as a test-time prompt for training-free medical image segmentation.
+
+This package exposes the full zero-shot pipeline described in the GazeRefine
+paper: a frozen DINOv3 backbone, gaze -> heatmap conversion, gaze-anchored
+foreground/background prototype construction, and recurrent refinement
+(contrastive background cleaning + kNN affinity propagation).
+
+Nothing in this package is trained. There are no learned weights other than
+the frozen, pretrained DINOv3 backbone loaded from `timm`.
+"""
+
+from .backbone import FrozenDINOv3
+from .gaze import get_scanpath, generate_gaze_heatmap
+from .model import GazeRefine, knn_affinity_refinement
+from .metrics import compute_metrics
+from .visualize import save_prediction, overlay_heatmap, overlay_mask
+
+__all__ = [
+    "FrozenDINOv3",
+    "get_scanpath",
+    "generate_gaze_heatmap",
+    "GazeRefine",
+    "knn_affinity_refinement",
+    "compute_metrics",
+    "save_prediction",
+    "overlay_heatmap",
+    "overlay_mask",
+]
+
+__version__ = "0.1.0"

notebooks/gazerefine/backbone.py

@@ -0,0 +1,78 @@
+"""
+backbone.py — Frozen DINOv3 feature extractor.
+
+GazeRefine never updates the backbone. We load a pretrained DINOv3 ViT via
+`timm`, freeze every parameter, and pull out raw patch tokens from one or
+more transformer blocks using forward hooks. No adapters, no projections,
+no fine-tuning.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+class FrozenDINOv3(nn.Module):
+    """Completely frozen DINOv3 ViT backbone that exposes raw patch tokens.
+
+    Parameters
+    ----------
+    model_name : str
+        Any DINOv3 variant available in `timm` (e.g.
+        ``"vit_large_patch16_dinov3.lvd1689m"``,
+        ``"vit_base_patch16_dinov3.lvd1689m"``). Larger backbones generally
+        give cleaner semantic separation but cost more memory/compute.
+    extract_mode : {"last", "all"}
+        - ``"last"``: use only the final block's patch tokens (fast, the
+          default used in our reported results).
+        - ``"all"``: pool patch tokens from 4 evenly-spaced blocks
+          (0, n/4, 3n/4, n-1) and average the resulting similarity maps in
+          ``GazeRefine``. This sometimes helps on harder modalities at the
+          cost of ~4x compute.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "vit_large_patch16_dinov3.lvd1689m",
+        extract_mode: str = "last",
+    ):
+        super().__init__()
+        import timm  # local import: keeps `timm` optional for users who only read the code
+
+        print(f"[GazeRefine] Loading frozen backbone: {model_name} (extract_mode={extract_mode})")
+        bb = timm.create_model(model_name, pretrained=True, num_classes=0)
+        for p in bb.parameters():
+            p.requires_grad_(False)
+        bb.eval()
+
+        self.backbone = bb
+        self.embed_dim = bb.embed_dim
+        self.num_blocks = len(bb.blocks)
+
+        if extract_mode == "last":
+            self.levels = [-1]
+        elif extract_mode == "all":
+            self.levels = [0, self.num_blocks // 4, (self.num_blocks * 3) // 4, self.num_blocks - 1]
+        else:
+            raise ValueError(f"extract_mode must be 'last' or 'all', got {extract_mode!r}")
+
+        print(f"[GazeRefine] Hooking transformer blocks at levels: {self.levels}")
+        self._feats: dict[int, torch.Tensor] = {}
+        for lvl in self.levels:
+            bb.blocks[lvl].register_forward_hook(self._make_hook(lvl))
+
+    def _make_hook(self, lvl: int):
+        def _hook_fn(module, inp, out):
+            # DINOv3 token layout: [CLS, register_1..register_k, patch_1..patch_N]
+            n_prefix = 1 + getattr(self.backbone, "num_register_tokens", 4)
+            self._feats[lvl] = out[:, n_prefix:, :]  # (B, N, D) — patch tokens only
+        return _hook_fn
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
+        """Run the frozen backbone and return a list of (B, N, D) patch-token tensors,
+        one per hooked level, in the order given by ``self.levels``."""
+        self._feats.clear()
+        self.backbone(x)
+        return [self._feats[lvl] for lvl in self.levels]

notebooks/gazerefine/constants.py

@@ -0,0 +1,13 @@
+"""constants.py — shared sizing and normalization constants.
+
+IMG_SIZE is fixed to an exact multiple of the ViT patch size so the patch
+grid divides evenly with no rounding/cropping artifacts.
+"""
+
+PATCH_SIZE = 16
+H_PATCH = 37 * 2                      # 74 patches per side
+IMG_SIZE = H_PATCH * PATCH_SIZE       # 1184 px  (use 518 = 37*14 if you switch to a /14 ViT)
+N_PATCHES = H_PATCH ** 2
+
+IMG_MEAN = [0.485, 0.456, 0.406]
+IMG_STD = [0.229, 0.224, 0.225]

notebooks/gazerefine/datasets.py

@@ -0,0 +1,154 @@
+"""
+datasets.py — Image/mask/fixation loaders for the two reported benchmarks,
+plus a small base class so adding a third dataset is just a few lines.
+
+Expected folder layout (same for both datasets):
+
+    <root>/
+        images/    one image per case (.jpg/.png for Kvasir, .dcm for prostate MRI)
+        masks/     matching binary mask, same basename, .png
+    <fixation_csv>
+        one row per fixation, see gazerefine.gaze for the expected columns.
+        The IMAGE column must match an images/ filename (Kvasir) or the
+        DICOM basename + ".jpg" (prostate — fixation collection was run on
+        JPEG-rendered slices while the model reads the original DICOM).
+"""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+import pandas as pd
+import torch
+from PIL import Image
+from torch.utils.data import Dataset
+import torchvision.transforms as T
+
+from .constants import IMG_SIZE, IMG_MEAN, IMG_STD
+from .gaze import get_scanpath, IMG_COL
+
+
+class _BaseGazeDataset(Dataset):
+    """Shared image/mask transform + fixation-grouping logic.
+
+    Subclasses only need to implement ``_load_image(name)`` and provide the
+    set of valid image ids that have a matching mask and fixation entries.
+    """
+
+    def __init__(self, root: str, fixation_csv: str, img_size: int = IMG_SIZE):
+        self.root = root
+        self.img_dir = os.path.join(root, "images")
+        self.mask_dir = os.path.join(root, "masks")
+        self.img_size = img_size
+
+        self.df = pd.read_csv(fixation_csv)
+        self.df.columns = self.df.columns.str.strip()
+        self.fix_df = self.df.groupby(IMG_COL)
+        self.max_len = int(self.df.groupby(IMG_COL).size().max())
+        print(f"[Dataset] max scanpath length = {self.max_len}")
+
+        self.image_ids: list[str] = []  # set by subclass __init__
+
+        self.img_tf = T.Compose([
+            T.Resize((img_size, img_size)),
+            T.ToTensor(),
+            T.Normalize(IMG_MEAN, IMG_STD),
+        ])
+        self.mask_tf = T.Compose([
+            T.Resize((img_size, img_size), interpolation=T.InterpolationMode.NEAREST),
+            T.ToTensor(),
+        ])
+
+    def __len__(self):
+        return len(self.image_ids)
+
+    def _load_image(self, name: str) -> Image.Image:
+        raise NotImplementedError
+
+    def _mask_path(self, name: str) -> str:
+        raise NotImplementedError
+
+    def _fixation_key(self, name: str) -> str:
+        """CSV IMAGE-column key for this sample. Override if it differs from
+        the on-disk basename (e.g. prostate MRI uses .jpg keys for .dcm files)."""
+        return name
+
+    def __getitem__(self, idx: int):
+        name = self.image_ids[idx]
+
+        image = self.img_tf(self._load_image(name))
+
+        mask = self.mask_tf(Image.open(self._mask_path(name)).convert("L"))
+        mask = (mask > 0.5).float()
+
+        fix_rows = self.fix_df.get_group(self._fixation_key(name))
+        fixation = get_scanpath(fix_rows, self.max_len)
+
+        return {"image": image, "fixation": fixation, "mask": mask, "name": name}
+
+
+class KvasirSEGDataset(_BaseGazeDataset):
+    """Kvasir-SEG colonoscopy polyp segmentation. images/masks share filenames
+    (e.g. ``cju0qkwl35piu0993l0dewei2.jpg`` in both folders)."""
+
+    def __init__(self, root: str, fixation_csv: str, img_size: int = IMG_SIZE):
+        super().__init__(root, fixation_csv, img_size)
+        img_files = set(os.listdir(self.img_dir))
+        mask_files = set(os.listdir(self.mask_dir))
+        csv_imgs = set(self.df[IMG_COL].unique())
+        self.image_ids = sorted(img_files & mask_files & csv_imgs)
+        if not self.image_ids:
+            raise RuntimeError("No overlap between images/, masks/ and the fixation CSV.")
+        print(f"[Dataset] Kvasir-SEG valid samples = {len(self.image_ids)}")
+
+    def _load_image(self, name: str) -> Image.Image:
+        return Image.open(os.path.join(self.img_dir, name)).convert("RGB")
+
+    def _mask_path(self, name: str) -> str:
+        return os.path.join(self.mask_dir, name)
+
+
+class ProstateMRIDataset(_BaseGazeDataset):
+    """NCI-ISBI prostate MRI. images/ holds DICOM (.dcm), masks/ holds PNG,
+    and the fixation CSV references each case as ``<basename>.jpg`` (the
+    format the eye-tracking session was actually rendered/displayed in)."""
+
+    def __init__(self, root: str, fixation_csv: str, img_size: int = IMG_SIZE):
+        super().__init__(root, fixation_csv, img_size)
+
+        dcm_basenames = {os.path.splitext(f)[0] for f in os.listdir(self.img_dir) if f.endswith(".dcm")}
+        png_basenames = {os.path.splitext(f)[0] for f in os.listdir(self.mask_dir) if f.endswith(".png")}
+        csv_basenames = {os.path.splitext(f)[0] for f in self.df[IMG_COL].unique()}
+
+        self.image_ids = sorted(dcm_basenames & png_basenames & csv_basenames)
+        if not self.image_ids:
+            raise RuntimeError("No overlap between images/ (.dcm), masks/ (.png) and the fixation CSV.")
+        print(f"[Dataset] Prostate MRI valid samples = {len(self.image_ids)}")
+
+    def _load_image(self, name: str) -> Image.Image:
+        import pydicom  # local import: optional dependency, only needed for DICOM datasets
+
+        dcm_path = os.path.join(self.img_dir, f"{name}.dcm")
+        dicom = pydicom.dcmread(dcm_path)
+        arr = dicom.pixel_array.astype(np.float32)
+
+        if arr.max() > arr.min():
+            arr = (arr - arr.min()) / (arr.max() - arr.min())
+        else:
+            arr = np.zeros_like(arr)
+
+        rgb = np.stack([arr, arr, arr], axis=-1)  # grayscale -> 3-channel, DINOv3 expects RGB
+        return Image.fromarray((rgb * 255).astype(np.uint8))
+
+    def _mask_path(self, name: str) -> str:
+        return os.path.join(self.mask_dir, f"{name}.png")
+
+    def _fixation_key(self, name: str) -> str:
+        return f"{name}.jpg"
+
+
+DATASET_REGISTRY = {
+    "kvasir": KvasirSEGDataset,
+    "prostate_mri": ProstateMRIDataset,
+}

notebooks/gazerefine/gaze.py

@@ -0,0 +1,119 @@
+"""
+gaze.py — Turn an eye-tracking scanpath into a spatial prior.
+
+Implements Eq. (gaze heatmap) from the paper:
+
+    H(u, v) = sum_m  d_m * exp( -((u - x_m)^2 + (v - y_m)^2) / (2 * sigma^2) )
+
+i.e. a duration-weighted sum of 2D Gaussians, one per fixation, evaluated on
+the DINOv3 patch grid. The result is min-max normalized and used directly as
+the soft foreground prior; ``1 - H`` is the background complement.
+
+Expected CSV columns (standard EyeLink-style exports, e.g. SR Research /
+Tobii fixation reports):
+
+    IMAGE                   image filename the fixation belongs to
+    CURRENT_FIX_INDEX       fixation order within the trial
+    CURRENT_FIX_X            fixation x in image-normalized [0, 1] coords
+    CURRENT_FIX_Y            fixation y in image-normalized [0, 1] coords
+    CURRENT_FIX_DURATION     fixation duration (ms or any consistent unit)
+
+If your fixation coordinates are in raw pixels rather than [0, 1], divide by
+the screen/image width and height before building the CSV — GazeRefine
+expects normalized coordinates so it stays resolution-independent.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+import torch
+
+TIME_COL = "CURRENT_FIX_DURATION"
+FIX_IDX_COL = "CURRENT_FIX_INDEX"
+FIX_X_COL = "CURRENT_FIX_X"
+FIX_Y_COL = "CURRENT_FIX_Y"
+IMG_COL = "IMAGE"
+
+
+def get_scanpath(df: pd.DataFrame, max_len: int, time_col: str = TIME_COL) -> torch.Tensor:
+    """Build a zero-padded ``(max_len, 3)`` scanpath tensor from a fixation-report
+    DataFrame already filtered down to a single image.
+
+    Columns of the output: ``[x_norm, y_norm, t_norm]``, all in ``[0, 1]``.
+    ``t_norm`` is the fixation duration min-max normalized within that image's
+    own scanpath (longest fixation -> 1.0). Rows are sorted by
+    ``CURRENT_FIX_INDEX`` so temporal order is preserved (only the spatial +
+    duration channels are currently used by the model, but order is kept for
+    future scanpath-aware extensions).
+    """
+    if len(df) == 0:
+        return torch.zeros(max_len, 3)
+
+    df = df.sort_values(FIX_IDX_COL)
+    x = df[FIX_X_COL].to_numpy(np.float32)
+    y = df[FIX_Y_COL].to_numpy(np.float32)
+    t = df[time_col].to_numpy(np.float32)
+
+    t_min, t_max = t.min(), t.max()
+    t = (t - t_min) / (t_max - t_min + 1e-8)
+
+    seq = np.stack([x, y, t], axis=1)
+    if len(seq) > max_len:
+        seq = seq[:max_len]
+    if len(seq) < max_len:
+        seq = np.concatenate([seq, np.zeros((max_len - len(seq), 3), np.float32)], axis=0)
+
+    return torch.tensor(seq, dtype=torch.float32)
+
+
+def generate_gaze_heatmap(
+    fixation: torch.Tensor, h_patch: int, sigma: float = 2.0
+) -> torch.Tensor:
+    """Convert a batch of scanpaths into duration-weighted Gaussian heatmaps on
+    the DINOv3 patch grid.
+
+    Parameters
+    ----------
+    fixation : (B, L, 3) tensor of ``[x_norm, y_norm, t_norm]``, zero-padded.
+    h_patch : side length of the square patch grid (e.g. 37 for a 518/14 ViT,
+        37*2=74 for the 518/16 setup used in our experiments — see
+        ``gazerefine.model.H_PATCH``).
+    sigma : standard deviation (in patch units) of each fixation's Gaussian.
+        Larger sigma -> a softer, more diffuse foreground prior; smaller
+        sigma -> a tighter prior centered exactly on the fixated patches.
+
+    Returns
+    -------
+    (B, h_patch, h_patch) tensor, each sample independently min-max
+    normalized to ``[0, 1]`` (an all-zero map stays all-zero if a sample has
+    no valid fixations).
+    """
+    B, L, _ = fixation.shape
+    device = fixation.device
+
+    y_coords = torch.arange(h_patch, device=device, dtype=torch.float32)
+    x_coords = torch.arange(h_patch, device=device, dtype=torch.float32)
+    grid_y, grid_x = torch.meshgrid(y_coords, x_coords, indexing="ij")
+
+    heatmaps = []
+    for b in range(B):
+        # zero-padding rows are exactly all-zero -> filter them out
+        valid = fixation[b].abs().sum(dim=-1) > 1e-8
+        valid_fix = fixation[b][valid]
+
+        h_b = torch.zeros(h_patch, h_patch, device=device)
+        if len(valid_fix) > 0:
+            for k in range(valid_fix.shape[0]):
+                x_k, y_k, t_k = valid_fix[k, 0], valid_fix[k, 1], valid_fix[k, 2]
+                gx = x_k * (h_patch - 1)
+                gy = y_k * (h_patch - 1)
+                dist_sq = (grid_x - gx) ** 2 + (grid_y - gy) ** 2
+                h_b += t_k * torch.exp(-dist_sq / (2 * sigma ** 2))
+
+            h_max = h_b.max()
+            if h_max > 1e-8:
+                h_b = h_b / h_max
+        heatmaps.append(h_b)
+
+    return torch.stack(heatmaps, dim=0)

notebooks/gazerefine/metrics.py

@@ -0,0 +1,27 @@
+"""metrics.py — Dice / IoU against a binarized prediction."""
+
+from __future__ import annotations
+
+import torch
+
+
+@torch.no_grad()
+def compute_metrics(preds: torch.Tensor, mask: torch.Tensor, thr: float = 0.5) -> dict:
+    """
+    preds, mask : (B, 1, H, W). ``preds`` is the continuous [0, 1] GazeRefine
+    output, ``mask`` is the ground-truth binary mask.
+    """
+    pred_bin = (preds > thr).float()
+    tp = (pred_bin * mask).sum((1, 2, 3))
+    fp = (pred_bin * (1 - mask)).sum((1, 2, 3))
+    fn = ((1 - pred_bin) * mask).sum((1, 2, 3))
+
+    dice_per = 2 * tp / (2 * tp + fp + fn + 1e-8)
+    iou_per = tp / (tp + fp + fn + 1e-8)
+
+    return {
+        "dice": dice_per.mean().item(),
+        "iou": iou_per.mean().item(),
+        "dice_per": dice_per,
+        "iou_per": iou_per,
+    }

notebooks/gazerefine/model.py

@@ -0,0 +1,242 @@
+"""
+model.py — The GazeRefine zero-shot segmentation model.
+
+This is a 1:1 refactor of the two task-specific scripts (Kvasir-SEG polyp /
+NCI-ISBI prostate MRI) into a single, dataset-agnostic module. The numerics
+are unchanged from the original experiments — only structure, naming and
+comments were cleaned up. Every per-dataset difference (sigma, kNN k,
+gaze-anchor weight, max_iters, ...) is now a constructor / config argument
+instead of a hardcoded default, see ``configs/*.yaml``.
+
+Maps onto the paper (Section 2) as follows:
+
+    gaze.generate_gaze_heatmap   -> H(u, v)                          (Eq. gaze prior)
+    W_fg, W_bg                   -> W_fg^(0), W_bg^(0)                (gaze-derived weights)
+    F_proto_init, B_proto_init   -> F^(0), B^(0)                      (initial prototypes)
+    contrast_method="difference" -> S_i^(t) = max(0, s_fg - alpha*s_bg)  (contrastive cleaning)
+    knn_affinity_refinement      -> S̄_i^(t)                           (kNN affinity propagation)
+    gaze_anchor_weight           -> lambda                            (anchor blending strength)
+    convergence check            -> ||F^(t+1) - F^(t)|| < eps
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .backbone import FrozenDINOv3
+from .gaze import generate_gaze_heatmap
+from .constants import IMG_SIZE, PATCH_SIZE
+
+EPS = 1e-12
+
+
+def knn_affinity_refinement(
+    Pv: torch.Tensor, S: torch.Tensor, k: int = 5, temperature: float = 0.05
+) -> torch.Tensor:
+    """Propagate a per-patch score map ``S`` across its k nearest neighbors in
+    frozen DINOv3 feature space (patch-to-patch affinity), encouraging
+    coherent, object-level responses instead of isolated high-confidence
+    patches. Corresponds to S̄ in the paper.
+
+    Pv : (B, N, D) raw (unnormalized) patch embeddings.
+    S  : (B, N)    current per-patch score to be smoothed.
+    """
+    B, N, D = Pv.shape
+    Pv_norm = F.normalize(Pv, p=2, dim=-1)
+
+    sim_matrix = torch.bmm(Pv_norm, Pv_norm.transpose(1, 2))            # (B, N, N)
+    topk_vals, topk_indices = torch.topk(sim_matrix, k=k, dim=-1)        # (B, N, k)
+    weights = F.softmax(topk_vals / temperature, dim=-1)                 # (B, N, k)
+
+    S_expanded = S.unsqueeze(1).expand(B, N, N)
+    S_neighbors = torch.gather(S_expanded, dim=2, index=topk_indices)    # (B, N, k)
+
+    return torch.sum(weights * S_neighbors, dim=-1)                     # (B, N)
+
+
+class GazeRefine(nn.Module):
+    """Zero-shot, training-free, gaze-guided segmentation model.
+
+    The model has exactly one set of *learned* weights: the frozen,
+    pretrained DINOv3 backbone. Everything else — prototype construction,
+    contrastive cleaning, kNN propagation, gaze anchoring — is a closed-form
+    operation re-run from scratch on every image at inference time.
+
+    Parameters
+    ----------
+    dino_name : timm DINOv3 checkpoint name.
+    img_size, patch_size : input resolution / ViT patch size. Must satisfy
+        ``img_size % patch_size == 0``.
+    sigma : Gaussian spread (in patch units) for the gaze heatmap.
+    extract_mode : ``"last"`` (1 block) or ``"all"`` (4 blocks, averaged).
+    contrast_method : ``"difference"`` (paper default, contrastive cleaning),
+        ``"softmax"`` (foreground/background softmax ratio), or
+        ``"original"`` (plain foreground cosine similarity, no background
+        suppression — kept for the ablation in Table 2).
+    temperature : softmax temperature, only used when contrast_method="softmax".
+    max_iters : maximum recurrent refinement iterations (T in the paper).
+    knn_refine : whether to apply kNN affinity propagation each iteration.
+    knn_k, knn_temp : kNN neighborhood size / softmax temperature.
+    gaze_anchor_weight : lambda — how strongly each iteration's prototypes
+        are pulled back toward the original gaze-only prototypes. Higher =
+        trust the raw fixations more; lower = let the model drift further
+        from the initial gaze region.
+    """
+
+    def __init__(
+        self,
+        dino_name: str = "vit_large_patch16_dinov3.lvd1689m",
+        img_size: int = IMG_SIZE,
+        patch_size: int = PATCH_SIZE,
+        sigma: float = 2.0,
+        extract_mode: str = "last",
+        contrast_method: str = "difference",
+        temperature: float = 0.05,
+        max_iters: int = 10,
+        knn_refine: bool = True,
+        knn_k: int = 5,
+        knn_temp: float = 0.1,
+        gaze_anchor_weight: float = 0.6,
+    ):
+        super().__init__()
+        assert img_size % patch_size == 0, "img_size must be a multiple of patch_size"
+        self.img_size = img_size
+        self.h_patch = img_size // patch_size
+        self.sigma = sigma
+        self.contrast_method = contrast_method
+        self.temperature = temperature
+        self.max_iters = max_iters
+        self.knn_refine = knn_refine
+        self.knn_k = knn_k
+        self.knn_temp = knn_temp
+        self.gaze_anchor_weight = gaze_anchor_weight
+
+        self.visual_enc = FrozenDINOv3(dino_name, extract_mode=extract_mode)
+
+    # ------------------------------------------------------------------ #
+    def _gaze_weights(self, fixation: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Scanpath -> heatmap -> normalized foreground/background spatial weights."""
+        gaze_heatmap = generate_gaze_heatmap(fixation, h_patch=self.h_patch, sigma=self.sigma)
+        H_flat = gaze_heatmap.view(gaze_heatmap.size(0), -1)
+
+        W_fg = H_flat / (H_flat.sum(dim=-1, keepdim=True) + EPS)
+        bg_w = 1.0 - H_flat
+        W_bg = bg_w / (bg_w.sum(dim=-1, keepdim=True) + EPS)
+        return gaze_heatmap, W_fg, W_bg
+
+    def _score(self, Pv_norm: torch.Tensor, F_proto_norm: torch.Tensor, B_proto_norm: torch.Tensor) -> torch.Tensor:
+        """Foreground/background contrastive scoring for one level, one iteration."""
+        if self.contrast_method == "original":
+            S = torch.bmm(Pv_norm, F_proto_norm.unsqueeze(-1)).squeeze(-1)
+            return torch.clamp(S, min=0.0)
+
+        sim_fg = torch.bmm(Pv_norm, F_proto_norm.unsqueeze(-1)).squeeze(-1)
+        sim_bg = torch.bmm(Pv_norm, B_proto_norm.unsqueeze(-1)).squeeze(-1)
+
+        if self.contrast_method == "difference":
+            return torch.clamp(sim_fg - sim_bg, min=0.0)
+        elif self.contrast_method == "softmax":
+            stacked = torch.stack([sim_fg, sim_bg], dim=-1) / self.temperature
+            probs = F.softmax(stacked, dim=-1)
+            return probs[:, :, 0]
+        raise ValueError(f"Unknown contrast_method: {self.contrast_method!r}")
+
+    def _refine_level(self, Pv: torch.Tensor, W_fg: torch.Tensor, W_bg: torch.Tensor) -> torch.Tensor:
+        """Run the full recurrent gaze-anchored refinement loop for one DINOv3 level.
+        Returns the final (B, N) per-patch foreground score map for that level."""
+        B = Pv.size(0)
+        device = Pv.device
+        Pv_norm = F.normalize(Pv, p=2, dim=-1)
+
+        # Initial gaze-only prototypes — F^(0), B^(0)
+        F_proto_init = torch.sum(W_fg.unsqueeze(-1) * Pv, dim=1)
+        B_proto_init = torch.sum(W_bg.unsqueeze(-1) * Pv, dim=1)
+
+        W_fg_curr, W_bg_curr = W_fg.clone(), W_bg.clone()
+        best_S = W_fg.clone()
+        S_prev = W_fg.clone()
+        active = torch.ones(B, dtype=torch.bool, device=device)
+
+        for _ in range(self.max_iters):
+            W_fg_norm = W_fg_curr / (W_fg_curr.sum(dim=-1, keepdim=True) + EPS)
+            W_bg_norm = W_bg_curr / (W_bg_curr.sum(dim=-1, keepdim=True) + EPS)
+
+            F_proto_curr = torch.sum(W_fg_norm.unsqueeze(-1) * Pv, dim=1)
+            B_proto_curr = torch.sum(W_bg_norm.unsqueeze(-1) * Pv, dim=1)
+
+            # Anchor toward the initial gaze-only prototypes — lambda blending
+            lam = self.gaze_anchor_weight
+            F_proto = (1 - lam) * F_proto_curr + lam * F_proto_init
+            B_proto = (1 - lam) * B_proto_curr + lam * B_proto_init
+            F_proto_norm = F.normalize(F_proto, p=2, dim=-1)
+            B_proto_norm = F.normalize(B_proto, p=2, dim=-1)
+
+            S_iter = self._score(Pv_norm, F_proto_norm, B_proto_norm)
+            if self.knn_refine:
+                S_iter = knn_affinity_refinement(Pv, S_iter, k=self.knn_k, temperature=self.knn_temp)
+
+            S_new = S_prev.clone()
+            S_new[active] = S_iter[active]
+
+            # Collapse prevention: if an active sample's map flattened to ~0,
+            # roll it back and freeze it instead of letting it degenerate.
+            collapsed = S_new.max(dim=-1)[0] < 1e-4
+            freeze = active & collapsed
+            S_new[freeze] = S_prev[freeze]
+            active = active & ~collapsed
+
+            best_S = S_new.clone()
+            if not active.any():
+                break
+
+            # Turn the refined score into the next iteration's spatial weights
+            S_sig = torch.sigmoid(S_new)
+            W_fg_next = S_sig / (S_sig.sum(dim=-1, keepdim=True) + EPS)
+            bg_w = 1.0 - S_sig
+            bg_w = torch.clamp(bg_w - bg_w.min(dim=-1, keepdim=True).values, min=0.0)
+            W_bg_next = bg_w / (bg_w.sum(dim=-1, keepdim=True) + EPS)
+
+            F_proto_next = torch.sum(W_fg_next.unsqueeze(-1) * Pv, dim=1)
+            F_proto_next_norm = F.normalize(F_proto_next, p=2, dim=-1)
+
+            # Convergence: foreground prototype stopped moving
+            diff = torch.abs(F_proto_next_norm - F_proto_norm).mean(dim=-1)
+            if torch.all(diff < 1e-6):
+                break
+
+            S_prev = S_new.clone()
+            W_fg_curr[active] = W_fg_next[active]
+            W_bg_curr[active] = W_bg_next[active]
+
+        return best_S
+
+    # ------------------------------------------------------------------ #
+    def forward(self, image: torch.Tensor, fixation: torch.Tensor) -> dict:
+        """
+        image : (B, 3, H, W) normalized RGB tensor (ImageNet mean/std).
+        fixation : (B, L, 3) zero-padded scanpath, see ``gazerefine.gaze.get_scanpath``.
+
+        Returns a dict with:
+            preds        (B, 1, H, W) final mask in [0, 1] — threshold at 0.5 for a binary mask
+            gaze_heatmap (B, h, w) the raw gaze prior, useful for visualization
+            final_map    (B, N) the un-upsampled patch-level score map
+        """
+        B = image.size(0)
+        Pv_list = self.visual_enc(image)
+        gaze_heatmap, W_fg, W_bg = self._gaze_weights(fixation)
+
+        all_maps = [self._refine_level(Pv, W_fg, W_bg) for Pv in Pv_list]
+        final_map = torch.stack(all_maps, dim=1).mean(dim=1)  # average across levels
+
+        grid = final_map.view(B, 1, self.h_patch, self.h_patch)
+        upsampled = F.interpolate(grid, size=(self.img_size, self.img_size), mode="bilinear", align_corners=False)
+
+        # per-sample min-max normalization -> ready-to-threshold mask
+        flat = upsampled.view(B, -1)
+        m_min = flat.min(dim=-1, keepdim=True).values.view(B, 1, 1, 1)
+        m_max = flat.max(dim=-1, keepdim=True).values.view(B, 1, 1, 1)
+        preds = (upsampled - m_min) / (m_max - m_min + 1e-8)
+
+        return dict(preds=preds, gaze_heatmap=gaze_heatmap, final_map=final_map)

notebooks/gazerefine/visualize.py

@@ -0,0 +1,86 @@
+"""visualize.py — Plotting / overlay helpers.
+
+``save_prediction`` reproduces the 4-panel (image / gaze / GT / prediction)
+grid used during batch evaluation. ``overlay_heatmap`` and ``overlay_mask``
+are lighter-weight, matplotlib-free PIL helpers meant for the Gradio demo
+and the example notebook, where you typically want a single composited
+image rather than a saved subplot file.
+"""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+import torch
+from PIL import Image
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import matplotlib.cm as cm
+
+from .constants import IMG_MEAN, IMG_STD
+
+
+def _denormalize(img: torch.Tensor) -> np.ndarray:
+    """(3, H, W) normalized tensor -> (H, W, 3) uint8-range float array in [0, 1]."""
+    mean = torch.tensor(IMG_MEAN, device=img.device).view(3, 1, 1)
+    std = torch.tensor(IMG_STD, device=img.device).view(3, 1, 1)
+    return (img * std + mean).clamp(0, 1).permute(1, 2, 0).cpu().numpy()
+
+
+@torch.no_grad()
+def save_prediction(
+    img: torch.Tensor,
+    gaze_heatmap: torch.Tensor,
+    gt: torch.Tensor,
+    pred: torch.Tensor,
+    name: str,
+    out_dir: str = "predictions",
+) -> None:
+    """Save a 4-panel [image | gaze heatmap | GT mask | predicted mask] figure."""
+    os.makedirs(out_dir, exist_ok=True)
+
+    img_np = _denormalize(img)
+    gaze_np = gaze_heatmap.squeeze().cpu().numpy()
+    gt_np = gt.squeeze().cpu().numpy()
+    pred_np = pred.squeeze().cpu().numpy()
+
+    fig, ax = plt.subplots(1, 4, figsize=(20, 5))
+    ax[0].imshow(img_np); ax[0].set_title("Image")
+    ax[1].imshow(gaze_np, cmap="jet"); ax[1].set_title("Gaze Heatmap")
+    ax[2].imshow(gt_np, cmap="gray"); ax[2].set_title("GT mask")
+    ax[3].imshow(pred_np, cmap="gray"); ax[3].set_title("Predicted mask")
+    for a in ax:
+        a.axis("off")
+    plt.tight_layout()
+    plt.savefig(os.path.join(out_dir, name), dpi=100, bbox_inches="tight")
+    plt.close()
+
+
+def overlay_heatmap(image: Image.Image, heatmap: np.ndarray, alpha: float = 0.45) -> Image.Image:
+    """Composite a [0, 1] heatmap (any H'xW', will be resized) onto a PIL image
+    using a jet colormap. Used to show the gaze prior over the original image."""
+    heatmap = np.asarray(heatmap, dtype=np.float32)
+    heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min() + 1e-8)
+
+    heat_img = Image.fromarray((cm.jet(heatmap)[:, :, :3] * 255).astype(np.uint8))
+    heat_img = heat_img.resize(image.size, resample=Image.BILINEAR)
+
+    base = image.convert("RGB")
+    return Image.blend(base, heat_img, alpha=alpha)
+
+
+def overlay_mask(
+    image: Image.Image, mask: np.ndarray, color: tuple[int, int, int] = (255, 60, 60), alpha: float = 0.45
+) -> Image.Image:
+    """Composite a binary/soft [0, 1] mask onto a PIL image as a solid color wash."""
+    mask = np.asarray(mask, dtype=np.float32)
+    mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)
+    mask_img = Image.fromarray((mask * 255).astype(np.uint8)).resize(image.size, resample=Image.NEAREST)
+
+    base = image.convert("RGB")
+    color_layer = Image.new("RGB", base.size, color)
+    composited = Image.composite(color_layer, base, mask_img.point(lambda p: int(p * alpha)))
+    return composited

notebooks/huggingface_space/README.md

@@ -0,0 +1,42 @@
+---
+title: GazeRefine
+emoji: 👁️
+colorFrom: blue
+colorTo: red
+sdk: gradio
+sdk_version: 4.44.0
+app_file: app.py
+pinned: false
+license: mit
+short_description: Zero-shot, training-free gaze-guided medical segmentation
+---
+
+# GazeRefine — Expert Gaze as a Test-Time Prompt
+
+Interactive demo for **GazeRefine**, a training-free, zero-shot framework
+that turns expert eye-gaze into an inference-time prompt for medical image
+segmentation. Frozen DINOv3 patch features + gaze-anchored
+foreground/background prototypes + recurrent contrastive cleaning + kNN
+affinity propagation — no masks, no clicks-as-boxes, no fine-tuning, no
+adapters, no prompt encoder.
+
+## How to use
+1. Upload a colonoscopy or grayscale-MRI-style image.
+2. Click on the image 1–5 times where a clinician's gaze would land on the
+   structure of interest (a polyp, the prostate, ...). Each click adds a
+   numbered fixation marker; the slider controls that fixation's relative
+   duration/weight before your next click.
+3. Pick a hyperparameter preset (tuned per-modality, see the paper).
+4. Press **Run GazeRefine** to get the gaze-prior overlay and the predicted
+   segmentation mask.
+
+## Notes
+- Inference uses a frozen `vit_large_patch16_dinov3.lvd1689m` backbone from
+  `timm`. First run will download the checkpoint.
+- CPU inference works but is slow; a GPU Space is recommended for a smooth
+  demo.
+- This Space is for research/demonstration only — it is **not** a clinical
+  diagnostic tool.
+
+Full code, configs, and the unified evaluation pipeline:
+[GitHub repository](https://github.com/<your-org>/gazerefine).

notebooks/huggingface_space/app.py

@@ -0,0 +1,138 @@
+"""
+app.py — GazeRefine interactive demo (Hugging Face Space).
+
+Upload a medical image, click on it to drop a few fixation points the way a
+clinician's gaze would land on the structure of interest, and GazeRefine
+turns that into a segmentation mask — fully zero-shot, no training, no
+prompt-specific architecture.
+
+Run locally with:
+    pip install -r huggingface_space/requirements.txt
+    python huggingface_space/app.py
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+from PIL import Image, ImageDraw
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))  # repo root
+from scripts.predict_single import predict
+
+PRESETS = {
+    "Colonoscopy / polyp (Kvasir-SEG settings)": dict(
+        sigma=2.0, contrast_method="difference", max_iters=5,
+        gaze_anchor_weight=0.5, knn_k=20, knn_temp=0.1,
+    ),
+    "Grayscale MRI / CT (prostate-MRI settings)": dict(
+        sigma=1.5, contrast_method="difference", max_iters=1,
+        gaze_anchor_weight=0.8, knn_k=3, knn_temp=0.1,
+    ),
+}
+DINO_NAME = "vit_large_patch16_dinov3.lvd1689m"
+
+POINT_COLORS = ["#ff3b30", "#ff9500", "#ffcc00", "#34c759", "#5ac8fa", "#007aff", "#af52de"]
+
+
+def draw_points(image: Image.Image, points: list[tuple[float, float, float]]) -> Image.Image:
+    """Render numbered fixation markers over the image for visual feedback."""
+    if image is None:
+        return None
+    vis = image.convert("RGB").copy()
+    draw = ImageDraw.Draw(vis)
+    w, h = vis.size
+    r = max(6, min(w, h) // 80)
+    for i, (x, y, dur) in enumerate(points):
+        cx, cy = x * w, y * h
+        color = POINT_COLORS[i % len(POINT_COLORS)]
+        rad = r * (0.6 + 0.8 * dur)  # bigger marker = longer fixation
+        draw.ellipse([cx - rad, cy - rad, cx + rad, cy + rad], outline=color, width=3)
+        draw.text((cx + rad + 2, cy - rad), str(i + 1), fill=color)
+    return vis
+
+
+def on_select(image: Image.Image, points: list, duration: float, evt: gr.SelectData):
+    if image is None:
+        gr.Warning("Upload an image first.")
+        return points, None
+    w, h = image.size
+    x_px, y_px = evt.index
+    points = points + [(x_px / w, y_px / h, duration)]
+    return points, draw_points(image, points)
+
+
+def on_clear(image: Image.Image):
+    return [], image
+
+
+def on_image_change(image: Image.Image):
+    # new image -> reset fixations
+    return [], image
+
+
+def run(image: Image.Image, points: list, preset_name: str, threshold: float):
+    if image is None:
+        gr.Warning("Upload an image first.")
+        return None, None, None
+    if len(points) == 0:
+        gr.Warning("Click on the image at least once to place a fixation.")
+        return None, None, None
+
+    cfg = PRESETS[preset_name]
+    out = predict(
+        image=image,
+        fixations=points,
+        dino_name=DINO_NAME,
+        threshold=threshold,
+        **cfg,
+    )
+    mask_only = Image.fromarray((np.clip(out["mask_bin"], 0, 1) * 255).astype(np.uint8))
+    return out["gaze_overlay"], out["mask_overlay"], mask_only
+
+
+with gr.Blocks(title="GazeRefine — gaze-guided zero-shot segmentation") as demo:
+    gr.Markdown(
+        """
+        # 👁️ GazeRefine — Expert Gaze as a Test-Time Prompt
+        Training-free, zero-shot medical image segmentation. Upload an image, **click on it
+        1–5 times** where a clinician would look, pick a preset, and run.
+        No masks, clicks-as-bounding-boxes, fine-tuning, or adapters — just a frozen DINOv3
+        backbone steered by your gaze. See the paper / code on GitHub (linked below).
+        """
+    )
+
+    points_state = gr.State([])
+
+    with gr.Row():
+        with gr.Column():
+            image_in = gr.Image(type="pil", label="1. Upload image, then click to place fixations", height=420)
+            with gr.Row():
+                duration_slider = gr.Slider(0.1, 1.0, value=1.0, step=0.1, label="Next fixation's relative duration")
+                clear_btn = gr.Button("Clear fixations")
+            preset = gr.Radio(list(PRESETS.keys()), value=list(PRESETS.keys())[0], label="2. Hyperparameter preset")
+            threshold = gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="3. Mask threshold")
+            run_btn = gr.Button("Run GazeRefine", variant="primary")
+
+        with gr.Column():
+            gaze_out = gr.Image(label="Gaze prior over image", height=260)
+            with gr.Row():
+                mask_overlay_out = gr.Image(label="Predicted mask (overlay)", height=260)
+                mask_only_out = gr.Image(label="Predicted mask (binary)", height=260)
+
+    image_in.upload(on_image_change, inputs=[image_in], outputs=[points_state, image_in])
+    image_in.select(on_select, inputs=[image_in, points_state, duration_slider], outputs=[points_state, image_in])
+    clear_btn.click(on_clear, inputs=[image_in], outputs=[points_state, image_in])
+    run_btn.click(run, inputs=[image_in, points_state, preset, threshold], outputs=[gaze_out, mask_overlay_out, mask_only_out])
+
+    gr.Markdown(
+        "Code: [GitHub](https://github.com/<your-org>/gazerefine) · "
+        "Method: GazeRefine — frozen DINOv3 + gaze-anchored prototypes + recurrent "
+        "foreground/background refinement, entirely training-free."
+    )
+
+if __name__ == "__main__":
+    demo.launch()

notebooks/huggingface_space/requirements.txt

@@ -0,0 +1,7 @@
+torch>=2.1
+torchvision>=0.16
+timm>=1.0.0
+numpy
+pandas
+pillow
+gradio>=4.0

notebooks/requirements.txt

@@ -0,0 +1,9 @@
+torch>=2.1
+torchvision>=0.16
+timm>=1.0.0          # provides DINOv3 checkpoints (vit_*_dinov3.*)
+numpy
+pandas
+pillow
+pyyaml
+matplotlib
+pydicom               # only required for the prostate MRI (DICOM) dataset

requirements.txt

@@ -0,0 +1,16 @@
+torch>=2.1
+torchvision>=0.16
+gradio==4.44.1
+starlette>=0.37.2,<0.39
+timm>=1.0.0
+numpy
+pandas
+pillow
+matplotlib
+pydicom
+openpyxl
+huggingface-hub==0.24.0
+pydub==0.25.1
+diffusers>=0.30
+transformers>=4.44
+accelerate>=0.33

scripts/predict_single.py

@@ -0,0 +1,311 @@
+"""
+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()

scripts/run_eval.py

@@ -0,0 +1,129 @@
+"""
+run_eval.py — Unified zero-shot evaluation entrypoint.
+
+Usage
+-----
+    python scripts/run_eval.py --config configs/kvasir.yaml
+    python scripts/run_eval.py --config configs/prostate_mri.yaml
+
+    # override anything from the YAML on the command line:
+    python scripts/run_eval.py --config configs/kvasir.yaml --root /data/Kvasir-SEG --max_iters 8
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+import yaml
+from torch.utils.data import DataLoader
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))  # repo root, for `import gazecorrect`
+
+from gazecorrect import GazeRefine, compute_metrics, save_prediction
+from gazecorrect.datasets import DATASET_REGISTRY
+
+
+def load_config(path: str, overrides: argparse.Namespace) -> dict:
+    with open(path) as f:
+        cfg = yaml.safe_load(f)
+    for k, v in vars(overrides).items():
+        if k != "config" and v is not None:
+            cfg[k] = v
+    return cfg
+
+
+@torch.no_grad()
+def eval_epoch(model, loader, device, threshold: float, pred_dir: str, save_num: int) -> dict:
+    model.eval()
+    all_dice, all_iou = [], []
+    n = len(loader)
+    saved = 0
+
+    for i, batch in enumerate(loader):
+        img = batch["image"].to(device)
+        fix = batch["fixation"].to(device)
+        msk = batch["mask"].to(device)
+
+        out = model(img, fix)
+        preds = out["preds"]
+        m = compute_metrics(preds, msk, thr=threshold)
+        all_dice.append(m["dice_per"].cpu())
+        all_iou.append(m["iou_per"].cpu())
+
+        if saved < save_num:
+            for ridx in range(img.size(0)):
+                if saved >= save_num:
+                    break
+                pred_bin = (preds[ridx] > threshold).float()
+                save_prediction(
+                    img[ridx], out["gaze_heatmap"][ridx], msk[ridx], pred_bin,
+                    name=f"zeroshot_sample_{batch['name'][ridx]}.png", out_dir=pred_dir,
+                )
+                saved += 1
+
+        if (i + 1) % max(1, n // 5) == 0:
+            print(f"  [{i + 1}/{n}] dice={m['dice']:.4f} iou={m['iou']:.4f}")
+
+    all_dice = torch.cat(all_dice)
+    all_iou = torch.cat(all_iou)
+    return {
+        "dice": all_dice.mean().item(), "dice_std": all_dice.std().item(),
+        "iou": all_iou.mean().item(), "iou_std": all_iou.std().item(),
+    }
+
+
+def main():
+    ap = argparse.ArgumentParser("GazeRefine — zero-shot gaze-guided segmentation")
+    ap.add_argument("--config", required=True, help="path to a YAML config, e.g. configs/kvasir.yaml")
+    # optional CLI overrides for anything in the YAML
+    ap.add_argument("--root", type=str, default=None)
+    ap.add_argument("--fixation_csv", type=str, default=None)
+    ap.add_argument("--pred_dir", type=str, default=None)
+    ap.add_argument("--dino_name", type=str, default=None)
+    ap.add_argument("--batch_size", type=int, default=None)
+    ap.add_argument("--num_workers", type=int, default=None)
+    ap.add_argument("--sigma", type=float, default=None)
+    ap.add_argument("--threshold", type=float, default=None)
+    ap.add_argument("--extract_mode", type=str, default=None, choices=["all", "last"])
+    ap.add_argument("--contrast_method", type=str, default=None, choices=["difference", "softmax", "original"])
+    ap.add_argument("--temperature", type=float, default=None)
+    ap.add_argument("--max_iters", type=int, default=None)
+    ap.add_argument("--gaze_anchor_weight", type=float, default=None)
+    ap.add_argument("--knn_k", type=int, default=None)
+    ap.add_argument("--save_num", type=int, default=None)
+    args = ap.parse_args()
+
+    cfg = load_config(args.config, args)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"[GazeRefine] device={device}, dataset={cfg['dataset']}")
+
+    DatasetCls = DATASET_REGISTRY[cfg["dataset"]]
+    dataset = DatasetCls(cfg["root"], cfg["fixation_csv"])
+    loader = DataLoader(dataset, batch_size=cfg["batch_size"], num_workers=cfg["num_workers"], shuffle=False, pin_memory=True)
+    print(f"[GazeRefine] total samples to evaluate: {len(dataset)}")
+
+    model = GazeRefine(
+        dino_name=cfg["dino_name"],
+        sigma=cfg["sigma"],
+        extract_mode=cfg["extract_mode"],
+        contrast_method=cfg["contrast_method"],
+        temperature=cfg["temperature"],
+        max_iters=cfg["max_iters"],
+        knn_refine=cfg["knn_refine"],
+        knn_k=cfg["knn_k"],
+        knn_temp=cfg["knn_temp"],
+        gaze_anchor_weight=cfg["gaze_anchor_weight"],
+    ).to(device)
+
+    print(f"[GazeRefine] running zero-shot evaluation on {cfg['dataset']}...")
+    metrics = eval_epoch(model, loader, device, threshold=cfg["threshold"], pred_dir=cfg["pred_dir"], save_num=cfg["save_num"])
+
+    print(f"\n{'═' * 60}")
+    print(f"[Results] Dice: {metrics['dice']:.4f} ± {metrics['dice_std']:.4f}")
+    print(f"[Results] IoU:  {metrics['iou']:.4f} ± {metrics['iou_std']:.4f}")
+    print(f"{'═' * 60}\n")
+
+
+if __name__ == "__main__":
+    main()