File size: 28,247 Bytes
4d03be3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
938fad3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4d03be3
 
938fad3
4d03be3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
938fad3
4d03be3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3f86f08
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4d03be3
 
 
 
 
 
 
3f86f08
4d03be3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3f86f08
4d03be3
 
 
 
 
 
 
3f86f08
4d03be3
 
3f86f08
4d03be3
 
 
 
 
3f86f08
 
4d03be3
 
 
 
 
 
 
 
 
 
 
3f86f08
4d03be3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3f86f08
4d03be3
 
3f86f08
4d03be3
 
3f86f08
4d03be3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3f86f08
4d03be3
 
 
 
 
 
 
 
 
 
3f86f08
 
 
 
 
 
 
 
 
4d03be3
3f86f08
4d03be3
 
3f86f08
4d03be3
 
 
 
 
 
 
 
 
 
 
 
faca6ce
4d03be3
 
faca6ce
4d03be3
 
 
 
 
 
faca6ce
4d03be3
 
 
faca6ce
4d03be3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
faca6ce
4d03be3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
faca6ce
4d03be3
 
 
 
3f86f08
 
4d03be3
 
3f86f08
 
4d03be3
 
 
 
 
 
faca6ce
4d03be3
 
 
 
 
 
 
 
ecfd712
 
4d03be3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ecfd712
4d03be3
3f86f08
 
 
4d03be3
3f86f08
 
 
d14c9eb
ecfd712
4d03be3
ecfd712
 
faca6ce
 
 
 
 
 
 
 
 
 
ecfd712
 
 
 
 
 
 
 
9319bc3
 
 
 
 
faca6ce
4d03be3
 
 
 
 
3f86f08
4d03be3
 
25e8ff2
2a095ab
4d03be3
 
 
 
 
 
 
3f86f08
4d03be3
 
 
 
3f86f08
4d03be3
3f86f08
4d03be3
 
 
3f86f08
4d03be3
 
 
 
 
 
 
 
 
 
3f86f08
4d03be3
 
 
 
 
faca6ce
4d03be3
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
"""
Gradio demo for GazeAlign β€” gaze-supervised medical image classification.

Workflow
--------
1. Upload an image (JPG / PNG / BMP / TIFF / WEBP / DICOM).
2. Provide a radiologist-style scanpath in **either** of two ways:
     β€’ click on the image to drop fixation points, or
     β€’ upload a fixation table (.csv / .xlsx / .xls) and map its columns.
3. Run the model to get the predicted class (+ per-class probabilities)
   and the learned gaze-conditioned attention mask.

Run locally with:  python app.py
Deployed as a HuggingFace Space, this file is the entry point.
"""

from __future__ import annotations

import sys
import types
import os
from pathlib import Path

# ── 1. audioop shim (Python 3.13 removed audioop; some deps import 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 (old/new signature) ──
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]

# ── 3. huggingface_hub HfFolder shim β€” MUST run *before* `import gradio` ──────
# Newer huggingface_hub versions removed `HfFolder`, but `gradio.oauth` does
# `from huggingface_hub import HfFolder, whoami` at import time, so importing
# gradio blows up unless we put a compatible `HfFolder` back first.
import huggingface_hub as _hfh

if not hasattr(_hfh, "HfFolder"):
    class _FakeHfFolder:
        @staticmethod
        def get_token():
            try:
                from huggingface_hub import get_token as _gt

                return _gt()
            except Exception:
                return None

        @staticmethod
        def save_token(token):
            return None

    _hfh.HfFolder = _FakeHfFolder  # type: ignore[attr-defined]
    sys.modules["huggingface_hub"].HfFolder = _FakeHfFolder  # type: ignore[assignment]

import gradio as gr

# ── 4. gradio_client schema shim (guards against bad additionalProperties) ───
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

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

# ── 5. Path setup β€” make the repo root importable ────────────────────────────
_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)

from GazeAlign import get_device, get_scanpath  # noqa: E402
from GazeAlign.visualize import heatmap_to_image, make_overlay, patch_to_image  # noqa: E402
from scripts.predict_single import GazeAlignPredictor  # noqa: E402

