Spaces:
Sleeping
Sleeping
Upload 37 files
Browse files- .gitattributes +35 -35
- README.md +42 -13
- app.py +644 -0
- gazerefine/__init__.py +32 -0
- gazerefine/__pycache__/__init__.cpython-310.pyc +0 -0
- gazerefine/__pycache__/backbone.cpython-310.pyc +0 -0
- gazerefine/__pycache__/constants.cpython-310.pyc +0 -0
- gazerefine/__pycache__/gaze.cpython-310.pyc +0 -0
- gazerefine/__pycache__/metrics.cpython-310.pyc +0 -0
- gazerefine/__pycache__/model.cpython-310.pyc +0 -0
- gazerefine/__pycache__/visualize.cpython-310.pyc +0 -0
- gazerefine/backbone.py +78 -0
- gazerefine/constants.py +13 -0
- gazerefine/datasets.py +154 -0
- gazerefine/gaze.py +191 -0
- gazerefine/metrics.py +27 -0
- gazerefine/model.py +242 -0
- gazerefine/visualize.py +86 -0
- notebooks/GazeRefine_Demo.ipynb +306 -0
- notebooks/LICENSE +21 -0
- notebooks/configs/kvasir.yaml +27 -0
- notebooks/configs/prostate_mri.yaml +27 -0
- notebooks/gazerefine/__init__.py +33 -0
- notebooks/gazerefine/backbone.py +78 -0
- notebooks/gazerefine/constants.py +13 -0
- notebooks/gazerefine/datasets.py +154 -0
- notebooks/gazerefine/gaze.py +119 -0
- notebooks/gazerefine/metrics.py +27 -0
- notebooks/gazerefine/model.py +242 -0
- notebooks/gazerefine/visualize.py +86 -0
- notebooks/huggingface_space/README.md +42 -0
- notebooks/huggingface_space/app.py +138 -0
- notebooks/huggingface_space/requirements.txt +7 -0
- notebooks/requirements.txt +9 -0
- requirements.txt +10 -0
- scripts/predict_single.py +311 -0
- scripts/run_eval.py +129 -0
.gitattributes
CHANGED
|
@@ -1,35 +1,35 @@
|
|
| 1 |
-
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
-
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
-
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
-
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
-
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
-
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
-
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
-
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
-
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
-
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
-
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
-
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
-
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
-
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
-
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
-
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
-
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
-
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
-
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
-
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
-
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
-
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
-
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
-
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
-
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
-
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
-
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
-
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
-
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
-
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
-
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
-
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
-
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
-
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
-
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
| 1 |
+
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
+
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
+
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
+
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
+
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
+
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
+
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
+
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
+
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
+
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
+
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
+
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
+
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
+
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
+
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
+
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
+
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
+
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
+
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
+
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
+
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
+
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
+
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
+
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
+
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
+
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
+
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
+
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
+
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
+
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
+
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
+
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
README.md
CHANGED
|
@@ -1,13 +1,42 @@
|
|
| 1 |
-
---
|
| 2 |
-
title: GazeRefine
|
| 3 |
-
emoji:
|
| 4 |
-
colorFrom:
|
| 5 |
-
colorTo:
|
| 6 |
-
sdk: gradio
|
| 7 |
-
sdk_version:
|
| 8 |
-
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
---
|
| 12 |
-
|
| 13 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
title: GazeRefine
|
| 3 |
+
emoji: 👁️
|
| 4 |
+
colorFrom: blue
|
| 5 |
+
colorTo: red
|
| 6 |
+
sdk: gradio
|
| 7 |
+
sdk_version: 4.44.1
|
| 8 |
+
app_file: app.py
|
| 9 |
+
pinned: false
|
| 10 |
+
license: mit
|
| 11 |
+
short_description: Zero-shot, training-free gaze-guided medical segmentation
|
| 12 |
+
---
|
| 13 |
+
|
| 14 |
+
# GazeRefine — Expert Gaze as a Test-Time Prompt
|
| 15 |
+
|
| 16 |
+
Interactive demo for **GazeRefine**, a training-free, zero-shot framework
|
| 17 |
+
that turns expert eye-gaze into an inference-time prompt for medical image
|
| 18 |
+
segmentation. Frozen DINOv3 patch features + gaze-anchored
|
| 19 |
+
foreground/background prototypes + recurrent contrastive cleaning + kNN
|
| 20 |
+
affinity propagation — no masks, no clicks-as-boxes, no fine-tuning, no
|
| 21 |
+
adapters, no prompt encoder.
|
| 22 |
+
|
| 23 |
+
## How to use
|
| 24 |
+
1. Upload a colonoscopy or grayscale-MRI-style image.
|
| 25 |
+
2. Click on the image 1–5 times where a clinician's gaze would land on the
|
| 26 |
+
structure of interest (a polyp, the prostate, ...). Each click adds a
|
| 27 |
+
numbered fixation marker; the slider controls that fixation's relative
|
| 28 |
+
duration/weight before your next click.
|
| 29 |
+
3. Pick a hyperparameter preset (tuned per-modality, see the paper).
|
| 30 |
+
4. Press **Run GazeRefine** to get the gaze-prior overlay and the predicted
|
| 31 |
+
segmentation mask.
|
| 32 |
+
|
| 33 |
+
## Notes
|
| 34 |
+
- Inference uses a frozen `vit_large_patch16_dinov3.lvd1689m` backbone from
|
| 35 |
+
`timm`. First run will download the checkpoint.
|
| 36 |
+
- CPU inference works but is slow; a GPU Space is recommended for a smooth
|
| 37 |
+
demo.
|
| 38 |
+
- This Space is for research/demonstration only — it is **not** a clinical
|
| 39 |
+
diagnostic tool.
|
| 40 |
+
|
| 41 |
+
Full code, configs, and the unified evaluation pipeline:
|
| 42 |
+
[GitHub repository](https://github.com/<your-org>/gazerefine).
|
app.py
ADDED
|
@@ -0,0 +1,644 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
app.py — GazeRefine interactive demo (Hugging Face Space).
|
| 3 |
+
|
| 4 |
+
Upload flow
|
| 5 |
+
-----------
|
| 6 |
+
* Standard images (.jpg / .png / etc.) — drag-and-drop or click on the main
|
| 7 |
+
gr.Image widget. The same widget also accepts clicks to place fixations, so
|
| 8 |
+
there is now only ONE image panel instead of two.
|
| 9 |
+
* DICOM files (.dcm) — use the separate "Upload DICOM" file picker. The file
|
| 10 |
+
is decoded with pydicom and converted to an RGB PIL image before being handed
|
| 11 |
+
to the same fixation / run pipeline.
|
| 12 |
+
* Fixation file (.csv / .xlsx / .xls) — use the separate "Upload fixation
|
| 13 |
+
file" picker. This can contain fixations for one or many images (e.g. an
|
| 14 |
+
eye-tracker export with one row per fixation). After upload, three
|
| 15 |
+
dropdowns let the user pick which column is the image-name/ID column and
|
| 16 |
+
which columns hold X / Y (and, optionally, duration). Rows are matched to
|
| 17 |
+
the currently loaded image by filename; X/Y values are auto-detected as
|
| 18 |
+
either normalised [0,1] or raw pixel coordinates.
|
| 19 |
+
"""
|
| 20 |
+
from __future__ import annotations
|
| 21 |
+
|
| 22 |
+
import sys
|
| 23 |
+
import types
|
| 24 |
+
import tempfile
|
| 25 |
+
import csv
|
| 26 |
+
import os
|
| 27 |
+
from pathlib import Path
|
| 28 |
+
|
| 29 |
+
# ── 1. audioop shim (Python 3.13 removed audioop; pydub needs it) ─────────────
|
| 30 |
+
if sys.version_info >= (3, 13):
|
| 31 |
+
for _mod in ("audioop", "pyaudioop"):
|
| 32 |
+
if _mod not in sys.modules:
|
| 33 |
+
sys.modules[_mod] = types.ModuleType(_mod)
|
| 34 |
+
|
| 35 |
+
# ── 2. Patch starlette Jinja2Templates.TemplateResponse ──────────────────────
|
| 36 |
+
import starlette.templating as _st
|
| 37 |
+
|
| 38 |
+
_orig_TR = _st.Jinja2Templates.TemplateResponse
|
| 39 |
+
|
| 40 |
+
def _compat_TR(self, *args, **kwargs):
|
| 41 |
+
if args and isinstance(args[0], str) and len(args) >= 2 and isinstance(args[1], dict):
|
| 42 |
+
name = args[0]
|
| 43 |
+
context = args[1]
|
| 44 |
+
status_code = args[2] if len(args) > 2 else kwargs.get("status_code", 200)
|
| 45 |
+
headers = kwargs.get("headers")
|
| 46 |
+
media_type = kwargs.get("media_type")
|
| 47 |
+
background = kwargs.get("background")
|
| 48 |
+
template = self.get_template(name)
|
| 49 |
+
return _st._TemplateResponse(
|
| 50 |
+
template, context,
|
| 51 |
+
status_code=status_code,
|
| 52 |
+
headers=headers,
|
| 53 |
+
media_type=media_type,
|
| 54 |
+
background=background,
|
| 55 |
+
)
|
| 56 |
+
return _orig_TR(self, *args, **kwargs)
|
| 57 |
+
|
| 58 |
+
_st.Jinja2Templates.TemplateResponse = _compat_TR # type: ignore[method-assign]
|
| 59 |
+
|
| 60 |
+
import gradio as gr
|
| 61 |
+
|
| 62 |
+
# ── 3. gradio_client schema shim ──────────────────────────────────────────────
|
| 63 |
+
try:
|
| 64 |
+
import gradio_client.utils as _gcu
|
| 65 |
+
_orig_inner = _gcu._json_schema_to_python_type
|
| 66 |
+
|
| 67 |
+
def _safe_inner(schema, defs=None):
|
| 68 |
+
if not isinstance(schema, dict):
|
| 69 |
+
return "Any"
|
| 70 |
+
if not isinstance(schema.get("additionalProperties"), dict):
|
| 71 |
+
schema = {k: v for k, v in schema.items() if k != "additionalProperties"}
|
| 72 |
+
return _orig_inner(schema, defs)
|
| 73 |
+
|
| 74 |
+
_gcu._json_schema_to_python_type = _safe_inner
|
| 75 |
+
except Exception:
|
| 76 |
+
pass
|
| 77 |
+
|
| 78 |
+
# ── 4. huggingface_hub HfFolder shim ─────────────────────────────────────────
|
| 79 |
+
try:
|
| 80 |
+
from huggingface_hub import HfFolder # noqa: F401
|
| 81 |
+
except ImportError:
|
| 82 |
+
import huggingface_hub as _hfh
|
| 83 |
+
class _FakeHfFolder:
|
| 84 |
+
@staticmethod
|
| 85 |
+
def get_token(): return None
|
| 86 |
+
_hfh.HfFolder = _FakeHfFolder # type: ignore[attr-defined]
|
| 87 |
+
sys.modules["huggingface_hub"].HfFolder = _FakeHfFolder # type: ignore[assignment]
|
| 88 |
+
|
| 89 |
+
import numpy as np
|
| 90 |
+
from PIL import Image, ImageDraw
|
| 91 |
+
|
| 92 |
+
# ── 5. Path setup ─────────────────────────────────────────────────────────────
|
| 93 |
+
_here = Path(__file__).resolve().parent
|
| 94 |
+
for _candidate in [_here] + list(_here.parents):
|
| 95 |
+
_s = str(_candidate)
|
| 96 |
+
if _s not in sys.path:
|
| 97 |
+
sys.path.insert(0, _s)
|
| 98 |
+
|
| 99 |
+
import scripts.predict_single as _predict_module # noqa: E402
|
| 100 |
+
from scripts.predict_single import predict # noqa: E402
|
| 101 |
+
|
| 102 |
+
# ── Monkey-patch load_fixation_csv ────────────────────────────────────────────
|
| 103 |
+
# Our single-image temp CSV has x,y,duration in raw pixel coordinates with no
|
| 104 |
+
# image_name column. The original loader expects a dataset CSV and returns an
|
| 105 |
+
# empty tensor when that column is absent.
|
| 106 |
+
# This patch detects the missing column, reads the CSV directly, normalises
|
| 107 |
+
# pixel → [0,1], and adds the batch dimension the model requires: (N,3)→(1,N,3).
|
| 108 |
+
|
| 109 |
+
import pandas as _pd
|
| 110 |
+
import torch as _torch
|
| 111 |
+
|
| 112 |
+
try:
|
| 113 |
+
from gazerefine.gaze import load_fixation_csv as _orig_load_fixation_csv
|
| 114 |
+
except Exception:
|
| 115 |
+
_orig_load_fixation_csv = None
|
| 116 |
+
|
| 117 |
+
def _patched_load_fixation_csv(csv_path, image_width=1, image_height=1, image_name=None):
|
| 118 |
+
df = _pd.read_csv(csv_path)
|
| 119 |
+
print(f"[PATCH] load_fixation_csv — columns: {list(df.columns)}, rows: {len(df)}")
|
| 120 |
+
|
| 121 |
+
if "image_name" in df.columns and _orig_load_fixation_csv is not None:
|
| 122 |
+
print("[PATCH] image_name column present — using original loader")
|
| 123 |
+
return _orig_load_fixation_csv(csv_path, image_width=image_width,
|
| 124 |
+
image_height=image_height, image_name=image_name)
|
| 125 |
+
|
| 126 |
+
x = df["x"].values.astype(float)
|
| 127 |
+
y = df["y"].values.astype(float)
|
| 128 |
+
dur = df["duration"].values.astype(float)
|
| 129 |
+
|
| 130 |
+
x_n = x / max(float(image_width), 1.0)
|
| 131 |
+
y_n = y / max(float(image_height), 1.0)
|
| 132 |
+
dur_n = dur / (dur.max() + 1e-8)
|
| 133 |
+
|
| 134 |
+
# model expects (B, N, 3) — add batch dim
|
| 135 |
+
fixations = _torch.tensor(
|
| 136 |
+
list(zip(x_n, y_n, dur_n)), dtype=_torch.float32
|
| 137 |
+
).unsqueeze(0) # (N, 3) → (1, N, 3)
|
| 138 |
+
|
| 139 |
+
print(f"[PATCH] tensor shape: {tuple(fixations.shape)}")
|
| 140 |
+
print(f"[PATCH] fixations (x_norm, y_norm, dur_norm):\n{fixations[0]}")
|
| 141 |
+
return fixations
|
| 142 |
+
|
| 143 |
+
_predict_module.load_fixation_csv = _patched_load_fixation_csv
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
# ─────────────────────────────────────────────────────────────────────────────
|
| 147 |
+
# Helpers
|
| 148 |
+
# ─────────────────────────────────────────────────────────────────────────────
|
| 149 |
+
|
| 150 |
+
PRESETS = {
|
| 151 |
+
"Colonoscopy / polyp (Kvasir-SEG settings)": "colonoscopy",
|
| 152 |
+
"Grayscale MRI / CT (prostate-MRI settings)": "mri",
|
| 153 |
+
}
|
| 154 |
+
POINT_COLORS = ["#ff3b30", "#ff9500", "#ffcc00", "#34c759", "#5ac8fa", "#007aff", "#af52de"]
|
| 155 |
+
|
| 156 |
+
_NO_COL = "— none —"
|
| 157 |
+
|
| 158 |
+
|
| 159 |
+
def dcm_to_pil(dcm_path: str) -> Image.Image:
|
| 160 |
+
"""Load a DICOM file and return an RGB PIL image."""
|
| 161 |
+
import pydicom
|
| 162 |
+
dcm = pydicom.dcmread(dcm_path)
|
| 163 |
+
arr = dcm.pixel_array.astype(np.float32)
|
| 164 |
+
arr = arr - arr.min()
|
| 165 |
+
arr = arr / (arr.max() + 1e-8)
|
| 166 |
+
arr = (arr * 255).astype(np.uint8)
|
| 167 |
+
# Handle multi-frame / greyscale / RGB DICOM
|
| 168 |
+
if arr.ndim == 2:
|
| 169 |
+
return Image.fromarray(arr, mode="L").convert("RGB")
|
| 170 |
+
if arr.ndim == 3 and arr.shape[0] in (1, 3, 4):
|
| 171 |
+
# (C, H, W) → (H, W, C)
|
| 172 |
+
arr = arr.transpose(1, 2, 0)
|
| 173 |
+
return Image.fromarray(arr).convert("RGB")
|
| 174 |
+
|
| 175 |
+
|
| 176 |
+
def draw_points(image: Image.Image, points: list) -> Image.Image:
|
| 177 |
+
"""Overlay fixation circles on a copy of `image`.
|
| 178 |
+
|
| 179 |
+
`points`: list of (x_px, y_px, duration) in original-image pixel coords.
|
| 180 |
+
"""
|
| 181 |
+
if image is None:
|
| 182 |
+
return None
|
| 183 |
+
vis = image.convert("RGB").copy()
|
| 184 |
+
draw = ImageDraw.Draw(vis)
|
| 185 |
+
w, h = vis.size
|
| 186 |
+
r = max(6, min(w, h) // 80)
|
| 187 |
+
for i, (x_px, y_px, dur) in enumerate(points):
|
| 188 |
+
color = POINT_COLORS[i % len(POINT_COLORS)]
|
| 189 |
+
rad = r * (0.6 + 0.8 * dur)
|
| 190 |
+
draw.ellipse(
|
| 191 |
+
[x_px - rad, y_px - rad, x_px + rad, y_px + rad],
|
| 192 |
+
outline=color, width=3,
|
| 193 |
+
)
|
| 194 |
+
draw.text((x_px + rad + 2, y_px - rad), str(i + 1), fill=color)
|
| 195 |
+
return vis
|
| 196 |
+
|
| 197 |
+
|
| 198 |
+
def read_table(path: str) -> "_pd.DataFrame":
|
| 199 |
+
"""Load a .csv / .xlsx / .xls fixation file into a DataFrame."""
|
| 200 |
+
ext = Path(path).suffix.lower()
|
| 201 |
+
if ext in (".xlsx", ".xls"):
|
| 202 |
+
return _pd.read_excel(path)
|
| 203 |
+
# Sniff delimiter for csv/tsv/txt — eye-tracker exports are sometimes
|
| 204 |
+
# tab-separated even with a .csv extension.
|
| 205 |
+
return _pd.read_csv(path, sep=None, engine="python")
|
| 206 |
+
|
| 207 |
+
|
| 208 |
+
def normalize_xy(x_vals: np.ndarray, y_vals: np.ndarray, img_w: int, img_h: int):
|
| 209 |
+
"""Convert X/Y column values to pixel coords for the given image size.
|
| 210 |
+
|
| 211 |
+
Values already in [0, 1] (inclusive, with a little slack for rounding)
|
| 212 |
+
are treated as normalised; anything else is assumed to already be raw
|
| 213 |
+
pixel coordinates and is left as-is (but clamped to the image bounds).
|
| 214 |
+
"""
|
| 215 |
+
looks_normalized = (
|
| 216 |
+
np.nanmax(x_vals) <= 1.05 and np.nanmax(y_vals) <= 1.05
|
| 217 |
+
and np.nanmin(x_vals) >= -0.05 and np.nanmin(y_vals) >= -0.05
|
| 218 |
+
)
|
| 219 |
+
if looks_normalized:
|
| 220 |
+
x_px = np.clip(x_vals, 0, 1) * img_w
|
| 221 |
+
y_px = np.clip(y_vals, 0, 1) * img_h
|
| 222 |
+
else:
|
| 223 |
+
x_px = np.clip(x_vals, 0, img_w)
|
| 224 |
+
y_px = np.clip(y_vals, 0, img_h)
|
| 225 |
+
return x_px, y_px
|
| 226 |
+
|
| 227 |
+
|
| 228 |
+
# ─────────────────────────────────────────────────────────────────────────────
|
| 229 |
+
# Event handlers
|
| 230 |
+
# ─────────────────────────────────────────────────────────────────────────────
|
| 231 |
+
|
| 232 |
+
|
| 233 |
+
_UPLOAD_LABEL = "Drop / click to load .jpg .png .bmp .tif .tiff .webp .dcm"
|
| 234 |
+
_FIXATION_LABEL = "Click to place fixations"
|
| 235 |
+
_FIXFILE_LABEL = "Upload fixation file (.csv / .xlsx / .xls) — optional"
|
| 236 |
+
|
| 237 |
+
|
| 238 |
+
def _resolve_path(file_obj):
|
| 239 |
+
"""Extract a filesystem path from whatever gr.File passes."""
|
| 240 |
+
if isinstance(file_obj, str):
|
| 241 |
+
return file_obj
|
| 242 |
+
if isinstance(file_obj, dict):
|
| 243 |
+
return file_obj.get("name") or file_obj.get("path") or file_obj.get("tmp_path") or ""
|
| 244 |
+
if hasattr(file_obj, "name"):
|
| 245 |
+
return file_obj.name
|
| 246 |
+
return ""
|
| 247 |
+
|
| 248 |
+
|
| 249 |
+
def on_file_upload(file_obj):
|
| 250 |
+
"""Load any image or DICOM and switch the panel to fixation-click mode."""
|
| 251 |
+
_no_change = (None, [], gr.update(), gr.update(), gr.update(), gr.update())
|
| 252 |
+
|
| 253 |
+
if file_obj is None:
|
| 254 |
+
return _no_change
|
| 255 |
+
|
| 256 |
+
image_name = ""
|
| 257 |
+
|
| 258 |
+
# gr.Image gives PIL/numpy; gr.File gives a path
|
| 259 |
+
if isinstance(file_obj, Image.Image):
|
| 260 |
+
pil = file_obj.convert("RGB")
|
| 261 |
+
elif isinstance(file_obj, np.ndarray):
|
| 262 |
+
pil = Image.fromarray(file_obj).convert("RGB")
|
| 263 |
+
else:
|
| 264 |
+
path = _resolve_path(file_obj)
|
| 265 |
+
if not path:
|
| 266 |
+
gr.Warning("Could not resolve file path.")
|
| 267 |
+
return _no_change
|
| 268 |
+
image_name = Path(path).name
|
| 269 |
+
ext = Path(path).suffix.lower()
|
| 270 |
+
try:
|
| 271 |
+
pil = dcm_to_pil(path) if ext == ".dcm" else Image.open(path).convert("RGB")
|
| 272 |
+
except Exception as e:
|
| 273 |
+
gr.Warning(f"Could not load file: {e}")
|
| 274 |
+
return _no_change
|
| 275 |
+
|
| 276 |
+
print(f"[DEBUG] on_file_upload — size={pil.size} name={image_name!r}")
|
| 277 |
+
# Switch: hide upload zone, show image panel + delete button
|
| 278 |
+
return (
|
| 279 |
+
pil, # orig_image_state
|
| 280 |
+
[], # points_state
|
| 281 |
+
image_name, # image_name_state
|
| 282 |
+
gr.update(visible=False), # upload_zone → hide
|
| 283 |
+
gr.update(value=pil, visible=True,
|
| 284 |
+
label=_FIXATION_LABEL), # image_panel → show with image
|
| 285 |
+
gr.update(visible=True), # delete_btn → show
|
| 286 |
+
)
|
| 287 |
+
|
| 288 |
+
|
| 289 |
+
def on_select(orig_image: Image.Image, points: list, duration: float, evt: gr.SelectData):
|
| 290 |
+
"""Record a fixation click in original-image pixel coords."""
|
| 291 |
+
if orig_image is None:
|
| 292 |
+
gr.Warning("Upload an image first.")
|
| 293 |
+
return points, gr.update()
|
| 294 |
+
x_px, y_px = float(evt.index[0]), float(evt.index[1])
|
| 295 |
+
new_points = points + [(x_px, y_px, duration)]
|
| 296 |
+
print(f"[DEBUG] fixation #{len(new_points)}: x={x_px:.1f} y={y_px:.1f} dur={duration}")
|
| 297 |
+
return new_points, draw_points(orig_image, new_points)
|
| 298 |
+
|
| 299 |
+
|
| 300 |
+
def on_clear(orig_image):
|
| 301 |
+
"""Remove all fixations but keep the current image."""
|
| 302 |
+
if orig_image is None:
|
| 303 |
+
return [], gr.update()
|
| 304 |
+
return [], gr.update(value=orig_image)
|
| 305 |
+
|
| 306 |
+
|
| 307 |
+
def on_delete():
|
| 308 |
+
"""Delete the current image and return to upload mode."""
|
| 309 |
+
return (
|
| 310 |
+
None, # orig_image_state
|
| 311 |
+
[], # points_state
|
| 312 |
+
"", # image_name_state
|
| 313 |
+
gr.update(value=None, visible=True), # upload_zone → show (reset)
|
| 314 |
+
gr.update(value=None, visible=False), # image_panel → hide
|
| 315 |
+
gr.update(visible=False), # delete_btn → hide
|
| 316 |
+
)
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
# ── Fixation-file upload → column mapping ────────────────────────────────────
|
| 320 |
+
|
| 321 |
+
def on_fixfile_upload(file_obj):
|
| 322 |
+
"""Load the fixation table and populate the column-mapping dropdowns."""
|
| 323 |
+
_hide = (
|
| 324 |
+
None, gr.update(visible=False),
|
| 325 |
+
gr.update(choices=[], value=None), gr.update(choices=[], value=None),
|
| 326 |
+
gr.update(choices=[], value=None), gr.update(choices=[], value=None),
|
| 327 |
+
gr.update(visible=False),
|
| 328 |
+
)
|
| 329 |
+
if file_obj is None:
|
| 330 |
+
return _hide
|
| 331 |
+
|
| 332 |
+
path = _resolve_path(file_obj)
|
| 333 |
+
if not path:
|
| 334 |
+
gr.Warning("Could not resolve fixation file path.")
|
| 335 |
+
return _hide
|
| 336 |
+
|
| 337 |
+
try:
|
| 338 |
+
df = read_table(path)
|
| 339 |
+
except Exception as e:
|
| 340 |
+
gr.Warning(f"Could not read fixation file: {e}")
|
| 341 |
+
return _hide
|
| 342 |
+
|
| 343 |
+
if df.empty or len(df.columns) == 0:
|
| 344 |
+
gr.Warning("Fixation file appears to be empty.")
|
| 345 |
+
return _hide
|
| 346 |
+
|
| 347 |
+
cols = [str(c) for c in df.columns]
|
| 348 |
+
print(f"[DEBUG] fixation file loaded — columns: {cols}, rows: {len(df)}")
|
| 349 |
+
|
| 350 |
+
def _guess(*keywords, fallback=None):
|
| 351 |
+
for c in cols:
|
| 352 |
+
cl = c.lower()
|
| 353 |
+
if any(k in cl for k in keywords):
|
| 354 |
+
return c
|
| 355 |
+
return fallback if fallback is not None else cols[0]
|
| 356 |
+
|
| 357 |
+
guess_id = _guess("image", "id", "name", "file", fallback=cols[0])
|
| 358 |
+
# exact / boundary-aware matches first (avoids "fix_index" matching "x"),
|
| 359 |
+
# then fall back to a bare trailing "x" / "y".
|
| 360 |
+
guess_x = _guess("fix_x", "pos_x", "gaze_x", fallback=None)
|
| 361 |
+
if guess_x is None:
|
| 362 |
+
guess_x = next((c for c in cols if c.lower().rstrip("_") .endswith("x")
|
| 363 |
+
and "index" not in c.lower()), cols[0])
|
| 364 |
+
guess_y = _guess("fix_y", "pos_y", "gaze_y", fallback=None)
|
| 365 |
+
if guess_y is None:
|
| 366 |
+
guess_y = next((c for c in cols if c.lower().rstrip("_").endswith("y")
|
| 367 |
+
and "index" not in c.lower()), cols[0])
|
| 368 |
+
|
| 369 |
+
dur_choices = [_NO_COL] + cols
|
| 370 |
+
guess_dur = _guess("duration", "dur", fallback=_NO_COL)
|
| 371 |
+
|
| 372 |
+
return (
|
| 373 |
+
df.to_json(), # fixfile_df_state (serialized)
|
| 374 |
+
gr.update(visible=True), # mapping_row → show
|
| 375 |
+
gr.update(choices=cols, value=guess_id), # id_col_dd
|
| 376 |
+
gr.update(choices=cols, value=guess_x), # x_col_dd
|
| 377 |
+
gr.update(choices=cols, value=guess_y), # y_col_dd
|
| 378 |
+
gr.update(choices=dur_choices, value=guess_dur), # dur_col_dd
|
| 379 |
+
gr.update(visible=True), # apply_fix_btn → show
|
| 380 |
+
)
|
| 381 |
+
|
| 382 |
+
|
| 383 |
+
def on_apply_fixfile(fixfile_json, id_col, x_col, y_col, dur_col,
|
| 384 |
+
orig_image, image_name):
|
| 385 |
+
"""Match rows to the currently loaded image (by filename) and load
|
| 386 |
+
them as fixation points, replacing whatever points are currently set.
|
| 387 |
+
|
| 388 |
+
If no rows match the loaded image's filename, nothing is loaded — the
|
| 389 |
+
existing points (if any) are left untouched, and the user is warned so
|
| 390 |
+
they can check the ID column / image filename instead of silently
|
| 391 |
+
getting fixations for the wrong image."""
|
| 392 |
+
if orig_image is None:
|
| 393 |
+
gr.Warning("Load an image first, then apply the fixation file.")
|
| 394 |
+
return gr.update(), gr.update()
|
| 395 |
+
if not fixfile_json:
|
| 396 |
+
gr.Warning("Upload a fixation file first.")
|
| 397 |
+
return gr.update(), gr.update()
|
| 398 |
+
if not id_col or not x_col or not y_col:
|
| 399 |
+
gr.Warning("Pick the ID, X and Y columns first.")
|
| 400 |
+
return gr.update(), gr.update()
|
| 401 |
+
if not image_name:
|
| 402 |
+
gr.Warning(
|
| 403 |
+
"Couldn't determine the loaded image's filename (this can "
|
| 404 |
+
"happen if the image was pasted/dropped without a filename). "
|
| 405 |
+
"Re-upload the image as a file and try again."
|
| 406 |
+
)
|
| 407 |
+
return gr.update(), gr.update()
|
| 408 |
+
|
| 409 |
+
df = _pd.read_json(fixfile_json)
|
| 410 |
+
|
| 411 |
+
# Match by exact filename first, then by stem-without-extension, so the
|
| 412 |
+
# mapping still works if the fixation file's IMAGE column omits the
|
| 413 |
+
# extension or uses a different one than the uploaded image.
|
| 414 |
+
mask = df[id_col].astype(str) == image_name
|
| 415 |
+
if not mask.any():
|
| 416 |
+
stem_no_ext = Path(image_name).stem
|
| 417 |
+
mask = df[id_col].astype(str).apply(lambda v: Path(str(v)).stem) == stem_no_ext
|
| 418 |
+
|
| 419 |
+
sub = df[mask]
|
| 420 |
+
|
| 421 |
+
if sub.empty:
|
| 422 |
+
gr.Warning(
|
| 423 |
+
f"No rows in the fixation file match the loaded image "
|
| 424 |
+
f"('{image_name}'). Nothing was loaded — check that the ID "
|
| 425 |
+
f"column values match the image filename."
|
| 426 |
+
)
|
| 427 |
+
return gr.update(), gr.update()
|
| 428 |
+
|
| 429 |
+
w, h = orig_image.size
|
| 430 |
+
x_vals = sub[x_col].astype(float).to_numpy()
|
| 431 |
+
y_vals = sub[y_col].astype(float).to_numpy()
|
| 432 |
+
x_px, y_px = normalize_xy(x_vals, y_vals, w, h)
|
| 433 |
+
|
| 434 |
+
if dur_col and dur_col != _NO_COL and dur_col in sub.columns:
|
| 435 |
+
dur_raw = sub[dur_col].astype(float).to_numpy()
|
| 436 |
+
dmax = float(np.nanmax(dur_raw)) if len(dur_raw) else 1.0
|
| 437 |
+
dur_n = dur_raw / (dmax + 1e-8)
|
| 438 |
+
else:
|
| 439 |
+
dur_n = np.full(len(sub), 1.0)
|
| 440 |
+
|
| 441 |
+
new_points = [
|
| 442 |
+
(float(xp), float(yp), float(d))
|
| 443 |
+
for xp, yp, d in zip(x_px, y_px, dur_n)
|
| 444 |
+
]
|
| 445 |
+
print(f"[DEBUG] loaded {len(new_points)} fixations from file for image '{image_name}'")
|
| 446 |
+
|
| 447 |
+
return new_points, draw_points(orig_image, new_points)
|
| 448 |
+
|
| 449 |
+
|
| 450 |
+
def run(orig_image: Image.Image, points: list, preset_name: str, threshold: float):
|
| 451 |
+
import traceback, uuid
|
| 452 |
+
print(f"[DEBUG] run — points={len(points)} preset={preset_name}")
|
| 453 |
+
|
| 454 |
+
if orig_image is None:
|
| 455 |
+
gr.Warning("Upload an image first.")
|
| 456 |
+
return None, None, None
|
| 457 |
+
if not points:
|
| 458 |
+
gr.Warning("Click on the image at least once to place a fixation (or load a fixation file).")
|
| 459 |
+
return None, None, None
|
| 460 |
+
|
| 461 |
+
preset_key = PRESETS[preset_name]
|
| 462 |
+
w, h = orig_image.size
|
| 463 |
+
|
| 464 |
+
shared_stem = f"gazerefine_{uuid.uuid4().hex}"
|
| 465 |
+
tmp_img_path = os.path.join(tempfile.gettempdir(), f"{shared_stem}.png")
|
| 466 |
+
fixation_csv_path = os.path.join(tempfile.gettempdir(), f"{shared_stem}.csv")
|
| 467 |
+
|
| 468 |
+
orig_image.convert("RGB").save(tmp_img_path)
|
| 469 |
+
|
| 470 |
+
with open(fixation_csv_path, "w", newline="") as f:
|
| 471 |
+
writer = csv.writer(f)
|
| 472 |
+
writer.writerow(["x", "y", "duration"])
|
| 473 |
+
for x_px, y_px, dur in points:
|
| 474 |
+
writer.writerow([x_px, y_px, dur])
|
| 475 |
+
|
| 476 |
+
print(f"[DEBUG] image {w}x{h} | {len(points)} fixations | preset={preset_key} thr={threshold}")
|
| 477 |
+
|
| 478 |
+
with open(fixation_csv_path) as f:
|
| 479 |
+
print(f"[DEBUG] CSV:\n{f.read()}")
|
| 480 |
+
|
| 481 |
+
try:
|
| 482 |
+
out = predict(
|
| 483 |
+
image_path=tmp_img_path,
|
| 484 |
+
fixation_csv=fixation_csv_path,
|
| 485 |
+
preset=preset_key,
|
| 486 |
+
threshold=threshold,
|
| 487 |
+
return_all=True,
|
| 488 |
+
)
|
| 489 |
+
except Exception as e:
|
| 490 |
+
print(f"[ERROR] predict() raised: {e}")
|
| 491 |
+
traceback.print_exc()
|
| 492 |
+
gr.Warning(f"Prediction failed: {e}")
|
| 493 |
+
return None, None, None
|
| 494 |
+
finally:
|
| 495 |
+
for p in (tmp_img_path, fixation_csv_path):
|
| 496 |
+
try:
|
| 497 |
+
os.unlink(p)
|
| 498 |
+
except OSError:
|
| 499 |
+
pass
|
| 500 |
+
|
| 501 |
+
mask_arr = np.array(out["mask"])
|
| 502 |
+
print(f"[DEBUG] mask non-zero: {(mask_arr > 0).sum()} / {mask_arr.size}")
|
| 503 |
+
return out["gaze_overlay"], out["mask_overlay"], out["mask"]
|
| 504 |
+
|
| 505 |
+
|
| 506 |
+
# ─────────────────────────────────────────────────────────────────────────────
|
| 507 |
+
# UI
|
| 508 |
+
# ─────────────────────────────────────────────────────────────────────────────
|
| 509 |
+
|
| 510 |
+
with gr.Blocks(title="GazeRefine — gaze-guided zero-shot segmentation") as demo:
|
| 511 |
+
gr.Markdown(
|
| 512 |
+
"""
|
| 513 |
+
# 👁️ GazeRefine — Expert Gaze as a Test-Time Prompt
|
| 514 |
+
Training-free, zero-shot medical image segmentation.
|
| 515 |
+
Upload an image or DICOM, click to place fixations (or load a fixation
|
| 516 |
+
file), then hit **Run**.
|
| 517 |
+
"""
|
| 518 |
+
)
|
| 519 |
+
|
| 520 |
+
orig_image_state = gr.State(None)
|
| 521 |
+
points_state = gr.State([])
|
| 522 |
+
image_name_state = gr.State("") # filename of the currently loaded image
|
| 523 |
+
fixfile_df_state = gr.State(None) # serialized DataFrame (to_json) of the uploaded fixation file
|
| 524 |
+
|
| 525 |
+
with gr.Row():
|
| 526 |
+
# ── Left column ───────────────────────────────────────────────────────
|
| 527 |
+
with gr.Column(scale=1):
|
| 528 |
+
|
| 529 |
+
# ── Upload zone (visible when no image loaded) ────────────────────
|
| 530 |
+
upload_zone = gr.File(
|
| 531 |
+
label=_UPLOAD_LABEL,
|
| 532 |
+
file_types=[".jpg", ".jpeg", ".png", ".bmp",
|
| 533 |
+
".tif", ".tiff", ".webp", ".gif", ".dcm"],
|
| 534 |
+
file_count="single",
|
| 535 |
+
visible=True,
|
| 536 |
+
elem_id="upload_zone",
|
| 537 |
+
)
|
| 538 |
+
|
| 539 |
+
# ── Image panel (hidden until image loaded; click to fixate) ──────
|
| 540 |
+
image_panel = gr.Image(
|
| 541 |
+
type="pil",
|
| 542 |
+
label=_FIXATION_LABEL,
|
| 543 |
+
height=430,
|
| 544 |
+
interactive=False, # no toolbar → .select fires on click
|
| 545 |
+
show_download_button=False,
|
| 546 |
+
visible=False,
|
| 547 |
+
elem_id="image_panel",
|
| 548 |
+
)
|
| 549 |
+
|
| 550 |
+
# ── Delete button (hidden until image loaded) ─────────────────────
|
| 551 |
+
delete_btn = gr.Button("🗑 Delete image — load another", visible=False, variant="secondary")
|
| 552 |
+
|
| 553 |
+
# ── Fixation file upload (optional alternative to manual clicks) ──
|
| 554 |
+
with gr.Accordion("📄 Load fixations from file", open=False):
|
| 555 |
+
fixfile_upload = gr.File(
|
| 556 |
+
label=_FIXFILE_LABEL,
|
| 557 |
+
file_types=[".csv", ".xlsx", ".xls", ".tsv", ".txt"],
|
| 558 |
+
file_count="single",
|
| 559 |
+
elem_id="fixfile_upload",
|
| 560 |
+
)
|
| 561 |
+
with gr.Row(visible=False) as mapping_row:
|
| 562 |
+
id_col_dd = gr.Dropdown(label="Image / ID column", choices=[])
|
| 563 |
+
x_col_dd = gr.Dropdown(label="X column", choices=[])
|
| 564 |
+
y_col_dd = gr.Dropdown(label="Y column", choices=[])
|
| 565 |
+
dur_col_dd = gr.Dropdown(label="Duration column (optional)", choices=[])
|
| 566 |
+
apply_fix_btn = gr.Button(
|
| 567 |
+
"📥 Load fixations for current image", visible=False,
|
| 568 |
+
)
|
| 569 |
+
|
| 570 |
+
# ── Controls ──────────────────────────────────────────────────────
|
| 571 |
+
with gr.Row():
|
| 572 |
+
duration_slider = gr.Slider(
|
| 573 |
+
0.1, 1.0, value=1.0, step=0.1,
|
| 574 |
+
label="Fixation duration weight",
|
| 575 |
+
)
|
| 576 |
+
clear_btn = gr.Button("✖ Clear fixations")
|
| 577 |
+
|
| 578 |
+
preset = gr.Radio(
|
| 579 |
+
list(PRESETS.keys()), value=list(PRESETS.keys())[0],
|
| 580 |
+
label="Preset",
|
| 581 |
+
)
|
| 582 |
+
threshold = gr.Slider(
|
| 583 |
+
0.1, 0.9, value=0.5, step=0.05,
|
| 584 |
+
label="Mask threshold",
|
| 585 |
+
)
|
| 586 |
+
run_btn = gr.Button("▶ Run GazeRefine", variant="primary")
|
| 587 |
+
|
| 588 |
+
# ── Right column: outputs ─────────────────────────────────────────────
|
| 589 |
+
with gr.Column(scale=1):
|
| 590 |
+
gaze_out = gr.Image(label="Gaze prior", height=260)
|
| 591 |
+
with gr.Row():
|
| 592 |
+
mask_overlay_out = gr.Image(label="Mask overlay", height=260)
|
| 593 |
+
mask_only_out = gr.Image(label="Binary mask", height=260)
|
| 594 |
+
|
| 595 |
+
# ── Event wiring ──────────────────────────────────────────────────────────
|
| 596 |
+
|
| 597 |
+
_upload_outputs = [orig_image_state, points_state, image_name_state,
|
| 598 |
+
upload_zone, image_panel, delete_btn]
|
| 599 |
+
|
| 600 |
+
upload_zone.upload(on_file_upload, inputs=[upload_zone], outputs=_upload_outputs)
|
| 601 |
+
upload_zone.change(on_file_upload, inputs=[upload_zone], outputs=_upload_outputs)
|
| 602 |
+
|
| 603 |
+
image_panel.select(
|
| 604 |
+
on_select,
|
| 605 |
+
inputs=[orig_image_state, points_state, duration_slider],
|
| 606 |
+
outputs=[points_state, image_panel],
|
| 607 |
+
)
|
| 608 |
+
|
| 609 |
+
clear_btn.click(
|
| 610 |
+
on_clear,
|
| 611 |
+
inputs=[orig_image_state],
|
| 612 |
+
outputs=[points_state, image_panel],
|
| 613 |
+
)
|
| 614 |
+
|
| 615 |
+
delete_btn.click(
|
| 616 |
+
on_delete,
|
| 617 |
+
outputs=[orig_image_state, points_state, image_name_state,
|
| 618 |
+
upload_zone, image_panel, delete_btn],
|
| 619 |
+
)
|
| 620 |
+
|
| 621 |
+
_fixfile_outputs = [fixfile_df_state, mapping_row, id_col_dd, x_col_dd, y_col_dd, dur_col_dd, apply_fix_btn]
|
| 622 |
+
|
| 623 |
+
fixfile_upload.upload(on_fixfile_upload, inputs=[fixfile_upload], outputs=_fixfile_outputs)
|
| 624 |
+
fixfile_upload.change(on_fixfile_upload, inputs=[fixfile_upload], outputs=_fixfile_outputs)
|
| 625 |
+
|
| 626 |
+
apply_fix_btn.click(
|
| 627 |
+
on_apply_fixfile,
|
| 628 |
+
inputs=[fixfile_df_state, id_col_dd, x_col_dd, y_col_dd, dur_col_dd,
|
| 629 |
+
orig_image_state, image_name_state],
|
| 630 |
+
outputs=[points_state, image_panel],
|
| 631 |
+
)
|
| 632 |
+
|
| 633 |
+
run_btn.click(
|
| 634 |
+
run,
|
| 635 |
+
inputs=[orig_image_state, points_state, preset, threshold],
|
| 636 |
+
outputs=[gaze_out, mask_overlay_out, mask_only_out],
|
| 637 |
+
)
|
| 638 |
+
|
| 639 |
+
gr.Markdown(
|
| 640 |
+
"Method: GazeRefine — frozen DINOv3 + gaze-anchored prototypes + recurrent "
|
| 641 |
+
"foreground/background refinement, entirely training-free."
|
| 642 |
+
)
|
| 643 |
+
|
| 644 |
+
demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
|
gazerefine/__init__.py
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
GazeRefine
|
| 3 |
+
==========
|
| 4 |
+
Expert gaze as a test-time prompt for training-free medical image segmentation.
|
| 5 |
+
|
| 6 |
+
Quick start
|
| 7 |
+
-----------
|
| 8 |
+
from scripts.predict_single import predict
|
| 9 |
+
mask = predict(image_path="image.png", fixation_csv="fixations.csv")
|
| 10 |
+
mask.save("output_mask.png")
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
from .backbone import FrozenDINOv3
|
| 14 |
+
from .gaze import get_scanpath, load_fixation_csv, generate_gaze_heatmap
|
| 15 |
+
from .model import GazeRefine, knn_affinity_refinement
|
| 16 |
+
from .metrics import compute_metrics
|
| 17 |
+
from .visualize import save_prediction, overlay_heatmap, overlay_mask
|
| 18 |
+
|
| 19 |
+
__all__ = [
|
| 20 |
+
"FrozenDINOv3",
|
| 21 |
+
"get_scanpath",
|
| 22 |
+
"load_fixation_csv",
|
| 23 |
+
"generate_gaze_heatmap",
|
| 24 |
+
"GazeRefine",
|
| 25 |
+
"knn_affinity_refinement",
|
| 26 |
+
"compute_metrics",
|
| 27 |
+
"save_prediction",
|
| 28 |
+
"overlay_heatmap",
|
| 29 |
+
"overlay_mask",
|
| 30 |
+
]
|
| 31 |
+
|
| 32 |
+
__version__ = "0.1.0"
|
gazerefine/__pycache__/__init__.cpython-310.pyc
ADDED
|
Binary file (916 Bytes). View file
|
|
|
gazerefine/__pycache__/backbone.cpython-310.pyc
ADDED
|
Binary file (3.51 kB). View file
|
|
|
gazerefine/__pycache__/constants.cpython-310.pyc
ADDED
|
Binary file (547 Bytes). View file
|
|
|
gazerefine/__pycache__/gaze.cpython-310.pyc
ADDED
|
Binary file (5.59 kB). View file
|
|
|
gazerefine/__pycache__/metrics.cpython-310.pyc
ADDED
|
Binary file (961 Bytes). View file
|
|
|
gazerefine/__pycache__/model.cpython-310.pyc
ADDED
|
Binary file (9.49 kB). View file
|
|
|
gazerefine/__pycache__/visualize.cpython-310.pyc
ADDED
|
Binary file (3.78 kB). View file
|
|
|
gazerefine/backbone.py
ADDED
|
@@ -0,0 +1,78 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
backbone.py — Frozen DINOv3 feature extractor.
|
| 3 |
+
|
| 4 |
+
GazeRefine never updates the backbone. We load a pretrained DINOv3 ViT via
|
| 5 |
+
`timm`, freeze every parameter, and pull out raw patch tokens from one or
|
| 6 |
+
more transformer blocks using forward hooks. No adapters, no projections,
|
| 7 |
+
no fine-tuning.
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
import torch
|
| 13 |
+
import torch.nn as nn
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
class FrozenDINOv3(nn.Module):
|
| 17 |
+
"""Completely frozen DINOv3 ViT backbone that exposes raw patch tokens.
|
| 18 |
+
|
| 19 |
+
Parameters
|
| 20 |
+
----------
|
| 21 |
+
model_name : str
|
| 22 |
+
Any DINOv3 variant available in `timm` (e.g.
|
| 23 |
+
``"vit_large_patch16_dinov3.lvd1689m"``,
|
| 24 |
+
``"vit_base_patch16_dinov3.lvd1689m"``). Larger backbones generally
|
| 25 |
+
give cleaner semantic separation but cost more memory/compute.
|
| 26 |
+
extract_mode : {"last", "all"}
|
| 27 |
+
- ``"last"``: use only the final block's patch tokens (fast, the
|
| 28 |
+
default used in our reported results).
|
| 29 |
+
- ``"all"``: pool patch tokens from 4 evenly-spaced blocks
|
| 30 |
+
(0, n/4, 3n/4, n-1) and average the resulting similarity maps in
|
| 31 |
+
``GazeRefine``. This sometimes helps on harder modalities at the
|
| 32 |
+
cost of ~4x compute.
|
| 33 |
+
"""
|
| 34 |
+
|
| 35 |
+
def __init__(
|
| 36 |
+
self,
|
| 37 |
+
model_name: str = "vit_large_patch16_dinov3.lvd1689m",
|
| 38 |
+
extract_mode: str = "last",
|
| 39 |
+
):
|
| 40 |
+
super().__init__()
|
| 41 |
+
import timm # local import: keeps `timm` optional for users who only read the code
|
| 42 |
+
|
| 43 |
+
print(f"[GazeRefine] Loading frozen backbone: {model_name} (extract_mode={extract_mode})")
|
| 44 |
+
bb = timm.create_model(model_name, pretrained=True, num_classes=0)
|
| 45 |
+
for p in bb.parameters():
|
| 46 |
+
p.requires_grad_(False)
|
| 47 |
+
bb.eval()
|
| 48 |
+
|
| 49 |
+
self.backbone = bb
|
| 50 |
+
self.embed_dim = bb.embed_dim
|
| 51 |
+
self.num_blocks = len(bb.blocks)
|
| 52 |
+
|
| 53 |
+
if extract_mode == "last":
|
| 54 |
+
self.levels = [-1]
|
| 55 |
+
elif extract_mode == "all":
|
| 56 |
+
self.levels = [0, self.num_blocks // 4, (self.num_blocks * 3) // 4, self.num_blocks - 1]
|
| 57 |
+
else:
|
| 58 |
+
raise ValueError(f"extract_mode must be 'last' or 'all', got {extract_mode!r}")
|
| 59 |
+
|
| 60 |
+
print(f"[GazeRefine] Hooking transformer blocks at levels: {self.levels}")
|
| 61 |
+
self._feats: dict[int, torch.Tensor] = {}
|
| 62 |
+
for lvl in self.levels:
|
| 63 |
+
bb.blocks[lvl].register_forward_hook(self._make_hook(lvl))
|
| 64 |
+
|
| 65 |
+
def _make_hook(self, lvl: int):
|
| 66 |
+
def _hook_fn(module, inp, out):
|
| 67 |
+
# DINOv3 token layout: [CLS, register_1..register_k, patch_1..patch_N]
|
| 68 |
+
n_prefix = 1 + getattr(self.backbone, "num_register_tokens", 4)
|
| 69 |
+
self._feats[lvl] = out[:, n_prefix:, :] # (B, N, D) — patch tokens only
|
| 70 |
+
return _hook_fn
|
| 71 |
+
|
| 72 |
+
@torch.no_grad()
|
| 73 |
+
def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
|
| 74 |
+
"""Run the frozen backbone and return a list of (B, N, D) patch-token tensors,
|
| 75 |
+
one per hooked level, in the order given by ``self.levels``."""
|
| 76 |
+
self._feats.clear()
|
| 77 |
+
self.backbone(x)
|
| 78 |
+
return [self._feats[lvl] for lvl in self.levels]
|
gazerefine/constants.py
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""constants.py — shared sizing and normalization constants.
|
| 2 |
+
|
| 3 |
+
IMG_SIZE is fixed to an exact multiple of the ViT patch size so the patch
|
| 4 |
+
grid divides evenly with no rounding/cropping artifacts.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
PATCH_SIZE = 16
|
| 8 |
+
H_PATCH = 37 * 2 # 74 patches per side
|
| 9 |
+
IMG_SIZE = H_PATCH * PATCH_SIZE # 1184 px (use 518 = 37*14 if you switch to a /14 ViT)
|
| 10 |
+
N_PATCHES = H_PATCH ** 2
|
| 11 |
+
|
| 12 |
+
IMG_MEAN = [0.485, 0.456, 0.406]
|
| 13 |
+
IMG_STD = [0.229, 0.224, 0.225]
|
gazerefine/datasets.py
ADDED
|
@@ -0,0 +1,154 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
datasets.py — Image/mask/fixation loaders for the two reported benchmarks,
|
| 3 |
+
plus a small base class so adding a third dataset is just a few lines.
|
| 4 |
+
|
| 5 |
+
Expected folder layout (same for both datasets):
|
| 6 |
+
|
| 7 |
+
<root>/
|
| 8 |
+
images/ one image per case (.jpg/.png for Kvasir, .dcm for prostate MRI)
|
| 9 |
+
masks/ matching binary mask, same basename, .png
|
| 10 |
+
<fixation_csv>
|
| 11 |
+
one row per fixation, see gazerefine.gaze for the expected columns.
|
| 12 |
+
The IMAGE column must match an images/ filename (Kvasir) or the
|
| 13 |
+
DICOM basename + ".jpg" (prostate — fixation collection was run on
|
| 14 |
+
JPEG-rendered slices while the model reads the original DICOM).
|
| 15 |
+
"""
|
| 16 |
+
|
| 17 |
+
from __future__ import annotations
|
| 18 |
+
|
| 19 |
+
import os
|
| 20 |
+
|
| 21 |
+
import numpy as np
|
| 22 |
+
import pandas as pd
|
| 23 |
+
import torch
|
| 24 |
+
from PIL import Image
|
| 25 |
+
from torch.utils.data import Dataset
|
| 26 |
+
import torchvision.transforms as T
|
| 27 |
+
|
| 28 |
+
from .constants import IMG_SIZE, IMG_MEAN, IMG_STD
|
| 29 |
+
from .gaze import get_scanpath, IMG_COL
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
class _BaseGazeDataset(Dataset):
|
| 33 |
+
"""Shared image/mask transform + fixation-grouping logic.
|
| 34 |
+
|
| 35 |
+
Subclasses only need to implement ``_load_image(name)`` and provide the
|
| 36 |
+
set of valid image ids that have a matching mask and fixation entries.
|
| 37 |
+
"""
|
| 38 |
+
|
| 39 |
+
def __init__(self, root: str, fixation_csv: str, img_size: int = IMG_SIZE):
|
| 40 |
+
self.root = root
|
| 41 |
+
self.img_dir = os.path.join(root, "images")
|
| 42 |
+
self.mask_dir = os.path.join(root, "masks")
|
| 43 |
+
self.img_size = img_size
|
| 44 |
+
|
| 45 |
+
self.df = pd.read_csv(fixation_csv)
|
| 46 |
+
self.df.columns = self.df.columns.str.strip()
|
| 47 |
+
self.fix_df = self.df.groupby(IMG_COL)
|
| 48 |
+
self.max_len = int(self.df.groupby(IMG_COL).size().max())
|
| 49 |
+
print(f"[Dataset] max scanpath length = {self.max_len}")
|
| 50 |
+
|
| 51 |
+
self.image_ids: list[str] = [] # set by subclass __init__
|
| 52 |
+
|
| 53 |
+
self.img_tf = T.Compose([
|
| 54 |
+
T.Resize((img_size, img_size)),
|
| 55 |
+
T.ToTensor(),
|
| 56 |
+
T.Normalize(IMG_MEAN, IMG_STD),
|
| 57 |
+
])
|
| 58 |
+
self.mask_tf = T.Compose([
|
| 59 |
+
T.Resize((img_size, img_size), interpolation=T.InterpolationMode.NEAREST),
|
| 60 |
+
T.ToTensor(),
|
| 61 |
+
])
|
| 62 |
+
|
| 63 |
+
def __len__(self):
|
| 64 |
+
return len(self.image_ids)
|
| 65 |
+
|
| 66 |
+
def _load_image(self, name: str) -> Image.Image:
|
| 67 |
+
raise NotImplementedError
|
| 68 |
+
|
| 69 |
+
def _mask_path(self, name: str) -> str:
|
| 70 |
+
raise NotImplementedError
|
| 71 |
+
|
| 72 |
+
def _fixation_key(self, name: str) -> str:
|
| 73 |
+
"""CSV IMAGE-column key for this sample. Override if it differs from
|
| 74 |
+
the on-disk basename (e.g. prostate MRI uses .jpg keys for .dcm files)."""
|
| 75 |
+
return name
|
| 76 |
+
|
| 77 |
+
def __getitem__(self, idx: int):
|
| 78 |
+
name = self.image_ids[idx]
|
| 79 |
+
|
| 80 |
+
image = self.img_tf(self._load_image(name))
|
| 81 |
+
|
| 82 |
+
mask = self.mask_tf(Image.open(self._mask_path(name)).convert("L"))
|
| 83 |
+
mask = (mask > 0.5).float()
|
| 84 |
+
|
| 85 |
+
fix_rows = self.fix_df.get_group(self._fixation_key(name))
|
| 86 |
+
fixation = get_scanpath(fix_rows, self.max_len)
|
| 87 |
+
|
| 88 |
+
return {"image": image, "fixation": fixation, "mask": mask, "name": name}
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
class KvasirSEGDataset(_BaseGazeDataset):
|
| 92 |
+
"""Kvasir-SEG colonoscopy polyp segmentation. images/masks share filenames
|
| 93 |
+
(e.g. ``cju0qkwl35piu0993l0dewei2.jpg`` in both folders)."""
|
| 94 |
+
|
| 95 |
+
def __init__(self, root: str, fixation_csv: str, img_size: int = IMG_SIZE):
|
| 96 |
+
super().__init__(root, fixation_csv, img_size)
|
| 97 |
+
img_files = set(os.listdir(self.img_dir))
|
| 98 |
+
mask_files = set(os.listdir(self.mask_dir))
|
| 99 |
+
csv_imgs = set(self.df[IMG_COL].unique())
|
| 100 |
+
self.image_ids = sorted(img_files & mask_files & csv_imgs)
|
| 101 |
+
if not self.image_ids:
|
| 102 |
+
raise RuntimeError("No overlap between images/, masks/ and the fixation CSV.")
|
| 103 |
+
print(f"[Dataset] Kvasir-SEG valid samples = {len(self.image_ids)}")
|
| 104 |
+
|
| 105 |
+
def _load_image(self, name: str) -> Image.Image:
|
| 106 |
+
return Image.open(os.path.join(self.img_dir, name)).convert("RGB")
|
| 107 |
+
|
| 108 |
+
def _mask_path(self, name: str) -> str:
|
| 109 |
+
return os.path.join(self.mask_dir, name)
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
class ProstateMRIDataset(_BaseGazeDataset):
|
| 113 |
+
"""NCI-ISBI prostate MRI. images/ holds DICOM (.dcm), masks/ holds PNG,
|
| 114 |
+
and the fixation CSV references each case as ``<basename>.jpg`` (the
|
| 115 |
+
format the eye-tracking session was actually rendered/displayed in)."""
|
| 116 |
+
|
| 117 |
+
def __init__(self, root: str, fixation_csv: str, img_size: int = IMG_SIZE):
|
| 118 |
+
super().__init__(root, fixation_csv, img_size)
|
| 119 |
+
|
| 120 |
+
dcm_basenames = {os.path.splitext(f)[0] for f in os.listdir(self.img_dir) if f.endswith(".dcm")}
|
| 121 |
+
png_basenames = {os.path.splitext(f)[0] for f in os.listdir(self.mask_dir) if f.endswith(".png")}
|
| 122 |
+
csv_basenames = {os.path.splitext(f)[0] for f in self.df[IMG_COL].unique()}
|
| 123 |
+
|
| 124 |
+
self.image_ids = sorted(dcm_basenames & png_basenames & csv_basenames)
|
| 125 |
+
if not self.image_ids:
|
| 126 |
+
raise RuntimeError("No overlap between images/ (.dcm), masks/ (.png) and the fixation CSV.")
|
| 127 |
+
print(f"[Dataset] Prostate MRI valid samples = {len(self.image_ids)}")
|
| 128 |
+
|
| 129 |
+
def _load_image(self, name: str) -> Image.Image:
|
| 130 |
+
import pydicom # local import: optional dependency, only needed for DICOM datasets
|
| 131 |
+
|
| 132 |
+
dcm_path = os.path.join(self.img_dir, f"{name}.dcm")
|
| 133 |
+
dicom = pydicom.dcmread(dcm_path)
|
| 134 |
+
arr = dicom.pixel_array.astype(np.float32)
|
| 135 |
+
|
| 136 |
+
if arr.max() > arr.min():
|
| 137 |
+
arr = (arr - arr.min()) / (arr.max() - arr.min())
|
| 138 |
+
else:
|
| 139 |
+
arr = np.zeros_like(arr)
|
| 140 |
+
|
| 141 |
+
rgb = np.stack([arr, arr, arr], axis=-1) # grayscale -> 3-channel, DINOv3 expects RGB
|
| 142 |
+
return Image.fromarray((rgb * 255).astype(np.uint8))
|
| 143 |
+
|
| 144 |
+
def _mask_path(self, name: str) -> str:
|
| 145 |
+
return os.path.join(self.mask_dir, f"{name}.png")
|
| 146 |
+
|
| 147 |
+
def _fixation_key(self, name: str) -> str:
|
| 148 |
+
return f"{name}.jpg"
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
DATASET_REGISTRY = {
|
| 152 |
+
"kvasir": KvasirSEGDataset,
|
| 153 |
+
"prostate_mri": ProstateMRIDataset,
|
| 154 |
+
}
|
gazerefine/gaze.py
ADDED
|
@@ -0,0 +1,191 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
gaze.py — Turn an eye-tracking scanpath into a spatial prior.
|
| 3 |
+
|
| 4 |
+
Implements the gaze heatmap from the paper:
|
| 5 |
+
|
| 6 |
+
H(u, v) = Σ_m d̃_m · exp( -((u-xm)² + (v-ym)²) / 2σ² )
|
| 7 |
+
|
| 8 |
+
i.e. a duration-weighted sum of 2D Gaussians, one per fixation, on the
|
| 9 |
+
DINOv3 patch grid. The result is min-max normalized and used as the soft
|
| 10 |
+
foreground prior W_fg; (1-H) gives the background complement W_bg.
|
| 11 |
+
|
| 12 |
+
───────────────────────────────────────────────────────────────────────────
|
| 13 |
+
Simple CSV format (what predict_single.py / the README document):
|
| 14 |
+
|
| 15 |
+
x,y,duration
|
| 16 |
+
340,221,180
|
| 17 |
+
356,228,145
|
| 18 |
+
...
|
| 19 |
+
|
| 20 |
+
x, y — fixation position in *raw pixel* coordinates of the original image.
|
| 21 |
+
load_fixation_csv() divides these by the image width/height to
|
| 22 |
+
produce the normalized [0,1] coords the model needs.
|
| 23 |
+
duration — fixation duration in any consistent unit (ms typical).
|
| 24 |
+
|
| 25 |
+
───────────────────────────────────────────────────────────────────────────
|
| 26 |
+
EyeLink/Tobii batch format (used by the full-dataset loaders in datasets.py):
|
| 27 |
+
|
| 28 |
+
IMAGE image filename the fixation belongs to
|
| 29 |
+
CURRENT_FIX_INDEX fixation order within the trial
|
| 30 |
+
CURRENT_FIX_X fixation x, already normalized to [0, 1]
|
| 31 |
+
CURRENT_FIX_Y fixation y, already normalized to [0, 1]
|
| 32 |
+
CURRENT_FIX_DURATION fixation duration (ms or any consistent unit)
|
| 33 |
+
"""
|
| 34 |
+
|
| 35 |
+
from __future__ import annotations
|
| 36 |
+
|
| 37 |
+
import numpy as np
|
| 38 |
+
import pandas as pd
|
| 39 |
+
import torch
|
| 40 |
+
from PIL import Image
|
| 41 |
+
|
| 42 |
+
# ── EyeLink-style column names (used by datasets.py / batch evaluation) ──
|
| 43 |
+
TIME_COL = "CURRENT_FIX_DURATION"
|
| 44 |
+
FIX_IDX_COL = "CURRENT_FIX_INDEX"
|
| 45 |
+
FIX_X_COL = "CURRENT_FIX_X"
|
| 46 |
+
FIX_Y_COL = "CURRENT_FIX_Y"
|
| 47 |
+
IMG_COL = "IMAGE"
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
# ═══════════════════════════════════════════════════════════════════════════
|
| 51 |
+
# Simple CSV loader (README-documented format)
|
| 52 |
+
# ═══════════════════════════════════════════════════════════════════════════
|
| 53 |
+
|
| 54 |
+
def load_fixation_csv(
|
| 55 |
+
csv_path: str,
|
| 56 |
+
image_width: int = None,
|
| 57 |
+
image_height: int = None,
|
| 58 |
+
image_name: str = None,
|
| 59 |
+
) -> torch.Tensor:
|
| 60 |
+
|
| 61 |
+
df = pd.read_csv(csv_path)
|
| 62 |
+
df.columns = df.columns.str.strip()
|
| 63 |
+
|
| 64 |
+
# --------------------------------------------------
|
| 65 |
+
# Match DICOM image IDs to fixation CSV image names
|
| 66 |
+
# Example:
|
| 67 |
+
# Prostate3T-01-0002_16.dcm
|
| 68 |
+
# -> Prostate3T-01-0002_16.jpg
|
| 69 |
+
# --------------------------------------------------
|
| 70 |
+
if "IMAGE" in df.columns and image_name is not None:
|
| 71 |
+
|
| 72 |
+
target_name = f"{image_name}.jpg"
|
| 73 |
+
|
| 74 |
+
df = df[df["IMAGE"] == target_name]
|
| 75 |
+
|
| 76 |
+
if len(df) == 0:
|
| 77 |
+
raise ValueError(
|
| 78 |
+
f"No fixations found for image: {image_name}"
|
| 79 |
+
)
|
| 80 |
+
|
| 81 |
+
# EyeLink/Tobii format
|
| 82 |
+
if (
|
| 83 |
+
"CURRENT_FIX_X" in df.columns
|
| 84 |
+
and "CURRENT_FIX_Y" in df.columns
|
| 85 |
+
and "CURRENT_FIX_DURATION" in df.columns
|
| 86 |
+
):
|
| 87 |
+
x = df["CURRENT_FIX_X"].to_numpy(np.float32)
|
| 88 |
+
y = df["CURRENT_FIX_Y"].to_numpy(np.float32)
|
| 89 |
+
t = df["CURRENT_FIX_DURATION"].to_numpy(np.float32)
|
| 90 |
+
|
| 91 |
+
# README format
|
| 92 |
+
elif "x" in df.columns and "y" in df.columns:
|
| 93 |
+
x = df["x"].to_numpy(np.float32) / image_width
|
| 94 |
+
y = df["y"].to_numpy(np.float32) / image_height
|
| 95 |
+
t = df["duration"].to_numpy(np.float32)
|
| 96 |
+
|
| 97 |
+
else:
|
| 98 |
+
raise ValueError(f"Unsupported CSV format: {df.columns}")
|
| 99 |
+
|
| 100 |
+
# duration normalization
|
| 101 |
+
t_min, t_max = t.min(), t.max()
|
| 102 |
+
t = (t - t_min) / (t_max - t_min + 1e-8)
|
| 103 |
+
|
| 104 |
+
seq = np.stack([x, y, t], axis=1).astype(np.float32)
|
| 105 |
+
|
| 106 |
+
return torch.tensor(
|
| 107 |
+
seq,
|
| 108 |
+
dtype=torch.float32
|
| 109 |
+
).unsqueeze(0)
|
| 110 |
+
|
| 111 |
+
# ═══════════════════════════════════════════════════════════════════════════
|
| 112 |
+
# Batch-evaluation helpers (used by datasets.py)
|
| 113 |
+
# ═══════════════════════════════════════════════════════════════════════════
|
| 114 |
+
|
| 115 |
+
def get_scanpath(df: pd.DataFrame, max_len: int, time_col: str = TIME_COL) -> torch.Tensor:
|
| 116 |
+
"""Build a zero-padded ``(max_len, 3)`` scanpath tensor from a
|
| 117 |
+
fixation-report DataFrame already filtered to a single image.
|
| 118 |
+
|
| 119 |
+
Columns: ``[CURRENT_FIX_X, CURRENT_FIX_Y, CURRENT_FIX_DURATION]``
|
| 120 |
+
already normalized to [0,1] as produced by the EyeLink/Tobii export
|
| 121 |
+
pipeline (datasets.py normalizes these before calling this function).
|
| 122 |
+
"""
|
| 123 |
+
if len(df) == 0:
|
| 124 |
+
return torch.zeros(max_len, 3)
|
| 125 |
+
|
| 126 |
+
df = df.sort_values(FIX_IDX_COL)
|
| 127 |
+
x = df[FIX_X_COL].to_numpy(np.float32)
|
| 128 |
+
y = df[FIX_Y_COL].to_numpy(np.float32)
|
| 129 |
+
t = df[time_col].to_numpy(np.float32)
|
| 130 |
+
|
| 131 |
+
t_min, t_max = t.min(), t.max()
|
| 132 |
+
t = (t - t_min) / (t_max - t_min + 1e-8)
|
| 133 |
+
|
| 134 |
+
seq = np.stack([x, y, t], axis=1)
|
| 135 |
+
if len(seq) > max_len:
|
| 136 |
+
seq = seq[:max_len]
|
| 137 |
+
if len(seq) < max_len:
|
| 138 |
+
seq = np.concatenate([seq, np.zeros((max_len - len(seq), 3), np.float32)], axis=0)
|
| 139 |
+
|
| 140 |
+
return torch.tensor(seq, dtype=torch.float32)
|
| 141 |
+
|
| 142 |
+
|
| 143 |
+
# ═══════════════════════════════════════════════════════════════════════════
|
| 144 |
+
# Heatmap generation
|
| 145 |
+
# ═══════════════════════════════════════════════════════════════════════════
|
| 146 |
+
|
| 147 |
+
def generate_gaze_heatmap(
|
| 148 |
+
fixation: torch.Tensor, h_patch: int, sigma: float = 2.0
|
| 149 |
+
) -> torch.Tensor:
|
| 150 |
+
"""Convert a batch of scanpaths into duration-weighted Gaussian heatmaps.
|
| 151 |
+
|
| 152 |
+
Parameters
|
| 153 |
+
----------
|
| 154 |
+
fixation : (B, L, 3) tensor [x_norm, y_norm, t_norm], zero-padded.
|
| 155 |
+
h_patch : patch-grid side length (74 for IMG_SIZE=1184 with patch_size=16).
|
| 156 |
+
sigma : Gaussian spread in patch units. Larger → softer, more diffuse
|
| 157 |
+
foreground prior; smaller → tighter, closer to exactly the
|
| 158 |
+
fixated patches.
|
| 159 |
+
|
| 160 |
+
Returns
|
| 161 |
+
-------
|
| 162 |
+
(B, h_patch, h_patch) float32 tensor, each sample min-max normalized
|
| 163 |
+
to [0, 1] independently.
|
| 164 |
+
"""
|
| 165 |
+
B, L, _ = fixation.shape
|
| 166 |
+
device = fixation.device
|
| 167 |
+
|
| 168 |
+
y_coords = torch.arange(h_patch, device=device, dtype=torch.float32)
|
| 169 |
+
x_coords = torch.arange(h_patch, device=device, dtype=torch.float32)
|
| 170 |
+
grid_y, grid_x = torch.meshgrid(y_coords, x_coords, indexing="ij")
|
| 171 |
+
|
| 172 |
+
heatmaps = []
|
| 173 |
+
for b in range(B):
|
| 174 |
+
valid = fixation[b].abs().sum(dim=-1) > 1e-8 # filter zero-padding rows
|
| 175 |
+
valid_fix = fixation[b][valid]
|
| 176 |
+
|
| 177 |
+
h_b = torch.zeros(h_patch, h_patch, device=device)
|
| 178 |
+
if len(valid_fix) > 0:
|
| 179 |
+
for k in range(valid_fix.shape[0]):
|
| 180 |
+
x_k, y_k, t_k = valid_fix[k, 0], valid_fix[k, 1], valid_fix[k, 2]
|
| 181 |
+
gx = x_k * (h_patch - 1)
|
| 182 |
+
gy = y_k * (h_patch - 1)
|
| 183 |
+
dist_sq = (grid_x - gx) ** 2 + (grid_y - gy) ** 2
|
| 184 |
+
h_b += t_k * torch.exp(-dist_sq / (2 * sigma ** 2))
|
| 185 |
+
|
| 186 |
+
h_max = h_b.max()
|
| 187 |
+
if h_max > 1e-8:
|
| 188 |
+
h_b = h_b / h_max
|
| 189 |
+
heatmaps.append(h_b)
|
| 190 |
+
|
| 191 |
+
return torch.stack(heatmaps, dim=0)
|
gazerefine/metrics.py
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""metrics.py — Dice / IoU against a binarized prediction."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import torch
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
@torch.no_grad()
|
| 9 |
+
def compute_metrics(preds: torch.Tensor, mask: torch.Tensor, thr: float = 0.5) -> dict:
|
| 10 |
+
"""
|
| 11 |
+
preds, mask : (B, 1, H, W). ``preds`` is the continuous [0, 1] GazeRefine
|
| 12 |
+
output, ``mask`` is the ground-truth binary mask.
|
| 13 |
+
"""
|
| 14 |
+
pred_bin = (preds > thr).float()
|
| 15 |
+
tp = (pred_bin * mask).sum((1, 2, 3))
|
| 16 |
+
fp = (pred_bin * (1 - mask)).sum((1, 2, 3))
|
| 17 |
+
fn = ((1 - pred_bin) * mask).sum((1, 2, 3))
|
| 18 |
+
|
| 19 |
+
dice_per = 2 * tp / (2 * tp + fp + fn + 1e-8)
|
| 20 |
+
iou_per = tp / (tp + fp + fn + 1e-8)
|
| 21 |
+
|
| 22 |
+
return {
|
| 23 |
+
"dice": dice_per.mean().item(),
|
| 24 |
+
"iou": iou_per.mean().item(),
|
| 25 |
+
"dice_per": dice_per,
|
| 26 |
+
"iou_per": iou_per,
|
| 27 |
+
}
|
gazerefine/model.py
ADDED
|
@@ -0,0 +1,242 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
model.py — The GazeRefine zero-shot segmentation model.
|
| 3 |
+
|
| 4 |
+
This is a 1:1 refactor of the two task-specific scripts (Kvasir-SEG polyp /
|
| 5 |
+
NCI-ISBI prostate MRI) into a single, dataset-agnostic module. The numerics
|
| 6 |
+
are unchanged from the original experiments — only structure, naming and
|
| 7 |
+
comments were cleaned up. Every per-dataset difference (sigma, kNN k,
|
| 8 |
+
gaze-anchor weight, max_iters, ...) is now a constructor / config argument
|
| 9 |
+
instead of a hardcoded default, see ``configs/*.yaml``.
|
| 10 |
+
|
| 11 |
+
Maps onto the paper (Section 2) as follows:
|
| 12 |
+
|
| 13 |
+
gaze.generate_gaze_heatmap -> H(u, v) (Eq. gaze prior)
|
| 14 |
+
W_fg, W_bg -> W_fg^(0), W_bg^(0) (gaze-derived weights)
|
| 15 |
+
F_proto_init, B_proto_init -> F^(0), B^(0) (initial prototypes)
|
| 16 |
+
contrast_method="difference" -> S_i^(t) = max(0, s_fg - alpha*s_bg) (contrastive cleaning)
|
| 17 |
+
knn_affinity_refinement -> S̄_i^(t) (kNN affinity propagation)
|
| 18 |
+
gaze_anchor_weight -> lambda (anchor blending strength)
|
| 19 |
+
convergence check -> ||F^(t+1) - F^(t)|| < eps
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
from __future__ import annotations
|
| 23 |
+
|
| 24 |
+
import torch
|
| 25 |
+
import torch.nn as nn
|
| 26 |
+
import torch.nn.functional as F
|
| 27 |
+
|
| 28 |
+
from .backbone import FrozenDINOv3
|
| 29 |
+
from .gaze import generate_gaze_heatmap
|
| 30 |
+
from .constants import IMG_SIZE, PATCH_SIZE
|
| 31 |
+
|
| 32 |
+
EPS = 1e-12
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def knn_affinity_refinement(
|
| 36 |
+
Pv: torch.Tensor, S: torch.Tensor, k: int = 5, temperature: float = 0.05
|
| 37 |
+
) -> torch.Tensor:
|
| 38 |
+
"""Propagate a per-patch score map ``S`` across its k nearest neighbors in
|
| 39 |
+
frozen DINOv3 feature space (patch-to-patch affinity), encouraging
|
| 40 |
+
coherent, object-level responses instead of isolated high-confidence
|
| 41 |
+
patches. Corresponds to S̄ in the paper.
|
| 42 |
+
|
| 43 |
+
Pv : (B, N, D) raw (unnormalized) patch embeddings.
|
| 44 |
+
S : (B, N) current per-patch score to be smoothed.
|
| 45 |
+
"""
|
| 46 |
+
B, N, D = Pv.shape
|
| 47 |
+
Pv_norm = F.normalize(Pv, p=2, dim=-1)
|
| 48 |
+
|
| 49 |
+
sim_matrix = torch.bmm(Pv_norm, Pv_norm.transpose(1, 2)) # (B, N, N)
|
| 50 |
+
topk_vals, topk_indices = torch.topk(sim_matrix, k=k, dim=-1) # (B, N, k)
|
| 51 |
+
weights = F.softmax(topk_vals / temperature, dim=-1) # (B, N, k)
|
| 52 |
+
|
| 53 |
+
S_expanded = S.unsqueeze(1).expand(B, N, N)
|
| 54 |
+
S_neighbors = torch.gather(S_expanded, dim=2, index=topk_indices) # (B, N, k)
|
| 55 |
+
|
| 56 |
+
return torch.sum(weights * S_neighbors, dim=-1) # (B, N)
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
class GazeRefine(nn.Module):
|
| 60 |
+
"""Zero-shot, training-free, gaze-guided segmentation model.
|
| 61 |
+
|
| 62 |
+
The model has exactly one set of *learned* weights: the frozen,
|
| 63 |
+
pretrained DINOv3 backbone. Everything else — prototype construction,
|
| 64 |
+
contrastive cleaning, kNN propagation, gaze anchoring — is a closed-form
|
| 65 |
+
operation re-run from scratch on every image at inference time.
|
| 66 |
+
|
| 67 |
+
Parameters
|
| 68 |
+
----------
|
| 69 |
+
dino_name : timm DINOv3 checkpoint name.
|
| 70 |
+
img_size, patch_size : input resolution / ViT patch size. Must satisfy
|
| 71 |
+
``img_size % patch_size == 0``.
|
| 72 |
+
sigma : Gaussian spread (in patch units) for the gaze heatmap.
|
| 73 |
+
extract_mode : ``"last"`` (1 block) or ``"all"`` (4 blocks, averaged).
|
| 74 |
+
contrast_method : ``"difference"`` (paper default, contrastive cleaning),
|
| 75 |
+
``"softmax"`` (foreground/background softmax ratio), or
|
| 76 |
+
``"original"`` (plain foreground cosine similarity, no background
|
| 77 |
+
suppression — kept for the ablation in Table 2).
|
| 78 |
+
temperature : softmax temperature, only used when contrast_method="softmax".
|
| 79 |
+
max_iters : maximum recurrent refinement iterations (T in the paper).
|
| 80 |
+
knn_refine : whether to apply kNN affinity propagation each iteration.
|
| 81 |
+
knn_k, knn_temp : kNN neighborhood size / softmax temperature.
|
| 82 |
+
gaze_anchor_weight : lambda — how strongly each iteration's prototypes
|
| 83 |
+
are pulled back toward the original gaze-only prototypes. Higher =
|
| 84 |
+
trust the raw fixations more; lower = let the model drift further
|
| 85 |
+
from the initial gaze region.
|
| 86 |
+
"""
|
| 87 |
+
|
| 88 |
+
def __init__(
|
| 89 |
+
self,
|
| 90 |
+
dino_name: str = "vit_large_patch16_dinov3.lvd1689m",
|
| 91 |
+
img_size: int = IMG_SIZE,
|
| 92 |
+
patch_size: int = PATCH_SIZE,
|
| 93 |
+
sigma: float = 2.0,
|
| 94 |
+
extract_mode: str = "last",
|
| 95 |
+
contrast_method: str = "difference",
|
| 96 |
+
temperature: float = 0.05,
|
| 97 |
+
max_iters: int = 10,
|
| 98 |
+
knn_refine: bool = True,
|
| 99 |
+
knn_k: int = 5,
|
| 100 |
+
knn_temp: float = 0.1,
|
| 101 |
+
gaze_anchor_weight: float = 0.6,
|
| 102 |
+
):
|
| 103 |
+
super().__init__()
|
| 104 |
+
assert img_size % patch_size == 0, "img_size must be a multiple of patch_size"
|
| 105 |
+
self.img_size = img_size
|
| 106 |
+
self.h_patch = img_size // patch_size
|
| 107 |
+
self.sigma = sigma
|
| 108 |
+
self.contrast_method = contrast_method
|
| 109 |
+
self.temperature = temperature
|
| 110 |
+
self.max_iters = max_iters
|
| 111 |
+
self.knn_refine = knn_refine
|
| 112 |
+
self.knn_k = knn_k
|
| 113 |
+
self.knn_temp = knn_temp
|
| 114 |
+
self.gaze_anchor_weight = gaze_anchor_weight
|
| 115 |
+
|
| 116 |
+
self.visual_enc = FrozenDINOv3(dino_name, extract_mode=extract_mode)
|
| 117 |
+
|
| 118 |
+
# ------------------------------------------------------------------ #
|
| 119 |
+
def _gaze_weights(self, fixation: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
| 120 |
+
"""Scanpath -> heatmap -> normalized foreground/background spatial weights."""
|
| 121 |
+
gaze_heatmap = generate_gaze_heatmap(fixation, h_patch=self.h_patch, sigma=self.sigma)
|
| 122 |
+
H_flat = gaze_heatmap.view(gaze_heatmap.size(0), -1)
|
| 123 |
+
|
| 124 |
+
W_fg = H_flat / (H_flat.sum(dim=-1, keepdim=True) + EPS)
|
| 125 |
+
bg_w = 1.0 - H_flat
|
| 126 |
+
W_bg = bg_w / (bg_w.sum(dim=-1, keepdim=True) + EPS)
|
| 127 |
+
return gaze_heatmap, W_fg, W_bg
|
| 128 |
+
|
| 129 |
+
def _score(self, Pv_norm: torch.Tensor, F_proto_norm: torch.Tensor, B_proto_norm: torch.Tensor) -> torch.Tensor:
|
| 130 |
+
"""Foreground/background contrastive scoring for one level, one iteration."""
|
| 131 |
+
if self.contrast_method == "original":
|
| 132 |
+
S = torch.bmm(Pv_norm, F_proto_norm.unsqueeze(-1)).squeeze(-1)
|
| 133 |
+
return torch.clamp(S, min=0.0)
|
| 134 |
+
|
| 135 |
+
sim_fg = torch.bmm(Pv_norm, F_proto_norm.unsqueeze(-1)).squeeze(-1)
|
| 136 |
+
sim_bg = torch.bmm(Pv_norm, B_proto_norm.unsqueeze(-1)).squeeze(-1)
|
| 137 |
+
|
| 138 |
+
if self.contrast_method == "difference":
|
| 139 |
+
return torch.clamp(sim_fg - sim_bg, min=0.0)
|
| 140 |
+
elif self.contrast_method == "softmax":
|
| 141 |
+
stacked = torch.stack([sim_fg, sim_bg], dim=-1) / self.temperature
|
| 142 |
+
probs = F.softmax(stacked, dim=-1)
|
| 143 |
+
return probs[:, :, 0]
|
| 144 |
+
raise ValueError(f"Unknown contrast_method: {self.contrast_method!r}")
|
| 145 |
+
|
| 146 |
+
def _refine_level(self, Pv: torch.Tensor, W_fg: torch.Tensor, W_bg: torch.Tensor) -> torch.Tensor:
|
| 147 |
+
"""Run the full recurrent gaze-anchored refinement loop for one DINOv3 level.
|
| 148 |
+
Returns the final (B, N) per-patch foreground score map for that level."""
|
| 149 |
+
B = Pv.size(0)
|
| 150 |
+
device = Pv.device
|
| 151 |
+
Pv_norm = F.normalize(Pv, p=2, dim=-1)
|
| 152 |
+
|
| 153 |
+
# Initial gaze-only prototypes — F^(0), B^(0)
|
| 154 |
+
F_proto_init = torch.sum(W_fg.unsqueeze(-1) * Pv, dim=1)
|
| 155 |
+
B_proto_init = torch.sum(W_bg.unsqueeze(-1) * Pv, dim=1)
|
| 156 |
+
|
| 157 |
+
W_fg_curr, W_bg_curr = W_fg.clone(), W_bg.clone()
|
| 158 |
+
best_S = W_fg.clone()
|
| 159 |
+
S_prev = W_fg.clone()
|
| 160 |
+
active = torch.ones(B, dtype=torch.bool, device=device)
|
| 161 |
+
|
| 162 |
+
for _ in range(self.max_iters):
|
| 163 |
+
W_fg_norm = W_fg_curr / (W_fg_curr.sum(dim=-1, keepdim=True) + EPS)
|
| 164 |
+
W_bg_norm = W_bg_curr / (W_bg_curr.sum(dim=-1, keepdim=True) + EPS)
|
| 165 |
+
|
| 166 |
+
F_proto_curr = torch.sum(W_fg_norm.unsqueeze(-1) * Pv, dim=1)
|
| 167 |
+
B_proto_curr = torch.sum(W_bg_norm.unsqueeze(-1) * Pv, dim=1)
|
| 168 |
+
|
| 169 |
+
# Anchor toward the initial gaze-only prototypes — lambda blending
|
| 170 |
+
lam = self.gaze_anchor_weight
|
| 171 |
+
F_proto = (1 - lam) * F_proto_curr + lam * F_proto_init
|
| 172 |
+
B_proto = (1 - lam) * B_proto_curr + lam * B_proto_init
|
| 173 |
+
F_proto_norm = F.normalize(F_proto, p=2, dim=-1)
|
| 174 |
+
B_proto_norm = F.normalize(B_proto, p=2, dim=-1)
|
| 175 |
+
|
| 176 |
+
S_iter = self._score(Pv_norm, F_proto_norm, B_proto_norm)
|
| 177 |
+
if self.knn_refine:
|
| 178 |
+
S_iter = knn_affinity_refinement(Pv, S_iter, k=self.knn_k, temperature=self.knn_temp)
|
| 179 |
+
|
| 180 |
+
S_new = S_prev.clone()
|
| 181 |
+
S_new[active] = S_iter[active]
|
| 182 |
+
|
| 183 |
+
# Collapse prevention: if an active sample's map flattened to ~0,
|
| 184 |
+
# roll it back and freeze it instead of letting it degenerate.
|
| 185 |
+
collapsed = S_new.max(dim=-1)[0] < 1e-4
|
| 186 |
+
freeze = active & collapsed
|
| 187 |
+
S_new[freeze] = S_prev[freeze]
|
| 188 |
+
active = active & ~collapsed
|
| 189 |
+
|
| 190 |
+
best_S = S_new.clone()
|
| 191 |
+
if not active.any():
|
| 192 |
+
break
|
| 193 |
+
|
| 194 |
+
# Turn the refined score into the next iteration's spatial weights
|
| 195 |
+
S_sig = torch.sigmoid(S_new)
|
| 196 |
+
W_fg_next = S_sig / (S_sig.sum(dim=-1, keepdim=True) + EPS)
|
| 197 |
+
bg_w = 1.0 - S_sig
|
| 198 |
+
bg_w = torch.clamp(bg_w - bg_w.min(dim=-1, keepdim=True).values, min=0.0)
|
| 199 |
+
W_bg_next = bg_w / (bg_w.sum(dim=-1, keepdim=True) + EPS)
|
| 200 |
+
|
| 201 |
+
F_proto_next = torch.sum(W_fg_next.unsqueeze(-1) * Pv, dim=1)
|
| 202 |
+
F_proto_next_norm = F.normalize(F_proto_next, p=2, dim=-1)
|
| 203 |
+
|
| 204 |
+
# Convergence: foreground prototype stopped moving
|
| 205 |
+
diff = torch.abs(F_proto_next_norm - F_proto_norm).mean(dim=-1)
|
| 206 |
+
if torch.all(diff < 1e-6):
|
| 207 |
+
break
|
| 208 |
+
|
| 209 |
+
S_prev = S_new.clone()
|
| 210 |
+
W_fg_curr[active] = W_fg_next[active]
|
| 211 |
+
W_bg_curr[active] = W_bg_next[active]
|
| 212 |
+
|
| 213 |
+
return best_S
|
| 214 |
+
|
| 215 |
+
# ------------------------------------------------------------------ #
|
| 216 |
+
def forward(self, image: torch.Tensor, fixation: torch.Tensor) -> dict:
|
| 217 |
+
"""
|
| 218 |
+
image : (B, 3, H, W) normalized RGB tensor (ImageNet mean/std).
|
| 219 |
+
fixation : (B, L, 3) zero-padded scanpath, see ``gazerefine.gaze.get_scanpath``.
|
| 220 |
+
|
| 221 |
+
Returns a dict with:
|
| 222 |
+
preds (B, 1, H, W) final mask in [0, 1] — threshold at 0.5 for a binary mask
|
| 223 |
+
gaze_heatmap (B, h, w) the raw gaze prior, useful for visualization
|
| 224 |
+
final_map (B, N) the un-upsampled patch-level score map
|
| 225 |
+
"""
|
| 226 |
+
B = image.size(0)
|
| 227 |
+
Pv_list = self.visual_enc(image)
|
| 228 |
+
gaze_heatmap, W_fg, W_bg = self._gaze_weights(fixation)
|
| 229 |
+
|
| 230 |
+
all_maps = [self._refine_level(Pv, W_fg, W_bg) for Pv in Pv_list]
|
| 231 |
+
final_map = torch.stack(all_maps, dim=1).mean(dim=1) # average across levels
|
| 232 |
+
|
| 233 |
+
grid = final_map.view(B, 1, self.h_patch, self.h_patch)
|
| 234 |
+
upsampled = F.interpolate(grid, size=(self.img_size, self.img_size), mode="bilinear", align_corners=False)
|
| 235 |
+
|
| 236 |
+
# per-sample min-max normalization -> ready-to-threshold mask
|
| 237 |
+
flat = upsampled.view(B, -1)
|
| 238 |
+
m_min = flat.min(dim=-1, keepdim=True).values.view(B, 1, 1, 1)
|
| 239 |
+
m_max = flat.max(dim=-1, keepdim=True).values.view(B, 1, 1, 1)
|
| 240 |
+
preds = (upsampled - m_min) / (m_max - m_min + 1e-8)
|
| 241 |
+
|
| 242 |
+
return dict(preds=preds, gaze_heatmap=gaze_heatmap, final_map=final_map)
|
gazerefine/visualize.py
ADDED
|
@@ -0,0 +1,86 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""visualize.py — Plotting / overlay helpers.
|
| 2 |
+
|
| 3 |
+
``save_prediction`` reproduces the 4-panel (image / gaze / GT / prediction)
|
| 4 |
+
grid used during batch evaluation. ``overlay_heatmap`` and ``overlay_mask``
|
| 5 |
+
are lighter-weight, matplotlib-free PIL helpers meant for the Gradio demo
|
| 6 |
+
and the example notebook, where you typically want a single composited
|
| 7 |
+
image rather than a saved subplot file.
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
import os
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import torch
|
| 16 |
+
from PIL import Image
|
| 17 |
+
|
| 18 |
+
import matplotlib
|
| 19 |
+
matplotlib.use("Agg")
|
| 20 |
+
import matplotlib.pyplot as plt
|
| 21 |
+
import matplotlib.cm as cm
|
| 22 |
+
|
| 23 |
+
from .constants import IMG_MEAN, IMG_STD
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def _denormalize(img: torch.Tensor) -> np.ndarray:
|
| 27 |
+
"""(3, H, W) normalized tensor -> (H, W, 3) uint8-range float array in [0, 1]."""
|
| 28 |
+
mean = torch.tensor(IMG_MEAN, device=img.device).view(3, 1, 1)
|
| 29 |
+
std = torch.tensor(IMG_STD, device=img.device).view(3, 1, 1)
|
| 30 |
+
return (img * std + mean).clamp(0, 1).permute(1, 2, 0).cpu().numpy()
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
@torch.no_grad()
|
| 34 |
+
def save_prediction(
|
| 35 |
+
img: torch.Tensor,
|
| 36 |
+
gaze_heatmap: torch.Tensor,
|
| 37 |
+
gt: torch.Tensor,
|
| 38 |
+
pred: torch.Tensor,
|
| 39 |
+
name: str,
|
| 40 |
+
out_dir: str = "predictions",
|
| 41 |
+
) -> None:
|
| 42 |
+
"""Save a 4-panel [image | gaze heatmap | GT mask | predicted mask] figure."""
|
| 43 |
+
os.makedirs(out_dir, exist_ok=True)
|
| 44 |
+
|
| 45 |
+
img_np = _denormalize(img)
|
| 46 |
+
gaze_np = gaze_heatmap.squeeze().cpu().numpy()
|
| 47 |
+
gt_np = gt.squeeze().cpu().numpy()
|
| 48 |
+
pred_np = pred.squeeze().cpu().numpy()
|
| 49 |
+
|
| 50 |
+
fig, ax = plt.subplots(1, 4, figsize=(20, 5))
|
| 51 |
+
ax[0].imshow(img_np); ax[0].set_title("Image")
|
| 52 |
+
ax[1].imshow(gaze_np, cmap="jet"); ax[1].set_title("Gaze Heatmap")
|
| 53 |
+
ax[2].imshow(gt_np, cmap="gray"); ax[2].set_title("GT mask")
|
| 54 |
+
ax[3].imshow(pred_np, cmap="gray"); ax[3].set_title("Predicted mask")
|
| 55 |
+
for a in ax:
|
| 56 |
+
a.axis("off")
|
| 57 |
+
plt.tight_layout()
|
| 58 |
+
plt.savefig(os.path.join(out_dir, name), dpi=100, bbox_inches="tight")
|
| 59 |
+
plt.close()
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def overlay_heatmap(image: Image.Image, heatmap: np.ndarray, alpha: float = 0.45) -> Image.Image:
|
| 63 |
+
"""Composite a [0, 1] heatmap (any H'xW', will be resized) onto a PIL image
|
| 64 |
+
using a jet colormap. Used to show the gaze prior over the original image."""
|
| 65 |
+
heatmap = np.asarray(heatmap, dtype=np.float32)
|
| 66 |
+
heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min() + 1e-8)
|
| 67 |
+
|
| 68 |
+
heat_img = Image.fromarray((cm.jet(heatmap)[:, :, :3] * 255).astype(np.uint8))
|
| 69 |
+
heat_img = heat_img.resize(image.size, resample=Image.BILINEAR)
|
| 70 |
+
|
| 71 |
+
base = image.convert("RGB")
|
| 72 |
+
return Image.blend(base, heat_img, alpha=alpha)
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
def overlay_mask(
|
| 76 |
+
image: Image.Image, mask: np.ndarray, color: tuple[int, int, int] = (255, 60, 60), alpha: float = 0.45
|
| 77 |
+
) -> Image.Image:
|
| 78 |
+
"""Composite a binary/soft [0, 1] mask onto a PIL image as a solid color wash."""
|
| 79 |
+
mask = np.asarray(mask, dtype=np.float32)
|
| 80 |
+
mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)
|
| 81 |
+
mask_img = Image.fromarray((mask * 255).astype(np.uint8)).resize(image.size, resample=Image.NEAREST)
|
| 82 |
+
|
| 83 |
+
base = image.convert("RGB")
|
| 84 |
+
color_layer = Image.new("RGB", base.size, color)
|
| 85 |
+
composited = Image.composite(color_layer, base, mask_img.point(lambda p: int(p * alpha)))
|
| 86 |
+
return composited
|
notebooks/GazeRefine_Demo.ipynb
ADDED
|
@@ -0,0 +1,306 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"cells": [
|
| 3 |
+
{
|
| 4 |
+
"cell_type": "markdown",
|
| 5 |
+
"metadata": {},
|
| 6 |
+
"source": [
|
| 7 |
+
"# GazeRefine \u2014 Interactive Walkthrough\n",
|
| 8 |
+
"\n",
|
| 9 |
+
"This notebook walks through the **full GazeRefine pipeline** step by step:\n",
|
| 10 |
+
"\n",
|
| 11 |
+
"1. Build a gaze heatmap from a scanpath\n",
|
| 12 |
+
"2. Construct gaze-anchored foreground/background prototypes in frozen DINOv3 feature space\n",
|
| 13 |
+
"3. Run the recurrent refinement loop (contrastive cleaning + kNN affinity propagation)\n",
|
| 14 |
+
"4. Visualize the final mask\n",
|
| 15 |
+
"\n",
|
| 16 |
+
"It ships with a **synthetic toy example** (a blob on a noisy background, with a few clicked\n",
|
| 17 |
+
"fixation points) so you can see every intermediate step without needing any dataset.\n",
|
| 18 |
+
"Swap in a real image + fixation CSV from `examples/` (see `examples/README.md`) to run it\n",
|
| 19 |
+
"on actual data \u2014 the API is identical.\n",
|
| 20 |
+
"\n",
|
| 21 |
+
"> Section 4 needs `torch` + `timm` (and a DINOv3 checkpoint download, so internet access)\n",
|
| 22 |
+
"> to call the real backbone. Sections 1\u20133 only need `numpy`/`matplotlib` and run anywhere,\n",
|
| 23 |
+
"> using a tiny **mock feature extractor** so you can still see the prototype/refinement math."
|
| 24 |
+
]
|
| 25 |
+
},
|
| 26 |
+
{
|
| 27 |
+
"cell_type": "code",
|
| 28 |
+
"execution_count": null,
|
| 29 |
+
"metadata": {},
|
| 30 |
+
"outputs": [],
|
| 31 |
+
"source": [
|
| 32 |
+
"import sys, os\n",
|
| 33 |
+
"from pathlib import Path\n",
|
| 34 |
+
"\n",
|
| 35 |
+
"REPO_ROOT = Path.cwd().parent if Path.cwd().name == \"notebooks\" else Path.cwd()\n",
|
| 36 |
+
"sys.path.insert(0, str(REPO_ROOT))\n",
|
| 37 |
+
"\n",
|
| 38 |
+
"import numpy as np\n",
|
| 39 |
+
"import matplotlib.pyplot as plt\n",
|
| 40 |
+
"\n",
|
| 41 |
+
"np.random.seed(0)\n",
|
| 42 |
+
"%matplotlib inline"
|
| 43 |
+
]
|
| 44 |
+
},
|
| 45 |
+
{
|
| 46 |
+
"cell_type": "markdown",
|
| 47 |
+
"metadata": {},
|
| 48 |
+
"source": [
|
| 49 |
+
"## 1. A synthetic example\n",
|
| 50 |
+
"\n",
|
| 51 |
+
"We draw a soft circular \"lesion\" on a noisy background and pick a few fixation points near\n",
|
| 52 |
+
"its center \u2014 mimicking a clinician glancing at the structure of interest a handful of times,\n",
|
| 53 |
+
"with the first, longest fixation landing closest to the center."
|
| 54 |
+
]
|
| 55 |
+
},
|
| 56 |
+
{
|
| 57 |
+
"cell_type": "code",
|
| 58 |
+
"execution_count": null,
|
| 59 |
+
"metadata": {},
|
| 60 |
+
"outputs": [],
|
| 61 |
+
"source": [
|
| 62 |
+
"IMG_SIZE = 224 # small, just for a fast visual demo (the real pipeline defaults to 1184 / DINOv3 patch grid)\n",
|
| 63 |
+
"\n",
|
| 64 |
+
"def make_toy_image(size=IMG_SIZE, cx=0.55, cy=0.45, r=0.12):\n",
|
| 65 |
+
" yy, xx = np.mgrid[0:size, 0:size] / size\n",
|
| 66 |
+
" dist = np.sqrt((xx - cx) ** 2 + (yy - cy) ** 2)\n",
|
| 67 |
+
" blob = np.clip(1 - dist / r, 0, 1) ** 2\n",
|
| 68 |
+
" background = 0.25 + 0.05 * np.random.randn(size, size)\n",
|
| 69 |
+
" img = np.clip(background + 0.6 * blob, 0, 1)\n",
|
| 70 |
+
" gt_mask = (dist < r).astype(np.float32)\n",
|
| 71 |
+
" return img, gt_mask\n",
|
| 72 |
+
"\n",
|
| 73 |
+
"toy_img, toy_mask = make_toy_image()\n",
|
| 74 |
+
"\n",
|
| 75 |
+
"# a few fixations clustered near the blob, decaying duration, with a touch of natural jitter\n\n",
|
| 76 |
+
"fixations = [\n",
|
| 77 |
+
" (0.55, 0.45, 1.0),\n",
|
| 78 |
+
" (0.57, 0.47, 0.7),\n",
|
| 79 |
+
" (0.53, 0.44, 0.5),\n",
|
| 80 |
+
" (0.50, 0.50, 0.3),\n",
|
| 81 |
+
"]\n",
|
| 82 |
+
"\n",
|
| 83 |
+
"fig, ax = plt.subplots(1, 2, figsize=(8, 4))\n",
|
| 84 |
+
"ax[0].imshow(toy_img, cmap=\"gray\"); ax[0].set_title(\"Synthetic image\")\n",
|
| 85 |
+
"for x, y, d in fixations:\n",
|
| 86 |
+
" ax[0].scatter(x * IMG_SIZE, y * IMG_SIZE, s=200 * d, c=\"red\", alpha=0.6, edgecolors=\"white\")\n",
|
| 87 |
+
"ax[1].imshow(toy_mask, cmap=\"gray\"); ax[1].set_title(\"Ground-truth mask (synthetic)\")\n",
|
| 88 |
+
"for a in ax: a.axis(\"off\")\n",
|
| 89 |
+
"plt.tight_layout(); plt.show()"
|
| 90 |
+
]
|
| 91 |
+
},
|
| 92 |
+
{
|
| 93 |
+
"cell_type": "markdown",
|
| 94 |
+
"metadata": {},
|
| 95 |
+
"source": [
|
| 96 |
+
"## 2. Scanpath \u2192 gaze heatmap\n",
|
| 97 |
+
"\n",
|
| 98 |
+
"`gazerefine.gaze.generate_gaze_heatmap` places a duration-weighted Gaussian at each fixation\n",
|
| 99 |
+
"and min-max normalizes the result. This becomes the soft foreground prior $W_{fg}^{(0)}$;\n",
|
| 100 |
+
"$1 - H$ becomes the background prior $W_{bg}^{(0)}$."
|
| 101 |
+
]
|
| 102 |
+
},
|
| 103 |
+
{
|
| 104 |
+
"cell_type": "code",
|
| 105 |
+
"execution_count": null,
|
| 106 |
+
"metadata": {},
|
| 107 |
+
"outputs": [],
|
| 108 |
+
"source": [
|
| 109 |
+
"import torch\n",
|
| 110 |
+
"from gazerefine.gaze import generate_gaze_heatmap\n",
|
| 111 |
+
"\n",
|
| 112 |
+
"H_PATCH = 28 # toy patch grid (real pipeline uses 74x74, see gazerefine/constants.py)\n",
|
| 113 |
+
"fixation_t = torch.tensor([fixations], dtype=torch.float32) # (1, L, 3)\n",
|
| 114 |
+
"\n",
|
| 115 |
+
"gaze_heatmap = generate_gaze_heatmap(fixation_t, h_patch=H_PATCH, sigma=2.0)[0].numpy()\n",
|
| 116 |
+
"\n",
|
| 117 |
+
"plt.figure(figsize=(4, 4))\n",
|
| 118 |
+
"plt.imshow(gaze_heatmap, cmap=\"jet\")\n",
|
| 119 |
+
"plt.title(\"Duration-weighted gaze heatmap (patch grid)\")\n",
|
| 120 |
+
"plt.axis(\"off\"); plt.colorbar(fraction=0.046); plt.show()"
|
| 121 |
+
]
|
| 122 |
+
},
|
| 123 |
+
{
|
| 124 |
+
"cell_type": "markdown",
|
| 125 |
+
"metadata": {},
|
| 126 |
+
"source": [
|
| 127 |
+
"## 3. Prototype construction & refinement \u2014 mechanics demo\n",
|
| 128 |
+
"\n",
|
| 129 |
+
"To see the **math** of gaze-anchored prototypes, contrastive cleaning, and kNN affinity\n",
|
| 130 |
+
"propagation without downloading a backbone, we substitute a tiny mock \"feature extractor\":\n",
|
| 131 |
+
"patches near the blob get embeddings clustered around one random direction, background\n",
|
| 132 |
+
"patches cluster around another, with noise \u2014 a cartoon stand-in for \"DINOv3 already groups\n",
|
| 133 |
+
"semantically similar patches together\", which is the real property GazeRefine exploits.\n",
|
| 134 |
+
"\n",
|
| 135 |
+
"**This section is for intuition only.** Section 4 below runs the real frozen DINOv3 backbone."
|
| 136 |
+
]
|
| 137 |
+
},
|
| 138 |
+
{
|
| 139 |
+
"cell_type": "code",
|
| 140 |
+
"execution_count": null,
|
| 141 |
+
"metadata": {},
|
| 142 |
+
"outputs": [],
|
| 143 |
+
"source": [
|
| 144 |
+
"from gazerefine.model import knn_affinity_refinement\n\n",
|
| 145 |
+
"\n",
|
| 146 |
+
"D = 64 # toy embedding dim\n\n",
|
| 147 |
+
"rng = np.random.default_rng(0)\n\n",
|
| 148 |
+
"fg_direction = rng.normal(size=D); fg_direction /= np.linalg.norm(fg_direction)\n\n",
|
| 149 |
+
"bg_direction = rng.normal(size=D); bg_direction /= np.linalg.norm(bg_direction)\n\n",
|
| 150 |
+
"\n\n",
|
| 151 |
+
"# resize the toy GT mask down to the patch grid to decide which mock cluster each patch belongs to\n\n",
|
| 152 |
+
"import torch.nn.functional as Fnn\n\n",
|
| 153 |
+
"mask_small = Fnn.interpolate(torch.tensor(toy_mask)[None, None], size=(H_PATCH, H_PATCH), mode=\"nearest\")[0, 0].numpy()\n\n",
|
| 154 |
+
"patch_is_fg = mask_small.flatten() > 0.5\n\n",
|
| 155 |
+
"\n\n",
|
| 156 |
+
"N = H_PATCH * H_PATCH\n\n",
|
| 157 |
+
"Pv = np.stack([\n\n",
|
| 158 |
+
" (fg_direction if patch_is_fg[i] else bg_direction) + 0.3 * rng.normal(size=D)\n\n",
|
| 159 |
+
" for i in range(N)\n\n",
|
| 160 |
+
"])\n\n",
|
| 161 |
+
"Pv = torch.tensor(Pv, dtype=torch.float32)[None] # (1, N, D)\n"
|
| 162 |
+
]
|
| 163 |
+
},
|
| 164 |
+
{
|
| 165 |
+
"cell_type": "code",
|
| 166 |
+
"execution_count": null,
|
| 167 |
+
"metadata": {},
|
| 168 |
+
"outputs": [],
|
| 169 |
+
"source": [
|
| 170 |
+
"H_flat = torch.tensor(gaze_heatmap.flatten(), dtype=torch.float32)[None] # (1, N)\n\n",
|
| 171 |
+
"eps = 1e-12\n\n",
|
| 172 |
+
"W_fg = H_flat / (H_flat.sum(-1, keepdim=True) + eps)\n\n",
|
| 173 |
+
"W_bg = (1 - H_flat) / (1 - H_flat).sum(-1, keepdim=True)\n\n",
|
| 174 |
+
"\n\n",
|
| 175 |
+
"Pv_norm = Fnn.normalize(Pv, p=2, dim=-1)\n\n",
|
| 176 |
+
"F_proto = Fnn.normalize((W_fg.unsqueeze(-1) * Pv).sum(1), p=2, dim=-1)\n\n",
|
| 177 |
+
"B_proto = Fnn.normalize((W_bg.unsqueeze(-1) * Pv).sum(1), p=2, dim=-1)\n\n",
|
| 178 |
+
"\n\n",
|
| 179 |
+
"sim_fg = torch.bmm(Pv_norm, F_proto.unsqueeze(-1)).squeeze(-1)\n\n",
|
| 180 |
+
"sim_bg = torch.bmm(Pv_norm, B_proto.unsqueeze(-1)).squeeze(-1)\n\n",
|
| 181 |
+
"S_clean = torch.clamp(sim_fg - sim_bg, min=0.0) # contrastive cleaning\n\n",
|
| 182 |
+
"S_knn = knn_affinity_refinement(Pv, S_clean, k=8, temperature=0.1) # kNN affinity propagation\n\n",
|
| 183 |
+
"\n\n",
|
| 184 |
+
"def to_grid(s):\n\n",
|
| 185 |
+
" g = s[0].reshape(H_PATCH, H_PATCH).detach().numpy()\n\n",
|
| 186 |
+
" return (g - g.min()) / (g.max() - g.min() + 1e-8)\n\n",
|
| 187 |
+
"\n\n",
|
| 188 |
+
"fig, ax = plt.subplots(1, 4, figsize=(16, 4))\n\n",
|
| 189 |
+
"ax[0].imshow(gaze_heatmap, cmap=\"jet\"); ax[0].set_title(\"Gaze prior H\")\n\n",
|
| 190 |
+
"ax[1].imshow(to_grid(sim_fg - sim_bg), cmap=\"viridis\"); ax[1].set_title(\"sim_fg \u2212 sim_bg (pre-clean)\")\n\n",
|
| 191 |
+
"ax[2].imshow(to_grid(S_clean), cmap=\"viridis\"); ax[2].set_title(\"After contrastive cleaning\")\n\n",
|
| 192 |
+
"ax[3].imshow(to_grid(S_knn), cmap=\"viridis\"); ax[3].set_title(\"After kNN propagation\")\n\n",
|
| 193 |
+
"for a in ax: a.axis(\"off\")\n\n",
|
| 194 |
+
"plt.tight_layout(); plt.show()\n"
|
| 195 |
+
]
|
| 196 |
+
},
|
| 197 |
+
{
|
| 198 |
+
"cell_type": "markdown",
|
| 199 |
+
"metadata": {},
|
| 200 |
+
"source": [
|
| 201 |
+
"Notice how the gaze prior alone is a soft, blurry blob centered on the clicked points,\n\n",
|
| 202 |
+
"but after contrastive cleaning + kNN propagation the response sharpens to the full extent\n\n",
|
| 203 |
+
"of the (mock-)semantically coherent foreground region \u2014 this is the localization gain\n\n",
|
| 204 |
+
"described in the paper, beyond the initially fixated patches."
|
| 205 |
+
]
|
| 206 |
+
},
|
| 207 |
+
{
|
| 208 |
+
"cell_type": "markdown",
|
| 209 |
+
"metadata": {},
|
| 210 |
+
"source": [
|
| 211 |
+
"## 4. The real pipeline \u2014 frozen DINOv3 + `GazeRefine`\n",
|
| 212 |
+
"\n",
|
| 213 |
+
"This cell calls the actual model end to end. It needs `torch`, `timm`, and internet access\n",
|
| 214 |
+
"the first time (to download the DINOv3 checkpoint). If those aren't available in your current\n",
|
| 215 |
+
"environment, this cell will print instructions instead of failing silently \u2014 run it in a\n",
|
| 216 |
+
"GPU-enabled environment with the repo's `requirements.txt` installed."
|
| 217 |
+
]
|
| 218 |
+
},
|
| 219 |
+
{
|
| 220 |
+
"cell_type": "code",
|
| 221 |
+
"execution_count": null,
|
| 222 |
+
"metadata": {},
|
| 223 |
+
"outputs": [],
|
| 224 |
+
"source": [
|
| 225 |
+
"try:\n\n",
|
| 226 |
+
" from gazerefine import GazeRefine\n\n",
|
| 227 |
+
" from gazerefine.constants import IMG_SIZE\n\n",
|
| 228 |
+
" import torchvision.transforms as T\n\n",
|
| 229 |
+
"\n\n",
|
| 230 |
+
" device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n",
|
| 231 |
+
"\n\n",
|
| 232 |
+
" model = GazeRefine(\n\n",
|
| 233 |
+
" dino_name=\"vit_large_patch16_dinov3.lvd1689m\",\n\n",
|
| 234 |
+
" sigma=2.0, contrast_method=\"difference\", max_iters=5,\n\n",
|
| 235 |
+
" gaze_anchor_weight=0.5, knn_k=20,\n\n",
|
| 236 |
+
" ).to(device).eval()\n\n",
|
| 237 |
+
"\n\n",
|
| 238 |
+
" img_3ch = np.stack([toy_img] * 3, axis=-1)\n\n",
|
| 239 |
+
" tf = T.Compose([T.ToPILImage(), T.Resize((IMG_SIZE, IMG_SIZE)), T.ToTensor(),\n\n",
|
| 240 |
+
" T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])\n\n",
|
| 241 |
+
" img_t = tf((img_3ch * 255).astype(np.uint8)).unsqueeze(0).to(device)\n\n",
|
| 242 |
+
"\n\n",
|
| 243 |
+
" fix_t = torch.tensor([fixations], dtype=torch.float32).to(device)\n\n",
|
| 244 |
+
"\n\n",
|
| 245 |
+
" with torch.no_grad():\n\n",
|
| 246 |
+
" out = model(img_t, fix_t)\n\n",
|
| 247 |
+
"\n\n",
|
| 248 |
+
" pred = out[\"preds\"][0, 0].cpu().numpy()\n\n",
|
| 249 |
+
" fig, ax = plt.subplots(1, 3, figsize=(12, 4))\n\n",
|
| 250 |
+
" ax[0].imshow(toy_img, cmap=\"gray\"); ax[0].set_title(\"Input\")\n\n",
|
| 251 |
+
" ax[1].imshow(out[\"gaze_heatmap\"][0].cpu().numpy(), cmap=\"jet\"); ax[1].set_title(\"Gaze prior\")\n\n",
|
| 252 |
+
" ax[2].imshow(pred > 0.5, cmap=\"gray\"); ax[2].set_title(\"GazeRefine mask\")\n\n",
|
| 253 |
+
" for a in ax: a.axis(\"off\")\n\n",
|
| 254 |
+
" plt.tight_layout(); plt.show()\n\n",
|
| 255 |
+
"\n\n",
|
| 256 |
+
"except ModuleNotFoundError as e:\n\n",
|
| 257 |
+
" print(f\"[skipped] missing dependency: {e}\")\n\n",
|
| 258 |
+
" print(\"Install requirements.txt (torch, timm, ...) and re-run this cell to see real DINOv3 output.\")\n\n",
|
| 259 |
+
"except Exception as e:\n\n",
|
| 260 |
+
" print(f\"[skipped] could not run the real backbone in this environment: {e}\")\n\n",
|
| 261 |
+
" print(\"This is most likely a missing internet connection for the first-time checkpoint download.\")\n"
|
| 262 |
+
]
|
| 263 |
+
},
|
| 264 |
+
{
|
| 265 |
+
"cell_type": "markdown",
|
| 266 |
+
"metadata": {},
|
| 267 |
+
"source": [
|
| 268 |
+
"## 5. Try it on a real image + fixation CSV\n",
|
| 269 |
+
"\n",
|
| 270 |
+
"Once you've added real assets under `examples/` (see `examples/README.md`), point\n",
|
| 271 |
+
"`scripts.predict_single.predict` at them directly:"
|
| 272 |
+
]
|
| 273 |
+
},
|
| 274 |
+
{
|
| 275 |
+
"cell_type": "code",
|
| 276 |
+
"execution_count": null,
|
| 277 |
+
"metadata": {},
|
| 278 |
+
"outputs": [],
|
| 279 |
+
"source": [
|
| 280 |
+
"# from scripts.predict_single import predict\n\n",
|
| 281 |
+
"#\n\n",
|
| 282 |
+
"# out = predict(\n\n",
|
| 283 |
+
"# image=\"../examples/images/kvasir_sample_01.jpg\",\n\n",
|
| 284 |
+
"# fixation_csv=\"../examples/fixations/kvasir_sample_01_fixations.csv\",\n\n",
|
| 285 |
+
"# image_key=\"kvasir_sample_01.jpg\",\n\n",
|
| 286 |
+
"# sigma=2.0, contrast_method=\"difference\", max_iters=5,\n\n",
|
| 287 |
+
"# gaze_anchor_weight=0.5, knn_k=20,\n\n",
|
| 288 |
+
"# )\n\n",
|
| 289 |
+
"# out[\"mask_overlay\"]\n"
|
| 290 |
+
]
|
| 291 |
+
}
|
| 292 |
+
],
|
| 293 |
+
"metadata": {
|
| 294 |
+
"kernelspec": {
|
| 295 |
+
"display_name": "Python 3",
|
| 296 |
+
"language": "python",
|
| 297 |
+
"name": "python3"
|
| 298 |
+
},
|
| 299 |
+
"language_info": {
|
| 300 |
+
"name": "python",
|
| 301 |
+
"pygments_lexer": "ipython3"
|
| 302 |
+
}
|
| 303 |
+
},
|
| 304 |
+
"nbformat": 4,
|
| 305 |
+
"nbformat_minor": 5
|
| 306 |
+
}
|
notebooks/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
MIT License
|
| 2 |
+
|
| 3 |
+
Copyright (c) 2026 GazeRefine Authors
|
| 4 |
+
|
| 5 |
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
| 6 |
+
of this software and associated documentation files (the "Software"), to deal
|
| 7 |
+
in the Software without restriction, including without limitation the rights
|
| 8 |
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
| 9 |
+
copies of the Software, and to permit persons to whom the Software is
|
| 10 |
+
furnished to do so, subject to the following conditions:
|
| 11 |
+
|
| 12 |
+
The above copyright notice and this permission notice shall be included in all
|
| 13 |
+
copies or substantial portions of the Software.
|
| 14 |
+
|
| 15 |
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
| 16 |
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
| 17 |
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
| 18 |
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
| 19 |
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
| 20 |
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
| 21 |
+
SOFTWARE.
|
notebooks/configs/kvasir.yaml
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# GazeRefine config — Kvasir-SEG (colonoscopy polyp segmentation)
|
| 2 |
+
# These are the exact hyperparameters used to obtain the 89.49% Dice
|
| 3 |
+
# reported in Table 1 of the paper.
|
| 4 |
+
|
| 5 |
+
dataset: kvasir
|
| 6 |
+
root: /path/to/Kvasir-SEG # expects <root>/images, <root>/masks
|
| 7 |
+
fixation_csv: /path/to/kvasir_fixation.csv
|
| 8 |
+
|
| 9 |
+
dino_name: vit_large_patch16_dinov3.lvd1689m
|
| 10 |
+
extract_mode: last # "last" | "all"
|
| 11 |
+
|
| 12 |
+
sigma: 2.0 # gaze Gaussian spread, in patch units
|
| 13 |
+
contrast_method: difference # "difference" | "softmax" | "original"
|
| 14 |
+
temperature: 0.05 # only used by contrast_method=softmax
|
| 15 |
+
max_iters: 5 # T, recurrent refinement iterations
|
| 16 |
+
gaze_anchor_weight: 0.5 # lambda, anchoring strength
|
| 17 |
+
|
| 18 |
+
knn_refine: true
|
| 19 |
+
knn_k: 20
|
| 20 |
+
knn_temp: 0.1
|
| 21 |
+
|
| 22 |
+
threshold: 0.5 # binarization threshold
|
| 23 |
+
|
| 24 |
+
batch_size: 4
|
| 25 |
+
num_workers: 4
|
| 26 |
+
save_num: 10
|
| 27 |
+
pred_dir: ./predictions/kvasir
|
notebooks/configs/prostate_mri.yaml
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# GazeRefine config — NCI-ISBI Prostate MRI segmentation
|
| 2 |
+
# These are the exact hyperparameters used to obtain the 75.44% Dice
|
| 3 |
+
# reported in Table 1 of the paper.
|
| 4 |
+
|
| 5 |
+
dataset: prostate_mri
|
| 6 |
+
root: /path/to/Pros_MRI # expects <root>/images (.dcm), <root>/masks (.png)
|
| 7 |
+
fixation_csv: /path/to/nci-isbi_fixation.csv
|
| 8 |
+
|
| 9 |
+
dino_name: vit_large_patch16_dinov3.lvd1689m
|
| 10 |
+
extract_mode: last # "last" | "all"
|
| 11 |
+
|
| 12 |
+
sigma: 1.5 # gaze Gaussian spread, in patch units
|
| 13 |
+
contrast_method: difference # "difference" | "softmax" | "original"
|
| 14 |
+
temperature: 0.05 # only used by contrast_method=softmax
|
| 15 |
+
max_iters: 1 # T, recurrent refinement iterations
|
| 16 |
+
gaze_anchor_weight: 0.8 # lambda, anchoring strength (trust gaze more — low-contrast MRI)
|
| 17 |
+
|
| 18 |
+
knn_refine: true
|
| 19 |
+
knn_k: 3
|
| 20 |
+
knn_temp: 0.1
|
| 21 |
+
|
| 22 |
+
threshold: 0.5 # binarization threshold
|
| 23 |
+
|
| 24 |
+
batch_size: 4
|
| 25 |
+
num_workers: 4
|
| 26 |
+
save_num: 10
|
| 27 |
+
pred_dir: ./predictions/prostate_mri
|
notebooks/gazerefine/__init__.py
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
GazeRefine
|
| 3 |
+
==========
|
| 4 |
+
Expert gaze as a test-time prompt for training-free medical image segmentation.
|
| 5 |
+
|
| 6 |
+
This package exposes the full zero-shot pipeline described in the GazeRefine
|
| 7 |
+
paper: a frozen DINOv3 backbone, gaze -> heatmap conversion, gaze-anchored
|
| 8 |
+
foreground/background prototype construction, and recurrent refinement
|
| 9 |
+
(contrastive background cleaning + kNN affinity propagation).
|
| 10 |
+
|
| 11 |
+
Nothing in this package is trained. There are no learned weights other than
|
| 12 |
+
the frozen, pretrained DINOv3 backbone loaded from `timm`.
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
from .backbone import FrozenDINOv3
|
| 16 |
+
from .gaze import get_scanpath, generate_gaze_heatmap
|
| 17 |
+
from .model import GazeRefine, knn_affinity_refinement
|
| 18 |
+
from .metrics import compute_metrics
|
| 19 |
+
from .visualize import save_prediction, overlay_heatmap, overlay_mask
|
| 20 |
+
|
| 21 |
+
__all__ = [
|
| 22 |
+
"FrozenDINOv3",
|
| 23 |
+
"get_scanpath",
|
| 24 |
+
"generate_gaze_heatmap",
|
| 25 |
+
"GazeRefine",
|
| 26 |
+
"knn_affinity_refinement",
|
| 27 |
+
"compute_metrics",
|
| 28 |
+
"save_prediction",
|
| 29 |
+
"overlay_heatmap",
|
| 30 |
+
"overlay_mask",
|
| 31 |
+
]
|
| 32 |
+
|
| 33 |
+
__version__ = "0.1.0"
|
notebooks/gazerefine/backbone.py
ADDED
|
@@ -0,0 +1,78 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
backbone.py — Frozen DINOv3 feature extractor.
|
| 3 |
+
|
| 4 |
+
GazeRefine never updates the backbone. We load a pretrained DINOv3 ViT via
|
| 5 |
+
`timm`, freeze every parameter, and pull out raw patch tokens from one or
|
| 6 |
+
more transformer blocks using forward hooks. No adapters, no projections,
|
| 7 |
+
no fine-tuning.
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
import torch
|
| 13 |
+
import torch.nn as nn
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
class FrozenDINOv3(nn.Module):
|
| 17 |
+
"""Completely frozen DINOv3 ViT backbone that exposes raw patch tokens.
|
| 18 |
+
|
| 19 |
+
Parameters
|
| 20 |
+
----------
|
| 21 |
+
model_name : str
|
| 22 |
+
Any DINOv3 variant available in `timm` (e.g.
|
| 23 |
+
``"vit_large_patch16_dinov3.lvd1689m"``,
|
| 24 |
+
``"vit_base_patch16_dinov3.lvd1689m"``). Larger backbones generally
|
| 25 |
+
give cleaner semantic separation but cost more memory/compute.
|
| 26 |
+
extract_mode : {"last", "all"}
|
| 27 |
+
- ``"last"``: use only the final block's patch tokens (fast, the
|
| 28 |
+
default used in our reported results).
|
| 29 |
+
- ``"all"``: pool patch tokens from 4 evenly-spaced blocks
|
| 30 |
+
(0, n/4, 3n/4, n-1) and average the resulting similarity maps in
|
| 31 |
+
``GazeRefine``. This sometimes helps on harder modalities at the
|
| 32 |
+
cost of ~4x compute.
|
| 33 |
+
"""
|
| 34 |
+
|
| 35 |
+
def __init__(
|
| 36 |
+
self,
|
| 37 |
+
model_name: str = "vit_large_patch16_dinov3.lvd1689m",
|
| 38 |
+
extract_mode: str = "last",
|
| 39 |
+
):
|
| 40 |
+
super().__init__()
|
| 41 |
+
import timm # local import: keeps `timm` optional for users who only read the code
|
| 42 |
+
|
| 43 |
+
print(f"[GazeRefine] Loading frozen backbone: {model_name} (extract_mode={extract_mode})")
|
| 44 |
+
bb = timm.create_model(model_name, pretrained=True, num_classes=0)
|
| 45 |
+
for p in bb.parameters():
|
| 46 |
+
p.requires_grad_(False)
|
| 47 |
+
bb.eval()
|
| 48 |
+
|
| 49 |
+
self.backbone = bb
|
| 50 |
+
self.embed_dim = bb.embed_dim
|
| 51 |
+
self.num_blocks = len(bb.blocks)
|
| 52 |
+
|
| 53 |
+
if extract_mode == "last":
|
| 54 |
+
self.levels = [-1]
|
| 55 |
+
elif extract_mode == "all":
|
| 56 |
+
self.levels = [0, self.num_blocks // 4, (self.num_blocks * 3) // 4, self.num_blocks - 1]
|
| 57 |
+
else:
|
| 58 |
+
raise ValueError(f"extract_mode must be 'last' or 'all', got {extract_mode!r}")
|
| 59 |
+
|
| 60 |
+
print(f"[GazeRefine] Hooking transformer blocks at levels: {self.levels}")
|
| 61 |
+
self._feats: dict[int, torch.Tensor] = {}
|
| 62 |
+
for lvl in self.levels:
|
| 63 |
+
bb.blocks[lvl].register_forward_hook(self._make_hook(lvl))
|
| 64 |
+
|
| 65 |
+
def _make_hook(self, lvl: int):
|
| 66 |
+
def _hook_fn(module, inp, out):
|
| 67 |
+
# DINOv3 token layout: [CLS, register_1..register_k, patch_1..patch_N]
|
| 68 |
+
n_prefix = 1 + getattr(self.backbone, "num_register_tokens", 4)
|
| 69 |
+
self._feats[lvl] = out[:, n_prefix:, :] # (B, N, D) — patch tokens only
|
| 70 |
+
return _hook_fn
|
| 71 |
+
|
| 72 |
+
@torch.no_grad()
|
| 73 |
+
def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
|
| 74 |
+
"""Run the frozen backbone and return a list of (B, N, D) patch-token tensors,
|
| 75 |
+
one per hooked level, in the order given by ``self.levels``."""
|
| 76 |
+
self._feats.clear()
|
| 77 |
+
self.backbone(x)
|
| 78 |
+
return [self._feats[lvl] for lvl in self.levels]
|
notebooks/gazerefine/constants.py
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""constants.py — shared sizing and normalization constants.
|
| 2 |
+
|
| 3 |
+
IMG_SIZE is fixed to an exact multiple of the ViT patch size so the patch
|
| 4 |
+
grid divides evenly with no rounding/cropping artifacts.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
PATCH_SIZE = 16
|
| 8 |
+
H_PATCH = 37 * 2 # 74 patches per side
|
| 9 |
+
IMG_SIZE = H_PATCH * PATCH_SIZE # 1184 px (use 518 = 37*14 if you switch to a /14 ViT)
|
| 10 |
+
N_PATCHES = H_PATCH ** 2
|
| 11 |
+
|
| 12 |
+
IMG_MEAN = [0.485, 0.456, 0.406]
|
| 13 |
+
IMG_STD = [0.229, 0.224, 0.225]
|
notebooks/gazerefine/datasets.py
ADDED
|
@@ -0,0 +1,154 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
datasets.py — Image/mask/fixation loaders for the two reported benchmarks,
|
| 3 |
+
plus a small base class so adding a third dataset is just a few lines.
|
| 4 |
+
|
| 5 |
+
Expected folder layout (same for both datasets):
|
| 6 |
+
|
| 7 |
+
<root>/
|
| 8 |
+
images/ one image per case (.jpg/.png for Kvasir, .dcm for prostate MRI)
|
| 9 |
+
masks/ matching binary mask, same basename, .png
|
| 10 |
+
<fixation_csv>
|
| 11 |
+
one row per fixation, see gazerefine.gaze for the expected columns.
|
| 12 |
+
The IMAGE column must match an images/ filename (Kvasir) or the
|
| 13 |
+
DICOM basename + ".jpg" (prostate — fixation collection was run on
|
| 14 |
+
JPEG-rendered slices while the model reads the original DICOM).
|
| 15 |
+
"""
|
| 16 |
+
|
| 17 |
+
from __future__ import annotations
|
| 18 |
+
|
| 19 |
+
import os
|
| 20 |
+
|
| 21 |
+
import numpy as np
|
| 22 |
+
import pandas as pd
|
| 23 |
+
import torch
|
| 24 |
+
from PIL import Image
|
| 25 |
+
from torch.utils.data import Dataset
|
| 26 |
+
import torchvision.transforms as T
|
| 27 |
+
|
| 28 |
+
from .constants import IMG_SIZE, IMG_MEAN, IMG_STD
|
| 29 |
+
from .gaze import get_scanpath, IMG_COL
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
class _BaseGazeDataset(Dataset):
|
| 33 |
+
"""Shared image/mask transform + fixation-grouping logic.
|
| 34 |
+
|
| 35 |
+
Subclasses only need to implement ``_load_image(name)`` and provide the
|
| 36 |
+
set of valid image ids that have a matching mask and fixation entries.
|
| 37 |
+
"""
|
| 38 |
+
|
| 39 |
+
def __init__(self, root: str, fixation_csv: str, img_size: int = IMG_SIZE):
|
| 40 |
+
self.root = root
|
| 41 |
+
self.img_dir = os.path.join(root, "images")
|
| 42 |
+
self.mask_dir = os.path.join(root, "masks")
|
| 43 |
+
self.img_size = img_size
|
| 44 |
+
|
| 45 |
+
self.df = pd.read_csv(fixation_csv)
|
| 46 |
+
self.df.columns = self.df.columns.str.strip()
|
| 47 |
+
self.fix_df = self.df.groupby(IMG_COL)
|
| 48 |
+
self.max_len = int(self.df.groupby(IMG_COL).size().max())
|
| 49 |
+
print(f"[Dataset] max scanpath length = {self.max_len}")
|
| 50 |
+
|
| 51 |
+
self.image_ids: list[str] = [] # set by subclass __init__
|
| 52 |
+
|
| 53 |
+
self.img_tf = T.Compose([
|
| 54 |
+
T.Resize((img_size, img_size)),
|
| 55 |
+
T.ToTensor(),
|
| 56 |
+
T.Normalize(IMG_MEAN, IMG_STD),
|
| 57 |
+
])
|
| 58 |
+
self.mask_tf = T.Compose([
|
| 59 |
+
T.Resize((img_size, img_size), interpolation=T.InterpolationMode.NEAREST),
|
| 60 |
+
T.ToTensor(),
|
| 61 |
+
])
|
| 62 |
+
|
| 63 |
+
def __len__(self):
|
| 64 |
+
return len(self.image_ids)
|
| 65 |
+
|
| 66 |
+
def _load_image(self, name: str) -> Image.Image:
|
| 67 |
+
raise NotImplementedError
|
| 68 |
+
|
| 69 |
+
def _mask_path(self, name: str) -> str:
|
| 70 |
+
raise NotImplementedError
|
| 71 |
+
|
| 72 |
+
def _fixation_key(self, name: str) -> str:
|
| 73 |
+
"""CSV IMAGE-column key for this sample. Override if it differs from
|
| 74 |
+
the on-disk basename (e.g. prostate MRI uses .jpg keys for .dcm files)."""
|
| 75 |
+
return name
|
| 76 |
+
|
| 77 |
+
def __getitem__(self, idx: int):
|
| 78 |
+
name = self.image_ids[idx]
|
| 79 |
+
|
| 80 |
+
image = self.img_tf(self._load_image(name))
|
| 81 |
+
|
| 82 |
+
mask = self.mask_tf(Image.open(self._mask_path(name)).convert("L"))
|
| 83 |
+
mask = (mask > 0.5).float()
|
| 84 |
+
|
| 85 |
+
fix_rows = self.fix_df.get_group(self._fixation_key(name))
|
| 86 |
+
fixation = get_scanpath(fix_rows, self.max_len)
|
| 87 |
+
|
| 88 |
+
return {"image": image, "fixation": fixation, "mask": mask, "name": name}
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
class KvasirSEGDataset(_BaseGazeDataset):
|
| 92 |
+
"""Kvasir-SEG colonoscopy polyp segmentation. images/masks share filenames
|
| 93 |
+
(e.g. ``cju0qkwl35piu0993l0dewei2.jpg`` in both folders)."""
|
| 94 |
+
|
| 95 |
+
def __init__(self, root: str, fixation_csv: str, img_size: int = IMG_SIZE):
|
| 96 |
+
super().__init__(root, fixation_csv, img_size)
|
| 97 |
+
img_files = set(os.listdir(self.img_dir))
|
| 98 |
+
mask_files = set(os.listdir(self.mask_dir))
|
| 99 |
+
csv_imgs = set(self.df[IMG_COL].unique())
|
| 100 |
+
self.image_ids = sorted(img_files & mask_files & csv_imgs)
|
| 101 |
+
if not self.image_ids:
|
| 102 |
+
raise RuntimeError("No overlap between images/, masks/ and the fixation CSV.")
|
| 103 |
+
print(f"[Dataset] Kvasir-SEG valid samples = {len(self.image_ids)}")
|
| 104 |
+
|
| 105 |
+
def _load_image(self, name: str) -> Image.Image:
|
| 106 |
+
return Image.open(os.path.join(self.img_dir, name)).convert("RGB")
|
| 107 |
+
|
| 108 |
+
def _mask_path(self, name: str) -> str:
|
| 109 |
+
return os.path.join(self.mask_dir, name)
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
class ProstateMRIDataset(_BaseGazeDataset):
|
| 113 |
+
"""NCI-ISBI prostate MRI. images/ holds DICOM (.dcm), masks/ holds PNG,
|
| 114 |
+
and the fixation CSV references each case as ``<basename>.jpg`` (the
|
| 115 |
+
format the eye-tracking session was actually rendered/displayed in)."""
|
| 116 |
+
|
| 117 |
+
def __init__(self, root: str, fixation_csv: str, img_size: int = IMG_SIZE):
|
| 118 |
+
super().__init__(root, fixation_csv, img_size)
|
| 119 |
+
|
| 120 |
+
dcm_basenames = {os.path.splitext(f)[0] for f in os.listdir(self.img_dir) if f.endswith(".dcm")}
|
| 121 |
+
png_basenames = {os.path.splitext(f)[0] for f in os.listdir(self.mask_dir) if f.endswith(".png")}
|
| 122 |
+
csv_basenames = {os.path.splitext(f)[0] for f in self.df[IMG_COL].unique()}
|
| 123 |
+
|
| 124 |
+
self.image_ids = sorted(dcm_basenames & png_basenames & csv_basenames)
|
| 125 |
+
if not self.image_ids:
|
| 126 |
+
raise RuntimeError("No overlap between images/ (.dcm), masks/ (.png) and the fixation CSV.")
|
| 127 |
+
print(f"[Dataset] Prostate MRI valid samples = {len(self.image_ids)}")
|
| 128 |
+
|
| 129 |
+
def _load_image(self, name: str) -> Image.Image:
|
| 130 |
+
import pydicom # local import: optional dependency, only needed for DICOM datasets
|
| 131 |
+
|
| 132 |
+
dcm_path = os.path.join(self.img_dir, f"{name}.dcm")
|
| 133 |
+
dicom = pydicom.dcmread(dcm_path)
|
| 134 |
+
arr = dicom.pixel_array.astype(np.float32)
|
| 135 |
+
|
| 136 |
+
if arr.max() > arr.min():
|
| 137 |
+
arr = (arr - arr.min()) / (arr.max() - arr.min())
|
| 138 |
+
else:
|
| 139 |
+
arr = np.zeros_like(arr)
|
| 140 |
+
|
| 141 |
+
rgb = np.stack([arr, arr, arr], axis=-1) # grayscale -> 3-channel, DINOv3 expects RGB
|
| 142 |
+
return Image.fromarray((rgb * 255).astype(np.uint8))
|
| 143 |
+
|
| 144 |
+
def _mask_path(self, name: str) -> str:
|
| 145 |
+
return os.path.join(self.mask_dir, f"{name}.png")
|
| 146 |
+
|
| 147 |
+
def _fixation_key(self, name: str) -> str:
|
| 148 |
+
return f"{name}.jpg"
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
DATASET_REGISTRY = {
|
| 152 |
+
"kvasir": KvasirSEGDataset,
|
| 153 |
+
"prostate_mri": ProstateMRIDataset,
|
| 154 |
+
}
|
notebooks/gazerefine/gaze.py
ADDED
|
@@ -0,0 +1,119 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
gaze.py — Turn an eye-tracking scanpath into a spatial prior.
|
| 3 |
+
|
| 4 |
+
Implements Eq. (gaze heatmap) from the paper:
|
| 5 |
+
|
| 6 |
+
H(u, v) = sum_m d_m * exp( -((u - x_m)^2 + (v - y_m)^2) / (2 * sigma^2) )
|
| 7 |
+
|
| 8 |
+
i.e. a duration-weighted sum of 2D Gaussians, one per fixation, evaluated on
|
| 9 |
+
the DINOv3 patch grid. The result is min-max normalized and used directly as
|
| 10 |
+
the soft foreground prior; ``1 - H`` is the background complement.
|
| 11 |
+
|
| 12 |
+
Expected CSV columns (standard EyeLink-style exports, e.g. SR Research /
|
| 13 |
+
Tobii fixation reports):
|
| 14 |
+
|
| 15 |
+
IMAGE image filename the fixation belongs to
|
| 16 |
+
CURRENT_FIX_INDEX fixation order within the trial
|
| 17 |
+
CURRENT_FIX_X fixation x in image-normalized [0, 1] coords
|
| 18 |
+
CURRENT_FIX_Y fixation y in image-normalized [0, 1] coords
|
| 19 |
+
CURRENT_FIX_DURATION fixation duration (ms or any consistent unit)
|
| 20 |
+
|
| 21 |
+
If your fixation coordinates are in raw pixels rather than [0, 1], divide by
|
| 22 |
+
the screen/image width and height before building the CSV — GazeRefine
|
| 23 |
+
expects normalized coordinates so it stays resolution-independent.
|
| 24 |
+
"""
|
| 25 |
+
|
| 26 |
+
from __future__ import annotations
|
| 27 |
+
|
| 28 |
+
import numpy as np
|
| 29 |
+
import pandas as pd
|
| 30 |
+
import torch
|
| 31 |
+
|
| 32 |
+
TIME_COL = "CURRENT_FIX_DURATION"
|
| 33 |
+
FIX_IDX_COL = "CURRENT_FIX_INDEX"
|
| 34 |
+
FIX_X_COL = "CURRENT_FIX_X"
|
| 35 |
+
FIX_Y_COL = "CURRENT_FIX_Y"
|
| 36 |
+
IMG_COL = "IMAGE"
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def get_scanpath(df: pd.DataFrame, max_len: int, time_col: str = TIME_COL) -> torch.Tensor:
|
| 40 |
+
"""Build a zero-padded ``(max_len, 3)`` scanpath tensor from a fixation-report
|
| 41 |
+
DataFrame already filtered down to a single image.
|
| 42 |
+
|
| 43 |
+
Columns of the output: ``[x_norm, y_norm, t_norm]``, all in ``[0, 1]``.
|
| 44 |
+
``t_norm`` is the fixation duration min-max normalized within that image's
|
| 45 |
+
own scanpath (longest fixation -> 1.0). Rows are sorted by
|
| 46 |
+
``CURRENT_FIX_INDEX`` so temporal order is preserved (only the spatial +
|
| 47 |
+
duration channels are currently used by the model, but order is kept for
|
| 48 |
+
future scanpath-aware extensions).
|
| 49 |
+
"""
|
| 50 |
+
if len(df) == 0:
|
| 51 |
+
return torch.zeros(max_len, 3)
|
| 52 |
+
|
| 53 |
+
df = df.sort_values(FIX_IDX_COL)
|
| 54 |
+
x = df[FIX_X_COL].to_numpy(np.float32)
|
| 55 |
+
y = df[FIX_Y_COL].to_numpy(np.float32)
|
| 56 |
+
t = df[time_col].to_numpy(np.float32)
|
| 57 |
+
|
| 58 |
+
t_min, t_max = t.min(), t.max()
|
| 59 |
+
t = (t - t_min) / (t_max - t_min + 1e-8)
|
| 60 |
+
|
| 61 |
+
seq = np.stack([x, y, t], axis=1)
|
| 62 |
+
if len(seq) > max_len:
|
| 63 |
+
seq = seq[:max_len]
|
| 64 |
+
if len(seq) < max_len:
|
| 65 |
+
seq = np.concatenate([seq, np.zeros((max_len - len(seq), 3), np.float32)], axis=0)
|
| 66 |
+
|
| 67 |
+
return torch.tensor(seq, dtype=torch.float32)
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
def generate_gaze_heatmap(
|
| 71 |
+
fixation: torch.Tensor, h_patch: int, sigma: float = 2.0
|
| 72 |
+
) -> torch.Tensor:
|
| 73 |
+
"""Convert a batch of scanpaths into duration-weighted Gaussian heatmaps on
|
| 74 |
+
the DINOv3 patch grid.
|
| 75 |
+
|
| 76 |
+
Parameters
|
| 77 |
+
----------
|
| 78 |
+
fixation : (B, L, 3) tensor of ``[x_norm, y_norm, t_norm]``, zero-padded.
|
| 79 |
+
h_patch : side length of the square patch grid (e.g. 37 for a 518/14 ViT,
|
| 80 |
+
37*2=74 for the 518/16 setup used in our experiments — see
|
| 81 |
+
``gazerefine.model.H_PATCH``).
|
| 82 |
+
sigma : standard deviation (in patch units) of each fixation's Gaussian.
|
| 83 |
+
Larger sigma -> a softer, more diffuse foreground prior; smaller
|
| 84 |
+
sigma -> a tighter prior centered exactly on the fixated patches.
|
| 85 |
+
|
| 86 |
+
Returns
|
| 87 |
+
-------
|
| 88 |
+
(B, h_patch, h_patch) tensor, each sample independently min-max
|
| 89 |
+
normalized to ``[0, 1]`` (an all-zero map stays all-zero if a sample has
|
| 90 |
+
no valid fixations).
|
| 91 |
+
"""
|
| 92 |
+
B, L, _ = fixation.shape
|
| 93 |
+
device = fixation.device
|
| 94 |
+
|
| 95 |
+
y_coords = torch.arange(h_patch, device=device, dtype=torch.float32)
|
| 96 |
+
x_coords = torch.arange(h_patch, device=device, dtype=torch.float32)
|
| 97 |
+
grid_y, grid_x = torch.meshgrid(y_coords, x_coords, indexing="ij")
|
| 98 |
+
|
| 99 |
+
heatmaps = []
|
| 100 |
+
for b in range(B):
|
| 101 |
+
# zero-padding rows are exactly all-zero -> filter them out
|
| 102 |
+
valid = fixation[b].abs().sum(dim=-1) > 1e-8
|
| 103 |
+
valid_fix = fixation[b][valid]
|
| 104 |
+
|
| 105 |
+
h_b = torch.zeros(h_patch, h_patch, device=device)
|
| 106 |
+
if len(valid_fix) > 0:
|
| 107 |
+
for k in range(valid_fix.shape[0]):
|
| 108 |
+
x_k, y_k, t_k = valid_fix[k, 0], valid_fix[k, 1], valid_fix[k, 2]
|
| 109 |
+
gx = x_k * (h_patch - 1)
|
| 110 |
+
gy = y_k * (h_patch - 1)
|
| 111 |
+
dist_sq = (grid_x - gx) ** 2 + (grid_y - gy) ** 2
|
| 112 |
+
h_b += t_k * torch.exp(-dist_sq / (2 * sigma ** 2))
|
| 113 |
+
|
| 114 |
+
h_max = h_b.max()
|
| 115 |
+
if h_max > 1e-8:
|
| 116 |
+
h_b = h_b / h_max
|
| 117 |
+
heatmaps.append(h_b)
|
| 118 |
+
|
| 119 |
+
return torch.stack(heatmaps, dim=0)
|
notebooks/gazerefine/metrics.py
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""metrics.py — Dice / IoU against a binarized prediction."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import torch
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
@torch.no_grad()
|
| 9 |
+
def compute_metrics(preds: torch.Tensor, mask: torch.Tensor, thr: float = 0.5) -> dict:
|
| 10 |
+
"""
|
| 11 |
+
preds, mask : (B, 1, H, W). ``preds`` is the continuous [0, 1] GazeRefine
|
| 12 |
+
output, ``mask`` is the ground-truth binary mask.
|
| 13 |
+
"""
|
| 14 |
+
pred_bin = (preds > thr).float()
|
| 15 |
+
tp = (pred_bin * mask).sum((1, 2, 3))
|
| 16 |
+
fp = (pred_bin * (1 - mask)).sum((1, 2, 3))
|
| 17 |
+
fn = ((1 - pred_bin) * mask).sum((1, 2, 3))
|
| 18 |
+
|
| 19 |
+
dice_per = 2 * tp / (2 * tp + fp + fn + 1e-8)
|
| 20 |
+
iou_per = tp / (tp + fp + fn + 1e-8)
|
| 21 |
+
|
| 22 |
+
return {
|
| 23 |
+
"dice": dice_per.mean().item(),
|
| 24 |
+
"iou": iou_per.mean().item(),
|
| 25 |
+
"dice_per": dice_per,
|
| 26 |
+
"iou_per": iou_per,
|
| 27 |
+
}
|
notebooks/gazerefine/model.py
ADDED
|
@@ -0,0 +1,242 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
model.py — The GazeRefine zero-shot segmentation model.
|
| 3 |
+
|
| 4 |
+
This is a 1:1 refactor of the two task-specific scripts (Kvasir-SEG polyp /
|
| 5 |
+
NCI-ISBI prostate MRI) into a single, dataset-agnostic module. The numerics
|
| 6 |
+
are unchanged from the original experiments — only structure, naming and
|
| 7 |
+
comments were cleaned up. Every per-dataset difference (sigma, kNN k,
|
| 8 |
+
gaze-anchor weight, max_iters, ...) is now a constructor / config argument
|
| 9 |
+
instead of a hardcoded default, see ``configs/*.yaml``.
|
| 10 |
+
|
| 11 |
+
Maps onto the paper (Section 2) as follows:
|
| 12 |
+
|
| 13 |
+
gaze.generate_gaze_heatmap -> H(u, v) (Eq. gaze prior)
|
| 14 |
+
W_fg, W_bg -> W_fg^(0), W_bg^(0) (gaze-derived weights)
|
| 15 |
+
F_proto_init, B_proto_init -> F^(0), B^(0) (initial prototypes)
|
| 16 |
+
contrast_method="difference" -> S_i^(t) = max(0, s_fg - alpha*s_bg) (contrastive cleaning)
|
| 17 |
+
knn_affinity_refinement -> S̄_i^(t) (kNN affinity propagation)
|
| 18 |
+
gaze_anchor_weight -> lambda (anchor blending strength)
|
| 19 |
+
convergence check -> ||F^(t+1) - F^(t)|| < eps
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
from __future__ import annotations
|
| 23 |
+
|
| 24 |
+
import torch
|
| 25 |
+
import torch.nn as nn
|
| 26 |
+
import torch.nn.functional as F
|
| 27 |
+
|
| 28 |
+
from .backbone import FrozenDINOv3
|
| 29 |
+
from .gaze import generate_gaze_heatmap
|
| 30 |
+
from .constants import IMG_SIZE, PATCH_SIZE
|
| 31 |
+
|
| 32 |
+
EPS = 1e-12
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def knn_affinity_refinement(
|
| 36 |
+
Pv: torch.Tensor, S: torch.Tensor, k: int = 5, temperature: float = 0.05
|
| 37 |
+
) -> torch.Tensor:
|
| 38 |
+
"""Propagate a per-patch score map ``S`` across its k nearest neighbors in
|
| 39 |
+
frozen DINOv3 feature space (patch-to-patch affinity), encouraging
|
| 40 |
+
coherent, object-level responses instead of isolated high-confidence
|
| 41 |
+
patches. Corresponds to S̄ in the paper.
|
| 42 |
+
|
| 43 |
+
Pv : (B, N, D) raw (unnormalized) patch embeddings.
|
| 44 |
+
S : (B, N) current per-patch score to be smoothed.
|
| 45 |
+
"""
|
| 46 |
+
B, N, D = Pv.shape
|
| 47 |
+
Pv_norm = F.normalize(Pv, p=2, dim=-1)
|
| 48 |
+
|
| 49 |
+
sim_matrix = torch.bmm(Pv_norm, Pv_norm.transpose(1, 2)) # (B, N, N)
|
| 50 |
+
topk_vals, topk_indices = torch.topk(sim_matrix, k=k, dim=-1) # (B, N, k)
|
| 51 |
+
weights = F.softmax(topk_vals / temperature, dim=-1) # (B, N, k)
|
| 52 |
+
|
| 53 |
+
S_expanded = S.unsqueeze(1).expand(B, N, N)
|
| 54 |
+
S_neighbors = torch.gather(S_expanded, dim=2, index=topk_indices) # (B, N, k)
|
| 55 |
+
|
| 56 |
+
return torch.sum(weights * S_neighbors, dim=-1) # (B, N)
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
class GazeRefine(nn.Module):
|
| 60 |
+
"""Zero-shot, training-free, gaze-guided segmentation model.
|
| 61 |
+
|
| 62 |
+
The model has exactly one set of *learned* weights: the frozen,
|
| 63 |
+
pretrained DINOv3 backbone. Everything else — prototype construction,
|
| 64 |
+
contrastive cleaning, kNN propagation, gaze anchoring — is a closed-form
|
| 65 |
+
operation re-run from scratch on every image at inference time.
|
| 66 |
+
|
| 67 |
+
Parameters
|
| 68 |
+
----------
|
| 69 |
+
dino_name : timm DINOv3 checkpoint name.
|
| 70 |
+
img_size, patch_size : input resolution / ViT patch size. Must satisfy
|
| 71 |
+
``img_size % patch_size == 0``.
|
| 72 |
+
sigma : Gaussian spread (in patch units) for the gaze heatmap.
|
| 73 |
+
extract_mode : ``"last"`` (1 block) or ``"all"`` (4 blocks, averaged).
|
| 74 |
+
contrast_method : ``"difference"`` (paper default, contrastive cleaning),
|
| 75 |
+
``"softmax"`` (foreground/background softmax ratio), or
|
| 76 |
+
``"original"`` (plain foreground cosine similarity, no background
|
| 77 |
+
suppression — kept for the ablation in Table 2).
|
| 78 |
+
temperature : softmax temperature, only used when contrast_method="softmax".
|
| 79 |
+
max_iters : maximum recurrent refinement iterations (T in the paper).
|
| 80 |
+
knn_refine : whether to apply kNN affinity propagation each iteration.
|
| 81 |
+
knn_k, knn_temp : kNN neighborhood size / softmax temperature.
|
| 82 |
+
gaze_anchor_weight : lambda — how strongly each iteration's prototypes
|
| 83 |
+
are pulled back toward the original gaze-only prototypes. Higher =
|
| 84 |
+
trust the raw fixations more; lower = let the model drift further
|
| 85 |
+
from the initial gaze region.
|
| 86 |
+
"""
|
| 87 |
+
|
| 88 |
+
def __init__(
|
| 89 |
+
self,
|
| 90 |
+
dino_name: str = "vit_large_patch16_dinov3.lvd1689m",
|
| 91 |
+
img_size: int = IMG_SIZE,
|
| 92 |
+
patch_size: int = PATCH_SIZE,
|
| 93 |
+
sigma: float = 2.0,
|
| 94 |
+
extract_mode: str = "last",
|
| 95 |
+
contrast_method: str = "difference",
|
| 96 |
+
temperature: float = 0.05,
|
| 97 |
+
max_iters: int = 10,
|
| 98 |
+
knn_refine: bool = True,
|
| 99 |
+
knn_k: int = 5,
|
| 100 |
+
knn_temp: float = 0.1,
|
| 101 |
+
gaze_anchor_weight: float = 0.6,
|
| 102 |
+
):
|
| 103 |
+
super().__init__()
|
| 104 |
+
assert img_size % patch_size == 0, "img_size must be a multiple of patch_size"
|
| 105 |
+
self.img_size = img_size
|
| 106 |
+
self.h_patch = img_size // patch_size
|
| 107 |
+
self.sigma = sigma
|
| 108 |
+
self.contrast_method = contrast_method
|
| 109 |
+
self.temperature = temperature
|
| 110 |
+
self.max_iters = max_iters
|
| 111 |
+
self.knn_refine = knn_refine
|
| 112 |
+
self.knn_k = knn_k
|
| 113 |
+
self.knn_temp = knn_temp
|
| 114 |
+
self.gaze_anchor_weight = gaze_anchor_weight
|
| 115 |
+
|
| 116 |
+
self.visual_enc = FrozenDINOv3(dino_name, extract_mode=extract_mode)
|
| 117 |
+
|
| 118 |
+
# ------------------------------------------------------------------ #
|
| 119 |
+
def _gaze_weights(self, fixation: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
| 120 |
+
"""Scanpath -> heatmap -> normalized foreground/background spatial weights."""
|
| 121 |
+
gaze_heatmap = generate_gaze_heatmap(fixation, h_patch=self.h_patch, sigma=self.sigma)
|
| 122 |
+
H_flat = gaze_heatmap.view(gaze_heatmap.size(0), -1)
|
| 123 |
+
|
| 124 |
+
W_fg = H_flat / (H_flat.sum(dim=-1, keepdim=True) + EPS)
|
| 125 |
+
bg_w = 1.0 - H_flat
|
| 126 |
+
W_bg = bg_w / (bg_w.sum(dim=-1, keepdim=True) + EPS)
|
| 127 |
+
return gaze_heatmap, W_fg, W_bg
|
| 128 |
+
|
| 129 |
+
def _score(self, Pv_norm: torch.Tensor, F_proto_norm: torch.Tensor, B_proto_norm: torch.Tensor) -> torch.Tensor:
|
| 130 |
+
"""Foreground/background contrastive scoring for one level, one iteration."""
|
| 131 |
+
if self.contrast_method == "original":
|
| 132 |
+
S = torch.bmm(Pv_norm, F_proto_norm.unsqueeze(-1)).squeeze(-1)
|
| 133 |
+
return torch.clamp(S, min=0.0)
|
| 134 |
+
|
| 135 |
+
sim_fg = torch.bmm(Pv_norm, F_proto_norm.unsqueeze(-1)).squeeze(-1)
|
| 136 |
+
sim_bg = torch.bmm(Pv_norm, B_proto_norm.unsqueeze(-1)).squeeze(-1)
|
| 137 |
+
|
| 138 |
+
if self.contrast_method == "difference":
|
| 139 |
+
return torch.clamp(sim_fg - sim_bg, min=0.0)
|
| 140 |
+
elif self.contrast_method == "softmax":
|
| 141 |
+
stacked = torch.stack([sim_fg, sim_bg], dim=-1) / self.temperature
|
| 142 |
+
probs = F.softmax(stacked, dim=-1)
|
| 143 |
+
return probs[:, :, 0]
|
| 144 |
+
raise ValueError(f"Unknown contrast_method: {self.contrast_method!r}")
|
| 145 |
+
|
| 146 |
+
def _refine_level(self, Pv: torch.Tensor, W_fg: torch.Tensor, W_bg: torch.Tensor) -> torch.Tensor:
|
| 147 |
+
"""Run the full recurrent gaze-anchored refinement loop for one DINOv3 level.
|
| 148 |
+
Returns the final (B, N) per-patch foreground score map for that level."""
|
| 149 |
+
B = Pv.size(0)
|
| 150 |
+
device = Pv.device
|
| 151 |
+
Pv_norm = F.normalize(Pv, p=2, dim=-1)
|
| 152 |
+
|
| 153 |
+
# Initial gaze-only prototypes — F^(0), B^(0)
|
| 154 |
+
F_proto_init = torch.sum(W_fg.unsqueeze(-1) * Pv, dim=1)
|
| 155 |
+
B_proto_init = torch.sum(W_bg.unsqueeze(-1) * Pv, dim=1)
|
| 156 |
+
|
| 157 |
+
W_fg_curr, W_bg_curr = W_fg.clone(), W_bg.clone()
|
| 158 |
+
best_S = W_fg.clone()
|
| 159 |
+
S_prev = W_fg.clone()
|
| 160 |
+
active = torch.ones(B, dtype=torch.bool, device=device)
|
| 161 |
+
|
| 162 |
+
for _ in range(self.max_iters):
|
| 163 |
+
W_fg_norm = W_fg_curr / (W_fg_curr.sum(dim=-1, keepdim=True) + EPS)
|
| 164 |
+
W_bg_norm = W_bg_curr / (W_bg_curr.sum(dim=-1, keepdim=True) + EPS)
|
| 165 |
+
|
| 166 |
+
F_proto_curr = torch.sum(W_fg_norm.unsqueeze(-1) * Pv, dim=1)
|
| 167 |
+
B_proto_curr = torch.sum(W_bg_norm.unsqueeze(-1) * Pv, dim=1)
|
| 168 |
+
|
| 169 |
+
# Anchor toward the initial gaze-only prototypes — lambda blending
|
| 170 |
+
lam = self.gaze_anchor_weight
|
| 171 |
+
F_proto = (1 - lam) * F_proto_curr + lam * F_proto_init
|
| 172 |
+
B_proto = (1 - lam) * B_proto_curr + lam * B_proto_init
|
| 173 |
+
F_proto_norm = F.normalize(F_proto, p=2, dim=-1)
|
| 174 |
+
B_proto_norm = F.normalize(B_proto, p=2, dim=-1)
|
| 175 |
+
|
| 176 |
+
S_iter = self._score(Pv_norm, F_proto_norm, B_proto_norm)
|
| 177 |
+
if self.knn_refine:
|
| 178 |
+
S_iter = knn_affinity_refinement(Pv, S_iter, k=self.knn_k, temperature=self.knn_temp)
|
| 179 |
+
|
| 180 |
+
S_new = S_prev.clone()
|
| 181 |
+
S_new[active] = S_iter[active]
|
| 182 |
+
|
| 183 |
+
# Collapse prevention: if an active sample's map flattened to ~0,
|
| 184 |
+
# roll it back and freeze it instead of letting it degenerate.
|
| 185 |
+
collapsed = S_new.max(dim=-1)[0] < 1e-4
|
| 186 |
+
freeze = active & collapsed
|
| 187 |
+
S_new[freeze] = S_prev[freeze]
|
| 188 |
+
active = active & ~collapsed
|
| 189 |
+
|
| 190 |
+
best_S = S_new.clone()
|
| 191 |
+
if not active.any():
|
| 192 |
+
break
|
| 193 |
+
|
| 194 |
+
# Turn the refined score into the next iteration's spatial weights
|
| 195 |
+
S_sig = torch.sigmoid(S_new)
|
| 196 |
+
W_fg_next = S_sig / (S_sig.sum(dim=-1, keepdim=True) + EPS)
|
| 197 |
+
bg_w = 1.0 - S_sig
|
| 198 |
+
bg_w = torch.clamp(bg_w - bg_w.min(dim=-1, keepdim=True).values, min=0.0)
|
| 199 |
+
W_bg_next = bg_w / (bg_w.sum(dim=-1, keepdim=True) + EPS)
|
| 200 |
+
|
| 201 |
+
F_proto_next = torch.sum(W_fg_next.unsqueeze(-1) * Pv, dim=1)
|
| 202 |
+
F_proto_next_norm = F.normalize(F_proto_next, p=2, dim=-1)
|
| 203 |
+
|
| 204 |
+
# Convergence: foreground prototype stopped moving
|
| 205 |
+
diff = torch.abs(F_proto_next_norm - F_proto_norm).mean(dim=-1)
|
| 206 |
+
if torch.all(diff < 1e-6):
|
| 207 |
+
break
|
| 208 |
+
|
| 209 |
+
S_prev = S_new.clone()
|
| 210 |
+
W_fg_curr[active] = W_fg_next[active]
|
| 211 |
+
W_bg_curr[active] = W_bg_next[active]
|
| 212 |
+
|
| 213 |
+
return best_S
|
| 214 |
+
|
| 215 |
+
# ------------------------------------------------------------------ #
|
| 216 |
+
def forward(self, image: torch.Tensor, fixation: torch.Tensor) -> dict:
|
| 217 |
+
"""
|
| 218 |
+
image : (B, 3, H, W) normalized RGB tensor (ImageNet mean/std).
|
| 219 |
+
fixation : (B, L, 3) zero-padded scanpath, see ``gazerefine.gaze.get_scanpath``.
|
| 220 |
+
|
| 221 |
+
Returns a dict with:
|
| 222 |
+
preds (B, 1, H, W) final mask in [0, 1] — threshold at 0.5 for a binary mask
|
| 223 |
+
gaze_heatmap (B, h, w) the raw gaze prior, useful for visualization
|
| 224 |
+
final_map (B, N) the un-upsampled patch-level score map
|
| 225 |
+
"""
|
| 226 |
+
B = image.size(0)
|
| 227 |
+
Pv_list = self.visual_enc(image)
|
| 228 |
+
gaze_heatmap, W_fg, W_bg = self._gaze_weights(fixation)
|
| 229 |
+
|
| 230 |
+
all_maps = [self._refine_level(Pv, W_fg, W_bg) for Pv in Pv_list]
|
| 231 |
+
final_map = torch.stack(all_maps, dim=1).mean(dim=1) # average across levels
|
| 232 |
+
|
| 233 |
+
grid = final_map.view(B, 1, self.h_patch, self.h_patch)
|
| 234 |
+
upsampled = F.interpolate(grid, size=(self.img_size, self.img_size), mode="bilinear", align_corners=False)
|
| 235 |
+
|
| 236 |
+
# per-sample min-max normalization -> ready-to-threshold mask
|
| 237 |
+
flat = upsampled.view(B, -1)
|
| 238 |
+
m_min = flat.min(dim=-1, keepdim=True).values.view(B, 1, 1, 1)
|
| 239 |
+
m_max = flat.max(dim=-1, keepdim=True).values.view(B, 1, 1, 1)
|
| 240 |
+
preds = (upsampled - m_min) / (m_max - m_min + 1e-8)
|
| 241 |
+
|
| 242 |
+
return dict(preds=preds, gaze_heatmap=gaze_heatmap, final_map=final_map)
|
notebooks/gazerefine/visualize.py
ADDED
|
@@ -0,0 +1,86 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""visualize.py — Plotting / overlay helpers.
|
| 2 |
+
|
| 3 |
+
``save_prediction`` reproduces the 4-panel (image / gaze / GT / prediction)
|
| 4 |
+
grid used during batch evaluation. ``overlay_heatmap`` and ``overlay_mask``
|
| 5 |
+
are lighter-weight, matplotlib-free PIL helpers meant for the Gradio demo
|
| 6 |
+
and the example notebook, where you typically want a single composited
|
| 7 |
+
image rather than a saved subplot file.
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
import os
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import torch
|
| 16 |
+
from PIL import Image
|
| 17 |
+
|
| 18 |
+
import matplotlib
|
| 19 |
+
matplotlib.use("Agg")
|
| 20 |
+
import matplotlib.pyplot as plt
|
| 21 |
+
import matplotlib.cm as cm
|
| 22 |
+
|
| 23 |
+
from .constants import IMG_MEAN, IMG_STD
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def _denormalize(img: torch.Tensor) -> np.ndarray:
|
| 27 |
+
"""(3, H, W) normalized tensor -> (H, W, 3) uint8-range float array in [0, 1]."""
|
| 28 |
+
mean = torch.tensor(IMG_MEAN, device=img.device).view(3, 1, 1)
|
| 29 |
+
std = torch.tensor(IMG_STD, device=img.device).view(3, 1, 1)
|
| 30 |
+
return (img * std + mean).clamp(0, 1).permute(1, 2, 0).cpu().numpy()
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
@torch.no_grad()
|
| 34 |
+
def save_prediction(
|
| 35 |
+
img: torch.Tensor,
|
| 36 |
+
gaze_heatmap: torch.Tensor,
|
| 37 |
+
gt: torch.Tensor,
|
| 38 |
+
pred: torch.Tensor,
|
| 39 |
+
name: str,
|
| 40 |
+
out_dir: str = "predictions",
|
| 41 |
+
) -> None:
|
| 42 |
+
"""Save a 4-panel [image | gaze heatmap | GT mask | predicted mask] figure."""
|
| 43 |
+
os.makedirs(out_dir, exist_ok=True)
|
| 44 |
+
|
| 45 |
+
img_np = _denormalize(img)
|
| 46 |
+
gaze_np = gaze_heatmap.squeeze().cpu().numpy()
|
| 47 |
+
gt_np = gt.squeeze().cpu().numpy()
|
| 48 |
+
pred_np = pred.squeeze().cpu().numpy()
|
| 49 |
+
|
| 50 |
+
fig, ax = plt.subplots(1, 4, figsize=(20, 5))
|
| 51 |
+
ax[0].imshow(img_np); ax[0].set_title("Image")
|
| 52 |
+
ax[1].imshow(gaze_np, cmap="jet"); ax[1].set_title("Gaze Heatmap")
|
| 53 |
+
ax[2].imshow(gt_np, cmap="gray"); ax[2].set_title("GT mask")
|
| 54 |
+
ax[3].imshow(pred_np, cmap="gray"); ax[3].set_title("Predicted mask")
|
| 55 |
+
for a in ax:
|
| 56 |
+
a.axis("off")
|
| 57 |
+
plt.tight_layout()
|
| 58 |
+
plt.savefig(os.path.join(out_dir, name), dpi=100, bbox_inches="tight")
|
| 59 |
+
plt.close()
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def overlay_heatmap(image: Image.Image, heatmap: np.ndarray, alpha: float = 0.45) -> Image.Image:
|
| 63 |
+
"""Composite a [0, 1] heatmap (any H'xW', will be resized) onto a PIL image
|
| 64 |
+
using a jet colormap. Used to show the gaze prior over the original image."""
|
| 65 |
+
heatmap = np.asarray(heatmap, dtype=np.float32)
|
| 66 |
+
heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min() + 1e-8)
|
| 67 |
+
|
| 68 |
+
heat_img = Image.fromarray((cm.jet(heatmap)[:, :, :3] * 255).astype(np.uint8))
|
| 69 |
+
heat_img = heat_img.resize(image.size, resample=Image.BILINEAR)
|
| 70 |
+
|
| 71 |
+
base = image.convert("RGB")
|
| 72 |
+
return Image.blend(base, heat_img, alpha=alpha)
|
| 73 |
+
|
| 74 |
+
|
| 75 |
+
def overlay_mask(
|
| 76 |
+
image: Image.Image, mask: np.ndarray, color: tuple[int, int, int] = (255, 60, 60), alpha: float = 0.45
|
| 77 |
+
) -> Image.Image:
|
| 78 |
+
"""Composite a binary/soft [0, 1] mask onto a PIL image as a solid color wash."""
|
| 79 |
+
mask = np.asarray(mask, dtype=np.float32)
|
| 80 |
+
mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)
|
| 81 |
+
mask_img = Image.fromarray((mask * 255).astype(np.uint8)).resize(image.size, resample=Image.NEAREST)
|
| 82 |
+
|
| 83 |
+
base = image.convert("RGB")
|
| 84 |
+
color_layer = Image.new("RGB", base.size, color)
|
| 85 |
+
composited = Image.composite(color_layer, base, mask_img.point(lambda p: int(p * alpha)))
|
| 86 |
+
return composited
|
notebooks/huggingface_space/README.md
ADDED
|
@@ -0,0 +1,42 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
title: GazeRefine
|
| 3 |
+
emoji: 👁️
|
| 4 |
+
colorFrom: blue
|
| 5 |
+
colorTo: red
|
| 6 |
+
sdk: gradio
|
| 7 |
+
sdk_version: 4.44.0
|
| 8 |
+
app_file: app.py
|
| 9 |
+
pinned: false
|
| 10 |
+
license: mit
|
| 11 |
+
short_description: Zero-shot, training-free gaze-guided medical segmentation
|
| 12 |
+
---
|
| 13 |
+
|
| 14 |
+
# GazeRefine — Expert Gaze as a Test-Time Prompt
|
| 15 |
+
|
| 16 |
+
Interactive demo for **GazeRefine**, a training-free, zero-shot framework
|
| 17 |
+
that turns expert eye-gaze into an inference-time prompt for medical image
|
| 18 |
+
segmentation. Frozen DINOv3 patch features + gaze-anchored
|
| 19 |
+
foreground/background prototypes + recurrent contrastive cleaning + kNN
|
| 20 |
+
affinity propagation — no masks, no clicks-as-boxes, no fine-tuning, no
|
| 21 |
+
adapters, no prompt encoder.
|
| 22 |
+
|
| 23 |
+
## How to use
|
| 24 |
+
1. Upload a colonoscopy or grayscale-MRI-style image.
|
| 25 |
+
2. Click on the image 1–5 times where a clinician's gaze would land on the
|
| 26 |
+
structure of interest (a polyp, the prostate, ...). Each click adds a
|
| 27 |
+
numbered fixation marker; the slider controls that fixation's relative
|
| 28 |
+
duration/weight before your next click.
|
| 29 |
+
3. Pick a hyperparameter preset (tuned per-modality, see the paper).
|
| 30 |
+
4. Press **Run GazeRefine** to get the gaze-prior overlay and the predicted
|
| 31 |
+
segmentation mask.
|
| 32 |
+
|
| 33 |
+
## Notes
|
| 34 |
+
- Inference uses a frozen `vit_large_patch16_dinov3.lvd1689m` backbone from
|
| 35 |
+
`timm`. First run will download the checkpoint.
|
| 36 |
+
- CPU inference works but is slow; a GPU Space is recommended for a smooth
|
| 37 |
+
demo.
|
| 38 |
+
- This Space is for research/demonstration only — it is **not** a clinical
|
| 39 |
+
diagnostic tool.
|
| 40 |
+
|
| 41 |
+
Full code, configs, and the unified evaluation pipeline:
|
| 42 |
+
[GitHub repository](https://github.com/<your-org>/gazerefine).
|
notebooks/huggingface_space/app.py
ADDED
|
@@ -0,0 +1,138 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
app.py — GazeRefine interactive demo (Hugging Face Space).
|
| 3 |
+
|
| 4 |
+
Upload a medical image, click on it to drop a few fixation points the way a
|
| 5 |
+
clinician's gaze would land on the structure of interest, and GazeRefine
|
| 6 |
+
turns that into a segmentation mask — fully zero-shot, no training, no
|
| 7 |
+
prompt-specific architecture.
|
| 8 |
+
|
| 9 |
+
Run locally with:
|
| 10 |
+
pip install -r huggingface_space/requirements.txt
|
| 11 |
+
python huggingface_space/app.py
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
from __future__ import annotations
|
| 15 |
+
|
| 16 |
+
import sys
|
| 17 |
+
from pathlib import Path
|
| 18 |
+
|
| 19 |
+
import gradio as gr
|
| 20 |
+
import numpy as np
|
| 21 |
+
from PIL import Image, ImageDraw
|
| 22 |
+
|
| 23 |
+
sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # repo root
|
| 24 |
+
from scripts.predict_single import predict
|
| 25 |
+
|
| 26 |
+
PRESETS = {
|
| 27 |
+
"Colonoscopy / polyp (Kvasir-SEG settings)": dict(
|
| 28 |
+
sigma=2.0, contrast_method="difference", max_iters=5,
|
| 29 |
+
gaze_anchor_weight=0.5, knn_k=20, knn_temp=0.1,
|
| 30 |
+
),
|
| 31 |
+
"Grayscale MRI / CT (prostate-MRI settings)": dict(
|
| 32 |
+
sigma=1.5, contrast_method="difference", max_iters=1,
|
| 33 |
+
gaze_anchor_weight=0.8, knn_k=3, knn_temp=0.1,
|
| 34 |
+
),
|
| 35 |
+
}
|
| 36 |
+
DINO_NAME = "vit_large_patch16_dinov3.lvd1689m"
|
| 37 |
+
|
| 38 |
+
POINT_COLORS = ["#ff3b30", "#ff9500", "#ffcc00", "#34c759", "#5ac8fa", "#007aff", "#af52de"]
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def draw_points(image: Image.Image, points: list[tuple[float, float, float]]) -> Image.Image:
|
| 42 |
+
"""Render numbered fixation markers over the image for visual feedback."""
|
| 43 |
+
if image is None:
|
| 44 |
+
return None
|
| 45 |
+
vis = image.convert("RGB").copy()
|
| 46 |
+
draw = ImageDraw.Draw(vis)
|
| 47 |
+
w, h = vis.size
|
| 48 |
+
r = max(6, min(w, h) // 80)
|
| 49 |
+
for i, (x, y, dur) in enumerate(points):
|
| 50 |
+
cx, cy = x * w, y * h
|
| 51 |
+
color = POINT_COLORS[i % len(POINT_COLORS)]
|
| 52 |
+
rad = r * (0.6 + 0.8 * dur) # bigger marker = longer fixation
|
| 53 |
+
draw.ellipse([cx - rad, cy - rad, cx + rad, cy + rad], outline=color, width=3)
|
| 54 |
+
draw.text((cx + rad + 2, cy - rad), str(i + 1), fill=color)
|
| 55 |
+
return vis
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def on_select(image: Image.Image, points: list, duration: float, evt: gr.SelectData):
|
| 59 |
+
if image is None:
|
| 60 |
+
gr.Warning("Upload an image first.")
|
| 61 |
+
return points, None
|
| 62 |
+
w, h = image.size
|
| 63 |
+
x_px, y_px = evt.index
|
| 64 |
+
points = points + [(x_px / w, y_px / h, duration)]
|
| 65 |
+
return points, draw_points(image, points)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def on_clear(image: Image.Image):
|
| 69 |
+
return [], image
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
def on_image_change(image: Image.Image):
|
| 73 |
+
# new image -> reset fixations
|
| 74 |
+
return [], image
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
def run(image: Image.Image, points: list, preset_name: str, threshold: float):
|
| 78 |
+
if image is None:
|
| 79 |
+
gr.Warning("Upload an image first.")
|
| 80 |
+
return None, None, None
|
| 81 |
+
if len(points) == 0:
|
| 82 |
+
gr.Warning("Click on the image at least once to place a fixation.")
|
| 83 |
+
return None, None, None
|
| 84 |
+
|
| 85 |
+
cfg = PRESETS[preset_name]
|
| 86 |
+
out = predict(
|
| 87 |
+
image=image,
|
| 88 |
+
fixations=points,
|
| 89 |
+
dino_name=DINO_NAME,
|
| 90 |
+
threshold=threshold,
|
| 91 |
+
**cfg,
|
| 92 |
+
)
|
| 93 |
+
mask_only = Image.fromarray((np.clip(out["mask_bin"], 0, 1) * 255).astype(np.uint8))
|
| 94 |
+
return out["gaze_overlay"], out["mask_overlay"], mask_only
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
with gr.Blocks(title="GazeRefine — gaze-guided zero-shot segmentation") as demo:
|
| 98 |
+
gr.Markdown(
|
| 99 |
+
"""
|
| 100 |
+
# 👁️ GazeRefine — Expert Gaze as a Test-Time Prompt
|
| 101 |
+
Training-free, zero-shot medical image segmentation. Upload an image, **click on it
|
| 102 |
+
1–5 times** where a clinician would look, pick a preset, and run.
|
| 103 |
+
No masks, clicks-as-bounding-boxes, fine-tuning, or adapters — just a frozen DINOv3
|
| 104 |
+
backbone steered by your gaze. See the paper / code on GitHub (linked below).
|
| 105 |
+
"""
|
| 106 |
+
)
|
| 107 |
+
|
| 108 |
+
points_state = gr.State([])
|
| 109 |
+
|
| 110 |
+
with gr.Row():
|
| 111 |
+
with gr.Column():
|
| 112 |
+
image_in = gr.Image(type="pil", label="1. Upload image, then click to place fixations", height=420)
|
| 113 |
+
with gr.Row():
|
| 114 |
+
duration_slider = gr.Slider(0.1, 1.0, value=1.0, step=0.1, label="Next fixation's relative duration")
|
| 115 |
+
clear_btn = gr.Button("Clear fixations")
|
| 116 |
+
preset = gr.Radio(list(PRESETS.keys()), value=list(PRESETS.keys())[0], label="2. Hyperparameter preset")
|
| 117 |
+
threshold = gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="3. Mask threshold")
|
| 118 |
+
run_btn = gr.Button("Run GazeRefine", variant="primary")
|
| 119 |
+
|
| 120 |
+
with gr.Column():
|
| 121 |
+
gaze_out = gr.Image(label="Gaze prior over image", height=260)
|
| 122 |
+
with gr.Row():
|
| 123 |
+
mask_overlay_out = gr.Image(label="Predicted mask (overlay)", height=260)
|
| 124 |
+
mask_only_out = gr.Image(label="Predicted mask (binary)", height=260)
|
| 125 |
+
|
| 126 |
+
image_in.upload(on_image_change, inputs=[image_in], outputs=[points_state, image_in])
|
| 127 |
+
image_in.select(on_select, inputs=[image_in, points_state, duration_slider], outputs=[points_state, image_in])
|
| 128 |
+
clear_btn.click(on_clear, inputs=[image_in], outputs=[points_state, image_in])
|
| 129 |
+
run_btn.click(run, inputs=[image_in, points_state, preset, threshold], outputs=[gaze_out, mask_overlay_out, mask_only_out])
|
| 130 |
+
|
| 131 |
+
gr.Markdown(
|
| 132 |
+
"Code: [GitHub](https://github.com/<your-org>/gazerefine) · "
|
| 133 |
+
"Method: GazeRefine — frozen DINOv3 + gaze-anchored prototypes + recurrent "
|
| 134 |
+
"foreground/background refinement, entirely training-free."
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
if __name__ == "__main__":
|
| 138 |
+
demo.launch()
|
notebooks/huggingface_space/requirements.txt
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
torch>=2.1
|
| 2 |
+
torchvision>=0.16
|
| 3 |
+
timm>=1.0.0
|
| 4 |
+
numpy
|
| 5 |
+
pandas
|
| 6 |
+
pillow
|
| 7 |
+
gradio>=4.0
|
notebooks/requirements.txt
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
torch>=2.1
|
| 2 |
+
torchvision>=0.16
|
| 3 |
+
timm>=1.0.0 # provides DINOv3 checkpoints (vit_*_dinov3.*)
|
| 4 |
+
numpy
|
| 5 |
+
pandas
|
| 6 |
+
pillow
|
| 7 |
+
pyyaml
|
| 8 |
+
matplotlib
|
| 9 |
+
pydicom # only required for the prostate MRI (DICOM) dataset
|
requirements.txt
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
torch>=2.1
|
| 2 |
+
torchvision>=0.16
|
| 3 |
+
timm>=1.0.0
|
| 4 |
+
numpy
|
| 5 |
+
pandas
|
| 6 |
+
pillow
|
| 7 |
+
matplotlib
|
| 8 |
+
pydicom
|
| 9 |
+
huggingface-hub==0.24.0
|
| 10 |
+
pydub==0.25.1
|
scripts/predict_single.py
ADDED
|
@@ -0,0 +1,311 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
predict_single.py — Run GazeRefine on a single image + fixation CSV.
|
| 3 |
+
|
| 4 |
+
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
|
| 5 |
+
Command-line usage (matches the README exactly)
|
| 6 |
+
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
|
| 7 |
+
|
| 8 |
+
python scripts/predict_single.py \\
|
| 9 |
+
--image examples/image.png \\
|
| 10 |
+
--fixations examples/fixations.csv \\
|
| 11 |
+
--output output_mask.png
|
| 12 |
+
|
| 13 |
+
Optional flags:
|
| 14 |
+
--preset colonoscopy | mri (default: colonoscopy)
|
| 15 |
+
--threshold 0.5 binarization threshold
|
| 16 |
+
--save_overlay also save a colour overlay PNG
|
| 17 |
+
--device cuda | cpu (auto-detected by default)
|
| 18 |
+
|
| 19 |
+
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
|
| 20 |
+
Python API (matches the README exactly)
|
| 21 |
+
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
|
| 22 |
+
|
| 23 |
+
from scripts.predict_single import predict
|
| 24 |
+
|
| 25 |
+
mask = predict(
|
| 26 |
+
image_path="image.png",
|
| 27 |
+
fixation_csv="fixations.csv",
|
| 28 |
+
)
|
| 29 |
+
# `mask` is a PIL Image of the binary segmentation mask.
|
| 30 |
+
# Save it:
|
| 31 |
+
mask.save("output_mask.png")
|
| 32 |
+
|
| 33 |
+
# Extended API — also get overlays and raw arrays:
|
| 34 |
+
result = predict(
|
| 35 |
+
image_path="image.png",
|
| 36 |
+
fixation_csv="fixations.csv",
|
| 37 |
+
preset="mri", # "colonoscopy" (default) or "mri"
|
| 38 |
+
threshold=0.5,
|
| 39 |
+
return_all=True,
|
| 40 |
+
)
|
| 41 |
+
result["mask"].save("mask.png")
|
| 42 |
+
result["gaze_overlay"].save("gaze.png")
|
| 43 |
+
result["mask_overlay"].save("overlay.png")
|
| 44 |
+
|
| 45 |
+
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
|
| 46 |
+
Fixation CSV format
|
| 47 |
+
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
|
| 48 |
+
|
| 49 |
+
x,y,duration
|
| 50 |
+
340,221,180
|
| 51 |
+
356,228,145
|
| 52 |
+
368,244,205
|
| 53 |
+
...
|
| 54 |
+
|
| 55 |
+
x, y — fixation position in *raw pixel* coordinates of the input image.
|
| 56 |
+
(These are automatically normalized by the image size internally.)
|
| 57 |
+
duration — fixation duration in any consistent unit (milliseconds typical).
|
| 58 |
+
The model only uses *relative* durations, so the unit does not matter.
|
| 59 |
+
"""
|
| 60 |
+
|
| 61 |
+
from __future__ import annotations
|
| 62 |
+
|
| 63 |
+
import argparse
|
| 64 |
+
import sys
|
| 65 |
+
from pathlib import Path
|
| 66 |
+
from typing import Sequence
|
| 67 |
+
|
| 68 |
+
import numpy as np
|
| 69 |
+
import pandas as pd
|
| 70 |
+
import torch
|
| 71 |
+
from PIL import Image
|
| 72 |
+
|
| 73 |
+
# allow `python scripts/predict_single.py` from any working directory
|
| 74 |
+
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
| 75 |
+
|
| 76 |
+
from gazerefine import GazeRefine, overlay_heatmap, overlay_mask
|
| 77 |
+
from gazerefine.constants import IMG_MEAN, IMG_STD, IMG_SIZE
|
| 78 |
+
from gazerefine.gaze import load_fixation_csv
|
| 79 |
+
import torchvision.transforms as T
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
# ── per-modality hyperparameter presets ────────────────────────────────────
|
| 83 |
+
# These match the exact settings used to produce the paper's Table 1 numbers.
|
| 84 |
+
PRESETS: dict[str, dict] = {
|
| 85 |
+
"colonoscopy": dict(
|
| 86 |
+
sigma=2.0,
|
| 87 |
+
contrast_method="difference",
|
| 88 |
+
max_iters=5,
|
| 89 |
+
gaze_anchor_weight=0.5,
|
| 90 |
+
knn_refine=True,
|
| 91 |
+
knn_k=20,
|
| 92 |
+
knn_temp=0.1,
|
| 93 |
+
),
|
| 94 |
+
"mri": dict(
|
| 95 |
+
sigma=1.5,
|
| 96 |
+
contrast_method="difference",
|
| 97 |
+
max_iters=1,
|
| 98 |
+
gaze_anchor_weight=0.8,
|
| 99 |
+
knn_refine=True,
|
| 100 |
+
knn_k=3,
|
| 101 |
+
knn_temp=0.1,
|
| 102 |
+
),
|
| 103 |
+
}
|
| 104 |
+
|
| 105 |
+
# one shared backbone name for both presets
|
| 106 |
+
DINO_NAME = "vit_large_patch16_dinov3.lvd1689m"
|
| 107 |
+
|
| 108 |
+
# module-level model cache: avoids reloading the backbone across repeated calls
|
| 109 |
+
# (useful when this module is imported by the Gradio Space or a notebook loop)
|
| 110 |
+
_MODEL_CACHE: dict[str, GazeRefine] = {}
|
| 111 |
+
|
| 112 |
+
|
| 113 |
+
def _get_model(preset: str, device: torch.device) -> GazeRefine:
|
| 114 |
+
"""Load (or return a cached) GazeRefine model for the given preset."""
|
| 115 |
+
if preset not in _MODEL_CACHE:
|
| 116 |
+
cfg = PRESETS[preset]
|
| 117 |
+
_MODEL_CACHE[preset] = GazeRefine(dino_name=DINO_NAME, **cfg).to(device).eval()
|
| 118 |
+
return _MODEL_CACHE[preset]
|
| 119 |
+
|
| 120 |
+
|
| 121 |
+
# ── main public function ────────────────────────────────────────────────────
|
| 122 |
+
import os
|
| 123 |
+
import pydicom
|
| 124 |
+
import numpy as np
|
| 125 |
+
from PIL import Image
|
| 126 |
+
@torch.no_grad()
|
| 127 |
+
def predict(
|
| 128 |
+
image_path: "str | Path | Image.Image",
|
| 129 |
+
fixation_csv: "str | Path",
|
| 130 |
+
preset: str = "colonoscopy",
|
| 131 |
+
threshold: float = 0.5,
|
| 132 |
+
device: str | None = None,
|
| 133 |
+
return_all: bool = False,
|
| 134 |
+
) -> "Image.Image | dict":
|
| 135 |
+
"""Run GazeRefine on one image and return the predicted segmentation mask.
|
| 136 |
+
|
| 137 |
+
Parameters
|
| 138 |
+
----------
|
| 139 |
+
image_path : path to the input image (.jpg / .jpeg / .png) **or** an
|
| 140 |
+
already-loaded ``PIL.Image`` (used by the Gradio Space).
|
| 141 |
+
fixation_csv : path to the fixation CSV (``x,y,duration`` columns,
|
| 142 |
+
pixel coordinates — see module docstring for the format).
|
| 143 |
+
preset : ``"colonoscopy"`` (default, Kvasir-SEG settings) or
|
| 144 |
+
``"mri"`` (NCI-ISBI prostate-MRI settings).
|
| 145 |
+
threshold : binarization cutoff applied to the [0, 1] soft mask.
|
| 146 |
+
device : ``"cuda"`` / ``"cpu"`` — auto-detected when ``None``.
|
| 147 |
+
return_all : when ``True``, return a dict with the binary mask PIL Image
|
| 148 |
+
**plus** ``gaze_overlay``, ``mask_overlay``, and the raw
|
| 149 |
+
numpy arrays ``preds`` and ``gaze_heatmap``.
|
| 150 |
+
When ``False`` (default), return only the mask PIL Image.
|
| 151 |
+
|
| 152 |
+
Returns
|
| 153 |
+
-------
|
| 154 |
+
``PIL.Image`` of the binary mask, **or** a dict (see ``return_all``).
|
| 155 |
+
"""
|
| 156 |
+
if preset not in PRESETS:
|
| 157 |
+
raise ValueError(f"preset must be one of {list(PRESETS)}, got {preset!r}")
|
| 158 |
+
|
| 159 |
+
_device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
|
| 160 |
+
|
| 161 |
+
# ── load the image ──
|
| 162 |
+
ext = os.path.splitext(str(image_path))[1].lower()
|
| 163 |
+
|
| 164 |
+
if isinstance(image_path, Image.Image):
|
| 165 |
+
pil_image = image_path.convert("RGB")
|
| 166 |
+
|
| 167 |
+
elif ext == ".dcm":
|
| 168 |
+
import pydicom
|
| 169 |
+
|
| 170 |
+
dcm = pydicom.dcmread(str(image_path))
|
| 171 |
+
arr = dcm.pixel_array.astype(np.float32)
|
| 172 |
+
|
| 173 |
+
# normalize properly (medical safe scaling)
|
| 174 |
+
arr = arr - arr.min()
|
| 175 |
+
arr = arr / (arr.max() + 1e-8)
|
| 176 |
+
arr = (arr * 255).astype(np.uint8)
|
| 177 |
+
|
| 178 |
+
pil_image = Image.fromarray(arr).convert("RGB")
|
| 179 |
+
|
| 180 |
+
else:
|
| 181 |
+
pil_image = Image.open(image_path).convert("RGB")
|
| 182 |
+
|
| 183 |
+
img_w, img_h = pil_image.size
|
| 184 |
+
|
| 185 |
+
# ── load fixations and normalize pixel → [0, 1] ──
|
| 186 |
+
from pathlib import Path
|
| 187 |
+
fix_t = load_fixation_csv(
|
| 188 |
+
str(fixation_csv),
|
| 189 |
+
image_width=img_w,
|
| 190 |
+
image_height=img_h,
|
| 191 |
+
image_name = Path(image_path).stem
|
| 192 |
+
) # (N, 3) float tensor: x, y, duration
|
| 193 |
+
fix_t = fix_t.to(_device)
|
| 194 |
+
|
| 195 |
+
# ── preprocess the image for DINOv3 ──
|
| 196 |
+
tf = T.Compose([
|
| 197 |
+
T.Resize((IMG_SIZE, IMG_SIZE)),
|
| 198 |
+
T.ToTensor(),
|
| 199 |
+
T.Normalize(IMG_MEAN, IMG_STD),
|
| 200 |
+
])
|
| 201 |
+
img_t = tf(pil_image).unsqueeze(0).to(_device) # (1, 3, IMG_SIZE, IMG_SIZE)
|
| 202 |
+
|
| 203 |
+
# ── run the model ──
|
| 204 |
+
model = _get_model(preset, _device)
|
| 205 |
+
out = model(img_t, fix_t)
|
| 206 |
+
|
| 207 |
+
# ── decode outputs ──
|
| 208 |
+
soft_mask = out["preds"][0, 0].cpu().numpy() # (H, W) float in [0, 1]
|
| 209 |
+
gaze = out["gaze_heatmap"][0].cpu().numpy() # (h, w) float in [0, 1]
|
| 210 |
+
bin_mask = (soft_mask > threshold).astype(np.uint8) * 255
|
| 211 |
+
|
| 212 |
+
mask_pil = Image.fromarray(bin_mask, mode="L")
|
| 213 |
+
|
| 214 |
+
if not return_all:
|
| 215 |
+
return mask_pil
|
| 216 |
+
|
| 217 |
+
return dict(
|
| 218 |
+
mask = mask_pil,
|
| 219 |
+
gaze_overlay = overlay_heatmap(pil_image, gaze),
|
| 220 |
+
mask_overlay = overlay_mask(pil_image, (bin_mask / 255).astype(np.float32)),
|
| 221 |
+
preds = soft_mask,
|
| 222 |
+
gaze_heatmap = gaze,
|
| 223 |
+
)
|
| 224 |
+
|
| 225 |
+
|
| 226 |
+
# ── CLI ─────────────────────────────────────────────────────────────────────
|
| 227 |
+
|
| 228 |
+
def _build_parser() -> argparse.ArgumentParser:
|
| 229 |
+
ap = argparse.ArgumentParser(
|
| 230 |
+
prog="predict_single.py",
|
| 231 |
+
description="GazeRefine — zero-shot gaze-guided segmentation on a single image.",
|
| 232 |
+
formatter_class=argparse.RawDescriptionHelpFormatter,
|
| 233 |
+
epilog="""
|
| 234 |
+
Examples
|
| 235 |
+
--------
|
| 236 |
+
# colonoscopy polyp (default preset):
|
| 237 |
+
python scripts/predict_single.py \\
|
| 238 |
+
--image examples/images/kvasir_sample.jpg \\
|
| 239 |
+
--fixations examples/fixations/kvasir_sample.csv \\
|
| 240 |
+
--output output_mask.png
|
| 241 |
+
|
| 242 |
+
# prostate MRI:
|
| 243 |
+
python scripts/predict_single.py \\
|
| 244 |
+
--image examples/images/prostate_sample.png \\
|
| 245 |
+
--fixations examples/fixations/prostate_sample.csv \\
|
| 246 |
+
--output output_mask.png \\
|
| 247 |
+
--preset mri
|
| 248 |
+
|
| 249 |
+
# save overlays too:
|
| 250 |
+
python scripts/predict_single.py \\
|
| 251 |
+
--image examples/images/kvasir_sample.jpg \\
|
| 252 |
+
--fixations examples/fixations/kvasir_sample.csv \\
|
| 253 |
+
--output output_mask.png \\
|
| 254 |
+
--save_overlay
|
| 255 |
+
""",
|
| 256 |
+
)
|
| 257 |
+
ap.add_argument("--image", required=True,
|
| 258 |
+
help="Path to the input image (.jpg / .jpeg / .png).")
|
| 259 |
+
ap.add_argument("--fixations", required=True,
|
| 260 |
+
help="Path to the fixation CSV (x,y,duration — pixel coordinates).")
|
| 261 |
+
ap.add_argument("--output", required=True,
|
| 262 |
+
help="Where to save the predicted binary mask (.png).")
|
| 263 |
+
ap.add_argument("--preset", default="colonoscopy",
|
| 264 |
+
choices=list(PRESETS),
|
| 265 |
+
help="Hyperparameter preset: 'colonoscopy' (default) or 'mri'.")
|
| 266 |
+
ap.add_argument("--threshold", type=float, default=0.5,
|
| 267 |
+
help="Binarization threshold applied to the soft mask (default: 0.5).")
|
| 268 |
+
ap.add_argument("--save_overlay", action="store_true",
|
| 269 |
+
help="Also save a colour overlay PNG next to --output.")
|
| 270 |
+
ap.add_argument("--device", default=None,
|
| 271 |
+
help="'cuda' or 'cpu' — auto-detected when not given.")
|
| 272 |
+
return ap
|
| 273 |
+
|
| 274 |
+
|
| 275 |
+
def main():
|
| 276 |
+
args = _build_parser().parse_args()
|
| 277 |
+
|
| 278 |
+
print(f"[GazeRefine] image : {args.image}")
|
| 279 |
+
print(f"[GazeRefine] fixations: {args.fixations}")
|
| 280 |
+
print(f"[GazeRefine] preset : {args.preset}")
|
| 281 |
+
print(f"[GazeRefine] threshold: {args.threshold}")
|
| 282 |
+
|
| 283 |
+
result = predict(
|
| 284 |
+
image_path = args.image,
|
| 285 |
+
fixation_csv = args.fixations,
|
| 286 |
+
preset = args.preset,
|
| 287 |
+
threshold = args.threshold,
|
| 288 |
+
device = args.device,
|
| 289 |
+
return_all = args.save_overlay,
|
| 290 |
+
)
|
| 291 |
+
|
| 292 |
+
output_path = Path(args.output)
|
| 293 |
+
output_path.parent.mkdir(parents=True, exist_ok=True)
|
| 294 |
+
|
| 295 |
+
if isinstance(result, dict):
|
| 296 |
+
result["mask"].save(output_path)
|
| 297 |
+
print(f"[GazeRefine] mask saved → {output_path}")
|
| 298 |
+
if args.save_overlay:
|
| 299 |
+
overlay_path = output_path.with_stem(output_path.stem + "_overlay")
|
| 300 |
+
result["mask_overlay"].save(overlay_path)
|
| 301 |
+
gaze_path = output_path.with_stem(output_path.stem + "_gaze")
|
| 302 |
+
result["gaze_overlay"].save(gaze_path)
|
| 303 |
+
print(f"[GazeRefine] overlay → {overlay_path}")
|
| 304 |
+
print(f"[GazeRefine] gaze prior → {gaze_path}")
|
| 305 |
+
else:
|
| 306 |
+
result.save(output_path)
|
| 307 |
+
print(f"[GazeRefine] mask saved → {output_path}")
|
| 308 |
+
|
| 309 |
+
|
| 310 |
+
if __name__ == "__main__":
|
| 311 |
+
main()
|
scripts/run_eval.py
ADDED
|
@@ -0,0 +1,129 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
run_eval.py — Unified zero-shot evaluation entrypoint.
|
| 3 |
+
|
| 4 |
+
Usage
|
| 5 |
+
-----
|
| 6 |
+
python scripts/run_eval.py --config configs/kvasir.yaml
|
| 7 |
+
python scripts/run_eval.py --config configs/prostate_mri.yaml
|
| 8 |
+
|
| 9 |
+
# override anything from the YAML on the command line:
|
| 10 |
+
python scripts/run_eval.py --config configs/kvasir.yaml --root /data/Kvasir-SEG --max_iters 8
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
import argparse
|
| 14 |
+
import sys
|
| 15 |
+
from pathlib import Path
|
| 16 |
+
|
| 17 |
+
import torch
|
| 18 |
+
import yaml
|
| 19 |
+
from torch.utils.data import DataLoader
|
| 20 |
+
|
| 21 |
+
sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # repo root, for `import gazerefine`
|
| 22 |
+
|
| 23 |
+
from gazerefine import GazeRefine, compute_metrics, save_prediction
|
| 24 |
+
from gazerefine.datasets import DATASET_REGISTRY
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def load_config(path: str, overrides: argparse.Namespace) -> dict:
|
| 28 |
+
with open(path) as f:
|
| 29 |
+
cfg = yaml.safe_load(f)
|
| 30 |
+
for k, v in vars(overrides).items():
|
| 31 |
+
if k != "config" and v is not None:
|
| 32 |
+
cfg[k] = v
|
| 33 |
+
return cfg
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
@torch.no_grad()
|
| 37 |
+
def eval_epoch(model, loader, device, threshold: float, pred_dir: str, save_num: int) -> dict:
|
| 38 |
+
model.eval()
|
| 39 |
+
all_dice, all_iou = [], []
|
| 40 |
+
n = len(loader)
|
| 41 |
+
saved = 0
|
| 42 |
+
|
| 43 |
+
for i, batch in enumerate(loader):
|
| 44 |
+
img = batch["image"].to(device)
|
| 45 |
+
fix = batch["fixation"].to(device)
|
| 46 |
+
msk = batch["mask"].to(device)
|
| 47 |
+
|
| 48 |
+
out = model(img, fix)
|
| 49 |
+
preds = out["preds"]
|
| 50 |
+
m = compute_metrics(preds, msk, thr=threshold)
|
| 51 |
+
all_dice.append(m["dice_per"].cpu())
|
| 52 |
+
all_iou.append(m["iou_per"].cpu())
|
| 53 |
+
|
| 54 |
+
if saved < save_num:
|
| 55 |
+
for ridx in range(img.size(0)):
|
| 56 |
+
if saved >= save_num:
|
| 57 |
+
break
|
| 58 |
+
pred_bin = (preds[ridx] > threshold).float()
|
| 59 |
+
save_prediction(
|
| 60 |
+
img[ridx], out["gaze_heatmap"][ridx], msk[ridx], pred_bin,
|
| 61 |
+
name=f"zeroshot_sample_{batch['name'][ridx]}.png", out_dir=pred_dir,
|
| 62 |
+
)
|
| 63 |
+
saved += 1
|
| 64 |
+
|
| 65 |
+
if (i + 1) % max(1, n // 5) == 0:
|
| 66 |
+
print(f" [{i + 1}/{n}] dice={m['dice']:.4f} iou={m['iou']:.4f}")
|
| 67 |
+
|
| 68 |
+
all_dice = torch.cat(all_dice)
|
| 69 |
+
all_iou = torch.cat(all_iou)
|
| 70 |
+
return {
|
| 71 |
+
"dice": all_dice.mean().item(), "dice_std": all_dice.std().item(),
|
| 72 |
+
"iou": all_iou.mean().item(), "iou_std": all_iou.std().item(),
|
| 73 |
+
}
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def main():
|
| 77 |
+
ap = argparse.ArgumentParser("GazeRefine — zero-shot gaze-guided segmentation")
|
| 78 |
+
ap.add_argument("--config", required=True, help="path to a YAML config, e.g. configs/kvasir.yaml")
|
| 79 |
+
# optional CLI overrides for anything in the YAML
|
| 80 |
+
ap.add_argument("--root", type=str, default=None)
|
| 81 |
+
ap.add_argument("--fixation_csv", type=str, default=None)
|
| 82 |
+
ap.add_argument("--pred_dir", type=str, default=None)
|
| 83 |
+
ap.add_argument("--dino_name", type=str, default=None)
|
| 84 |
+
ap.add_argument("--batch_size", type=int, default=None)
|
| 85 |
+
ap.add_argument("--num_workers", type=int, default=None)
|
| 86 |
+
ap.add_argument("--sigma", type=float, default=None)
|
| 87 |
+
ap.add_argument("--threshold", type=float, default=None)
|
| 88 |
+
ap.add_argument("--extract_mode", type=str, default=None, choices=["all", "last"])
|
| 89 |
+
ap.add_argument("--contrast_method", type=str, default=None, choices=["difference", "softmax", "original"])
|
| 90 |
+
ap.add_argument("--temperature", type=float, default=None)
|
| 91 |
+
ap.add_argument("--max_iters", type=int, default=None)
|
| 92 |
+
ap.add_argument("--gaze_anchor_weight", type=float, default=None)
|
| 93 |
+
ap.add_argument("--knn_k", type=int, default=None)
|
| 94 |
+
ap.add_argument("--save_num", type=int, default=None)
|
| 95 |
+
args = ap.parse_args()
|
| 96 |
+
|
| 97 |
+
cfg = load_config(args.config, args)
|
| 98 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 99 |
+
print(f"[GazeRefine] device={device}, dataset={cfg['dataset']}")
|
| 100 |
+
|
| 101 |
+
DatasetCls = DATASET_REGISTRY[cfg["dataset"]]
|
| 102 |
+
dataset = DatasetCls(cfg["root"], cfg["fixation_csv"])
|
| 103 |
+
loader = DataLoader(dataset, batch_size=cfg["batch_size"], num_workers=cfg["num_workers"], shuffle=False, pin_memory=True)
|
| 104 |
+
print(f"[GazeRefine] total samples to evaluate: {len(dataset)}")
|
| 105 |
+
|
| 106 |
+
model = GazeRefine(
|
| 107 |
+
dino_name=cfg["dino_name"],
|
| 108 |
+
sigma=cfg["sigma"],
|
| 109 |
+
extract_mode=cfg["extract_mode"],
|
| 110 |
+
contrast_method=cfg["contrast_method"],
|
| 111 |
+
temperature=cfg["temperature"],
|
| 112 |
+
max_iters=cfg["max_iters"],
|
| 113 |
+
knn_refine=cfg["knn_refine"],
|
| 114 |
+
knn_k=cfg["knn_k"],
|
| 115 |
+
knn_temp=cfg["knn_temp"],
|
| 116 |
+
gaze_anchor_weight=cfg["gaze_anchor_weight"],
|
| 117 |
+
).to(device)
|
| 118 |
+
|
| 119 |
+
print(f"[GazeRefine] running zero-shot evaluation on {cfg['dataset']}...")
|
| 120 |
+
metrics = eval_epoch(model, loader, device, threshold=cfg["threshold"], pred_dir=cfg["pred_dir"], save_num=cfg["save_num"])
|
| 121 |
+
|
| 122 |
+
print(f"\n{'═' * 60}")
|
| 123 |
+
print(f"[Results] Dice: {metrics['dice']:.4f} ± {metrics['dice_std']:.4f}")
|
| 124 |
+
print(f"[Results] IoU: {metrics['iou']:.4f} ± {metrics['iou_std']:.4f}")
|
| 125 |
+
print(f"{'═' * 60}\n")
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
if __name__ == "__main__":
|
| 129 |
+
main()
|