Upload 41 files

README.md

@@ -1,5 +1,5 @@
 ---
-title: GazeRefine
+title: GazeCorrect
 emoji: 👁️
 colorFrom: blue
 colorTo: red
@@ -8,31 +8,22 @@ sdk_version: 4.44.1
 app_file: app.py
 pinned: false
 license: mit
-short_description: Zero-shot, training-free gaze-guided medical segmentation
+short_description: Gaze-guided Gaussian-noise image regeneration
 ---
 
-# GazeRefine — Expert Gaze as a Test-Time Prompt
+# GazeCorrect — Gaze-Guided Regeneration
 
-Interactive demo for **GazeRefine**, a training-free, zero-shot framework
-that turns expert eye-gaze into an inference-time prompt for medical image
-segmentation. Frozen DINOv3 patch features + gaze-anchored
-foreground/background prototypes + recurrent contrastive cleaning + kNN
-affinity propagation — no masks, no clicks-as-boxes, no fine-tuning, no
-adapters, no prompt encoder.
+Interactive demo that converts clinician gaze into a duration-weighted Gaussian
+attention map. It adds Gaussian noise in the attended region, then uses a
+text description and RoentGen-v2 image-to-image generation to regenerate that
+region. No DINOv3 or segmentation pipeline is used.
 
 ## How to use
-1. Upload a colonoscopy or grayscale-MRI-style image.
-2. Click on the image 1–5 times where a clinician's gaze would land on the
-   structure of interest (a polyp, the prostate, ...). Each click adds a
-   numbered fixation marker; the slider controls that fixation's relative
-   duration/weight before your next click.
-3. Pick a hyperparameter preset (tuned per-modality, see the paper).
-4. Press **Run GazeRefine** to get the gaze-prior overlay and the predicted
-   segmentation mask.
-5. Optionally open **Correct the gaze-attended region**. The predicted mask is
-   feathered, Gaussian noise is added only inside that region, and the noised
-   image is sent to RoentGen-v2 image-to-image inference. The generated result
-   is composited back only inside the mask, retaining the original background.
+1. Upload a chest X-ray.
+2. Click gaze points or upload a CSV containing `x,y,duration`.
+3. Enter a disease or radiology description.
+4. Press **Generate** to receive the gaze attention map, attention-weighted
+   Gaussian-noise image, and corrected regenerated image.
 
 ## Synthetic chest X-ray correction
 
@@ -63,11 +54,11 @@ and correction callbacks with `@spaces.GPU`. Keep `app.py` unchanged at the
 top level; removing those decorators causes the ZeroGPU startup error
 “No @spaces.GPU function detected”.
 
-## Notes
-- Inference uses a frozen `vit_large_patch16_dinov3.lvd1689m` backbone from
-  `timm`. First run will download the checkpoint.
-- CPU inference works but is slow; a GPU Space is recommended for a smooth
-  demo.
+## Notes
+- RoentGen-v2 is trained for chest X-rays; do not use this workflow for other
+  medical modalities.
+- The first generation downloads the model. A GPU or ZeroGPU Space is
+  recommended.
 - This Space is for research/demonstration only — it is **not** a clinical
   diagnostic tool.
 

app.py

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

requirements.txt

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