# ─────────────────────────────────────────────────────────────────────────────
# Config
# ─────────────────────────────────────────────────────────────────────────────

PRESETS_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "configs", "presets.yaml")
# Friendly label β†’ preset key in configs/presets.yaml. Add rows here as you
# train GazeAlign on new modalities.
PRESETS = {
    "Chest X-ray β€” CHF / Normal / Pneumonia": "cxr",
}
POINT_COLORS = ["#ff3b30", "#ff9500", "#ffcc00", "#34c759", "#5ac8fa", "#007aff", "#af52de"]
_NO_COL = "β€” none β€”"

_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"

_DEVICE = get_device()
_PREDICTORS: dict[str, GazeAlignPredictor] = {}


def get_predictor(preset_key: str) -> GazeAlignPredictor:
    """Lazily build & cache one predictor per preset."""
    if preset_key not in _PREDICTORS:
        _PREDICTORS[preset_key] = GazeAlignPredictor.from_preset(
            preset_key, presets_path=PRESETS_PATH, device=str(_DEVICE)
        )
    return _PREDICTORS[preset_key]


# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────


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)
    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 + connecting saccades on a copy of `image`.

    `points`: list of (x_px, y_px, weight) in original-image pixel coords,
    `weight` in [0, 1] (relative dwell / recency, controls circle size).
    """
    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)
    prev = None
    for (x_px, y_px, _weight) in points:
        color = POINT_COLORS[0]
        if prev is not None:
            draw.line([prev, (x_px, y_px)], fill=color, width=2)
        prev = (x_px, y_px)
    for i, (x_px, y_px, weight) in enumerate(points):
        color = POINT_COLORS[i % len(POINT_COLORS)]
        rad = r * (0.6 + 0.8 * float(weight))
        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 β€” 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] (with rounding slack) are treated as
    normalised; otherwise they are assumed to be raw pixels and 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


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 _status(points: list) -> str:
    """Small feedback line so it's obvious when fixations register."""
    n = len(points) if points else 0
    if not n:
        return "_No fixations yet β€” click the image, or load a fixation file below._"
    return f"**{n}** fixation(s) placed."


def gaze_duration_heatmap(points, height, width, img_w, img_h, sigma=None):
    """Gaussian-splat heatmap of the fixations, each blob weighted by its
    dwell / duration (the 3rd component of each point), rendered at
    (height, width). This is the *observed* gaze heatmap β€” not the model's
    learned attention.

    points: list of (x_px, y_px, weight) in original-image pixels.
    """
    hm = np.zeros((height, width), dtype=np.float32)
    if not points:
        return hm
    if sigma is None:
        sigma = max(height, width) / 22.0
    sx, sy = width / max(img_w, 1), height / max(img_h, 1)
    rad = max(int(sigma * 3), 1)
    for x_px, y_px, wgt in points:
        cx, cy = int(round(x_px * sx)), int(round(y_px * sy))
        if not (0 <= cx < width and 0 <= cy < height):
            continue
        x0, x1 = max(cx - rad, 0), min(cx + rad + 1, width)
        y0, y1 = max(cy - rad, 0), min(cy + rad + 1, height)
        xv, yv = np.meshgrid(np.arange(x0, x1), np.arange(y0, y1))
        g = np.exp(-((xv - cx) ** 2 + (yv - cy) ** 2) / (2.0 * sigma ** 2))
        # +0.15 floor so short-dwell fixations still register a little.
        hm[y0:y1, x0:x1] += g.astype(np.float32) * (0.15 + float(wgt))
    if hm.max() > 0:
        hm /= hm.max()
    return hm


# ─────────────────────────────────────────────────────────────────────────────
# Event handlers β€” image
# ─────────────────────────────────────────────────────────────────────────────


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())
    if file_obj is None:
        return _no_change

    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:  # noqa: BLE001
        gr.Warning(f"Could not load file: {e}")
        return _no_change

    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
        gr.update(visible=True),                                    # delete_btn β†’ show
        _status([]),                                                # fix_status
    )


def on_select(orig_image: Image.Image, points: list, weight: 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(), _status(points)
    x_px, y_px = float(evt.index[0]), float(evt.index[1])
    new_points = points + [(x_px, y_px, float(weight))]
    return new_points, draw_points(orig_image, new_points), _status(new_points)


def on_clear(orig_image):
    """Remove all fixations but keep the current image."""
    if orig_image is None:
        return [], gr.update(), _status([])
    return [], gr.update(value=orig_image), _status([])


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
        _status([]),                             # fix_status
    )


# ─────────────────────────────────────────────────────────────────────────────
# Event handlers β€” fixation file
# ─────────────────────────────────────────────────────────────────────────────


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:  # noqa: BLE001
        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]

    def _guess(*keywords, fallback=None):
        # Priority-ordered: try each keyword across ALL columns before moving
        # to the next, so e.g. "dicom" wins over a stray "id" in "SESSION_ID".
        for k in keywords:
            for c in cols:
                if k in c.lower():
                    return c
        return fallback if fallback is not None else cols[0]

    guess_id = _guess("dicom", "image", "id", "name", "file", fallback=cols[0])
    guess_x = _guess("x_original", "x_orig", "fix_x", "pos_x", "gaze_x", "x_pixel", 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("y_original", "y_orig", "fix_y", "pos_y", "gaze_y", "y_pixel", 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],
        )

    time_choices = [_NO_COL] + cols
    guess_time = _guess("time", "secs", "duration", "dur", "timestamp", fallback=_NO_COL)

    return (
        df.to_json(),                                       # fixfile_df_state
        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=time_choices, value=guess_time),  # time_col_dd
        gr.update(visible=True),                            # apply_fix_btn β†’ show
    )


def on_apply_fixfile(fixfile_json, id_col, x_col, y_col, time_col, orig_image, image_name):
    """Match rows to the loaded image (by filename) and load them as
    fixation points, replacing whatever points are currently set.

    If rows can't be matched by filename but the file holds a single image's
    worth of fixations, all rows are used (handy for single-image CSVs whose
    ID column doesn't match the uploaded filename)."""
    if orig_image is None:
        gr.Warning("Load an image first, then apply the fixation file.")
        return gr.update(), gr.update(), gr.update()
    if not fixfile_json:
        gr.Warning("Upload a fixation file first.")
        return gr.update(), gr.update(), gr.update()
    if not x_col or not y_col:
        gr.Warning("Pick the X and Y columns first.")
        return gr.update(), gr.update(), gr.update()

    df = pd.read_json(fixfile_json)

    sub = df
    if id_col and image_name:
        mask = df[id_col].astype(str) == image_name
        if not mask.any():
            stem = Path(image_name).stem
            mask = df[id_col].astype(str).apply(lambda v: Path(str(v)).stem) == stem
        if mask.any():
            sub = df[mask]
        elif df[id_col].nunique() > 1:
            gr.Warning(
                f"No rows match the loaded image ('{image_name}') and the file "
                f"has several ids β€” using ALL rows. Check the ID column."
            )

    if sub.empty:
        gr.Warning("No usable fixation rows found.")
        return gr.update(), gr.update(), gr.update()

    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 time_col and time_col != _NO_COL and time_col in sub.columns:
        t_raw = sub[time_col].astype(float).to_numpy()
        order = np.argsort(t_raw)  # chronological order
        x_px, y_px, t_raw = x_px[order], y_px[order], t_raw[order]
        # Per-fixation dwell = gap to the next fixation (last one gets the
        # median gap); normalised to [0,1] so it weights the duration heatmap.
        if len(t_raw) > 1:
            dwell = np.diff(t_raw, append=t_raw[-1] + np.median(np.diff(t_raw)))
            dwell = np.clip(dwell, 0, None)
            dmax = float(dwell.max())
            weight = dwell / dmax if dmax > 0 else np.ones_like(dwell)
        else:
            weight = np.ones(1)
    else:
        weight = np.ones(len(sub))

    new_points = [(float(xp), float(yp), float(wt)) for xp, yp, wt in zip(x_px, y_px, weight)]
    return new_points, draw_points(orig_image, new_points), _status(new_points)


# ─────────────────────────────────────────────────────────────────────────────
# Inference
# ─────────────────────────────────────────────────────────────────────────────


def run(orig_image: Image.Image, points: list, preset_name: str):
    import traceback

    if orig_image is None:
        gr.Warning("Upload an image first.")
        return None, "", None
    if not points or len(points) < 2:
        gr.Warning("Provide at least 2 fixations (click the image or load a fixation file).")
        return None, "", None

    preset_key = PRESETS[preset_name]
    try:
        predictor = get_predictor(preset_key)
    except FileNotFoundError as e:
        gr.Warning(str(e))
        return None, f"**Checkpoint not found** for preset `{preset_key}`.", None
    except Exception as e:  # noqa: BLE001
        traceback.print_exc()
        gr.Warning(f"Could not load model: {e}")
        return None, "", None

    w, h = orig_image.size
    # Build a MIMIC-style scanpath dataframe. The 3rd component (weight)
    # drives a monotonically increasing time axis for the scanpath encoder.
    weights = np.asarray([p[2] for p in points], dtype=float)
    times = np.cumsum(np.clip(weights, 1e-3, None))
    df = pd.DataFrame(
        {
            "DICOM_ID": ["webdemo"] * len(points),
            "X_ORIGINAL": [p[0] for p in points],
            "Y_ORIGINAL": [p[1] for p in points],
            "Time (in secs)": times,
        }
    )

    scanpath = get_scanpath(df, "webdemo", img_height=h, img_width=w)
    if scanpath is None or scanpath.numel() == 0:
        gr.Warning("Could not build a scanpath from the fixations.")
        return None, "", None
    scanpath = scanpath[:200].to(predictor.device)

    img_tensor = predictor.transform(np.array(orig_image)).unsqueeze(0).to(predictor.device)
    try:
        with torch.no_grad():
            _, patch_tokens, _ = predictor.image_encoder(img_tensor)
            _, sp_emb, _ = predictor.scanpath_encoder([scanpath])
            patch_mask = torch.sigmoid(predictor.mask_generator(sp_emb))  # [1, g, g]

            B, N, D = patch_tokens.shape
            feat_attended = (patch_tokens * patch_mask.view(B, N, 1)).mean(dim=1)
            logits = predictor.classifier(feat_attended)
            probs = torch.softmax(logits, dim=1)[0].cpu().numpy()
    except Exception as e:  # noqa: BLE001
        traceback.print_exc()
        gr.Warning(f"Prediction failed: {e}")
        return None, "", None

    class_probs = {c: float(p) for c, p in zip(predictor.classes, probs)}
    predicted_class = max(class_probs, key=class_probs.get)

    # Output visual: the *observed* gaze-fixation heatmap, each fixation
    # weighted by its dwell/duration β€” overlaid on the image and shown raw.
    img_size = predictor.img_size
    display_img = np.array(orig_image.resize((img_size, img_size)))
    gaze_hm = gaze_duration_heatmap(points, img_size, img_size, w, h)
    overlay = make_overlay(display_img, gaze_hm)

    prob_lines = "\n".join(
        f"- **{c}**: {p:.3f}" for c, p in sorted(class_probs.items(), key=lambda kv: -kv[1])
    )
    summary = f"### Predicted: **{predicted_class}**\n\n{prob_lines}"

    return class_probs, summary, overlay


# ─────────────────────────────────────────────────────────────────────────────
# UI
# ─────────────────────────────────────────────────────────────────────────────

_CSS = """
#run-btn {font-weight: 600;}
#upload-zone {min-height: 240px;}
#upload-zone .center {min-height: 220px;}
.footer-note {opacity: 0.7; font-size: 0.85rem;}
"""

with gr.Blocks(title="GazeAlign", css=_CSS) as demo:
    gr.Markdown(
        """
        # πŸ‘οΈ GazeAlign β€” Gaze-Supervised Medical Image Classification

        **1.** Upload an image Β· **2.** Add fixations by *clicking* the image
        **or** *uploading a fixation table (.csv / .xlsx)* Β· **3.** Run the model.
        """
    )

    orig_image_state = gr.State(None)
    points_state = gr.State([])
    image_name_state = gr.State("")
    fixfile_df_state = gr.State(None)

    with gr.Row():
        # ── Left: image + fixations ──────────────────────────────────────────
        with gr.Column(scale=1):
            upload_zone = gr.File(
                label=_UPLOAD_LABEL,
                file_types=[".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp", ".dcm"],
                type="filepath",
                elem_id="upload-zone",
            )
            # interactive=False β†’ a pure display surface that reports click
            # coordinates via .select (an interactive Image opens an editor
            # instead and never fires reliable pixel coords).
            image_panel = gr.Image(
                label=_FIXATION_LABEL, type="pil", interactive=False,
                visible=False, height=440,
            )
            delete_btn = gr.Button("πŸ—‘  Delete image / load another", visible=False)
            fix_status = gr.Markdown("")

            # Load-from-file menu β€” collapsed; click the header to reveal.
            with gr.Accordion("πŸ“„  Load fixations from file", open=False):
                fixfile = gr.File(
                    label=_FIXFILE_LABEL, file_types=[".csv", ".xlsx", ".xls"], type="filepath"
                )
                with gr.Row(visible=False) as mapping_row:
                    id_col_dd = gr.Dropdown(label="ID column", choices=[])
                    x_col_dd = gr.Dropdown(label="X column", choices=[])
                    y_col_dd = gr.Dropdown(label="Y column", choices=[])
                    time_col_dd = gr.Dropdown(label="Time column (optional)", choices=[])
                apply_fix_btn = gr.Button("Apply fixation file", visible=False)

            # Duration weight + clear on one row (screenshot layout).
            with gr.Row():
                weight_slider = gr.Slider(
                    0.0, 1.0, value=1.0, step=0.05,
                    label="Fixation duration weight", scale=3,
                )
                clear_btn = gr.Button("βœ• Clear fixations", scale=2)

            preset_dd = gr.Dropdown(
                choices=list(PRESETS.keys()),
                value=list(PRESETS.keys())[0],
                label="Model / modality preset",
            )
            run_btn = gr.Button("Run GazeAlign", variant="primary", elem_id="run-btn")

        # ── Right: results ───────────────────────────────────────────────────
        with gr.Column(scale=1):
            label_output = gr.Label(label="Predicted class (probabilities)", num_top_classes=5)
            summary_output = gr.Markdown()
            overlay_output = gr.Image(label="Gaze-fixation heatmap (dwell-weighted) β€” overlay")

    gr.Markdown(
        "<div class='footer-note'>See the "
        "<a href='https://github.com/anonymous-IA/GazeAlign'>GitHub repo</a> "
        "for training and evaluation code.</div>"
    )

    # ── wiring ──
    upload_zone.upload(
        on_file_upload,
        [upload_zone],
        [orig_image_state, points_state, image_name_state, upload_zone, image_panel, delete_btn, fix_status],
    )
    image_panel.select(
        on_select,
        [orig_image_state, points_state, weight_slider],
        [points_state, image_panel, fix_status],
    )
    clear_btn.click(on_clear, [orig_image_state], [points_state, image_panel, fix_status])
    delete_btn.click(
        on_delete,
        None,
        [orig_image_state, points_state, image_name_state, upload_zone, image_panel, delete_btn, fix_status],
    )

    fixfile.upload(
        on_fixfile_upload,
        [fixfile],
        [fixfile_df_state, mapping_row, id_col_dd, x_col_dd, y_col_dd, time_col_dd, apply_fix_btn],
    )
    apply_fix_btn.click(
        on_apply_fixfile,
        [fixfile_df_state, id_col_dd, x_col_dd, y_col_dd, time_col_dd, orig_image_state, image_name_state],
        [points_state, image_panel, fix_status],
    )

    run_btn.click(
        run,
        [orig_image_state, points_state, preset_dd],
        [label_output, summary_output, overlay_output],
    )


if __name__ == "__main__":
    demo.launch()