# SPDX-FileCopyrightText: Copyright © 2026 Idiap Research Institute <contact@idiap.ch>
#
# SPDX-License-Identifier: GPL-3.0-or-later
"""Sequence (multi-frame) data loading for the AngioReport dataset.
This module mirrors :py:mod:`mednet.data.classify.angioreport` but loads
multiple frames per exam so temporal models can process longitudinal signal.
"""
from __future__ import annotations
import json
import os
import pathlib
import typing
from importlib.resources import abc as resources_abc
import PIL.Image
import torch
import torch.nn.functional
from torchvision import tv_tensors
from torchvision.transforms.v2.functional import to_dtype, to_image
from ...utils.rc import load_rc
from ..classify.angioreport import (
CONFIGURATION_KEY_DATADIR,
DATABASE_SLUG,
binarize_findings,
)
from ..datamodule import CachingDataModule
from ..typing import DatabaseSplit, Sample
from ..typing import RawDataLoader as BaseDataLoader
from .hyperftype_split import (
FrameStamp,
elapsed_seconds_by_file,
frame_names,
parse_frames_field,
)
BOUND_EARLY_MID = float(os.environ.get("MEDNET_TEMPORAL_BOUND_EARLY_MID", "103.0"))
BOUND_MID_LATE = float(os.environ.get("MEDNET_TEMPORAL_BOUND_MID_LATE", "518.0"))
# phase_id value for an intentional padding timestep (``one_per_phase`` missing phase).
PHASE_ID_PADDING = -2
NUM_PHASES = 3
def _phase_id_from_seconds(seconds: float) -> int:
if seconds < BOUND_EARLY_MID:
return 0
if seconds < BOUND_MID_LATE:
return 1
return 2
def _uniform_pick_indices(indices: list[int], k: int) -> list[int]:
if k <= 0 or not indices:
return []
if len(indices) <= k:
return list(indices)
if k == 1:
return [indices[-1]]
return [indices[round(i * (len(indices) - 1) / (k - 1))] for i in range(k)]
def _pick_one_index_per_phase_slot(
phase_to_indices: dict[int, list[int]],
) -> list[int | None]:
"""One frame index per phase slot (early, mid, late); ``None`` if phase is empty.
Picks the first early frame, the centre mid frame, and the first late frame.
Parameters
----------
phase_to_indices
Mapping from phase id (0/1/2) to sorted frame indices within that phase.
Returns
-------
list[int | None]
Three-element list; ``None`` for phases with no available frames.
"""
slots: list[int | None] = []
for phase in (0, 1, 2):
indices = phase_to_indices.get(phase, [])
if not indices:
slots.append(None)
continue
if phase == 0:
slots.append(indices[0])
elif phase == 1:
slots.append(indices[len(indices) // 2])
else:
slots.append(indices[0])
return slots
[docs]
class GroupedJSONSplit(DatabaseSplit):
"""Read a frame-level split JSON and expose one sample per exam.
Supports two split formats:
1) Legacy frame-level entries:
``["Train/Train/<exam>/<frame>.jpg", label]``
grouped into one sample per exam.
2) Temporal exam-level entries:
``{"exam": "Train/Train/<exam>", "label": y, "frames": [...]}``
where ``frames`` is either legacy filenames or objects with
``file`` + ``elapsed_seconds``.
Parameters
----------
path
Path to the split JSON (local file or importlib traversable resource).
"""
def __init__(self, path: pathlib.Path | resources_abc.Traversable) -> None:
with path.open() as f:
raw = json.load(f)
self._splits: dict[
str, list[tuple[str, typing.Any, list[FrameStamp] | None]]
] = {}
for split_name, entries in raw.items():
# New temporal format: one explicit exam row with frames.
if entries and isinstance(entries[0], dict):
rows: list[tuple[str, typing.Any, list[FrameStamp] | None]] = []
for row in entries:
exam_folder = str(row["exam"])
label = row["label"]
frames_raw = row.get("frames")
frames = (
parse_frames_field(frames_raw)
if frames_raw is not None
else None
)
rows.append((exam_folder, label, frames))
self._splits[split_name] = rows
continue
# Legacy format: frame-level entries, collapse to one exam sample.
grouped: dict[str, typing.Any] = {}
for frame_path, label in entries:
exam_folder = str(pathlib.Path(frame_path).parent)
grouped[exam_folder] = label
self._splits[split_name] = [
(exam_folder, label, None) for exam_folder, label in grouped.items()
]
def __getitem__(
self, key: str
) -> list[tuple[str, typing.Any, list[FrameStamp] | None]]:
return self._splits[key]
def __iter__(self):
return iter(self._splits)
def __len__(self) -> int:
return len(self._splits)
[docs]
class SequenceRawDataLoader(BaseDataLoader):
"""A specialized raw-data-loader for temporal AngioReport sequences.
Per-frame timestamps (and hence phases) come from the split JSON exam rows
(see ``mednet.data.temporal.hyperftype_split``).
Parameters
----------
problem_type
Problem type for target formatting.
modality
Modality to crop from side-by-side frames.
max_frames
Optional cap on number of frames per sequence.
max_frames_span
How to pick frames when the sequence is longer than ``max_frames``:
``"uniform"`` keeps an even temporal coverage, ``"last"`` keeps the most
recent frames, ``"phase_stratified"`` balances early/mid/late bins, and
``"one_per_phase"`` keeps at most one frame per phase (see
``_select_one_per_phase_slots``).
"""
datadir: pathlib.Path
@staticmethod
def _frame_sort_key(path: pathlib.Path) -> tuple[int, int | str]:
stem = path.stem
if stem.isdigit():
return (0, int(stem))
return (1, stem)
def __init__(
self,
problem_type: typing.Literal["binary", "multiclass", "multilabel"],
modality: typing.Literal["FA", "ICGA"] = "FA",
max_frames: int | None = 32,
max_frames_span: typing.Literal[
"uniform", "last", "phase_stratified", "one_per_phase"
] = "uniform",
) -> None:
self.datadir = pathlib.Path(
load_rc().get(
CONFIGURATION_KEY_DATADIR,
os.path.realpath(os.curdir),
)
)
self.problem_type = problem_type
self.modality = modality
self.max_frames = max_frames
self.max_frames_span = max_frames_span
def _load_one_frame(self, path: pathlib.Path) -> tv_tensors.Image:
image = PIL.Image.open(path).convert("L")
width, height = image.size
box = (
(0, 0, width - width / 2, height - 50)
if self.modality == "FA"
else (width - width / 2, 0, width, height - 50)
)
image = image.crop(box)
image = to_dtype(to_image(image), torch.float32, scale=True)
return tv_tensors.Image(image)
def _borrow_index_for_phase(
self,
phase: int,
phase_to_indices: dict[int, list[int]],
all_indices: list[int],
) -> int:
if phase == 0: # early: prefer first mid, then first late
if phase_to_indices[1]:
return phase_to_indices[1][0]
if phase_to_indices[2]:
return phase_to_indices[2][0]
elif phase == 1: # mid: prefer last early, then first late
if phase_to_indices[0]:
return phase_to_indices[0][-1]
if phase_to_indices[2]:
return phase_to_indices[2][0]
else: # late: prefer last mid, then last early
if phase_to_indices[1]:
return phase_to_indices[1][-1]
if phase_to_indices[0]:
return phase_to_indices[0][-1]
return all_indices[-1]
def _select_one_per_phase_slots(
self,
frame_paths: list[pathlib.Path],
frame_stamps: list[FrameStamp] | None,
) -> list[tuple[pathlib.Path | None, int]]:
"""Return exactly three (path, phase_id) slots; missing phases use ``None``.
Parameters
----------
frame_paths
Sorted list of available frame paths for the exam.
frame_stamps
Optional per-frame timestamp metadata from the split JSON.
Returns
-------
list[tuple[pathlib.Path | None, int]]
Three ``(path, phase_id)`` pairs for early, mid, and late phases.
"""
phase_to_indices: dict[int, list[int]] = {0: [], 1: [], 2: []}
for i, path in enumerate(frame_paths):
phase = self._phase_for_frame(path.name, frame_stamps)
if phase in (0, 1, 2):
phase_to_indices[phase].append(i)
index_slots = _pick_one_index_per_phase_slot(phase_to_indices)
if any(idx is not None for idx in index_slots):
return [
(
frame_paths[idx] if idx is not None else None,
phase if idx is not None else PHASE_ID_PADDING,
)
for idx, phase in zip(index_slots, (0, 1, 2), strict=True)
]
# No stamped phases: single last frame in the late slot.
return [
(None, PHASE_ID_PADDING),
(None, PHASE_ID_PADDING),
(frame_paths[-1], 2),
]
def _select_frames(
self,
frame_paths: list[pathlib.Path],
frame_stamps: list[FrameStamp] | None,
) -> list[pathlib.Path]:
if self.max_frames is None or len(frame_paths) <= self.max_frames:
return frame_paths
if self.max_frames_span == "last":
return frame_paths[-self.max_frames :]
if self.max_frames_span == "phase_stratified":
all_indices = list(range(len(frame_paths)))
phase_to_indices: dict[int, list[int]] = {0: [], 1: [], 2: []}
for i, path in enumerate(frame_paths):
phase = self._phase_for_frame(path.name, frame_stamps)
if phase in (0, 1, 2):
phase_to_indices[phase].append(i)
# 12 -> 4/4/4 by default, with remainder assigned early->mid->late.
base = self.max_frames // 3
rem = self.max_frames % 3
quotas = {0: base, 1: base, 2: base}
for p in (0, 1, 2):
if rem <= 0:
break
quotas[p] += 1
rem -= 1
selected: list[int] = []
for phase in (0, 1, 2):
q = quotas[phase]
phase_indices = phase_to_indices[phase]
picked = _uniform_pick_indices(phase_indices, q)
if len(picked) < q:
borrow = self._borrow_index_for_phase(
phase=phase,
phase_to_indices=phase_to_indices,
all_indices=all_indices,
)
picked += [borrow] * (q - len(picked))
selected.extend(picked)
# Guardrail: if something went wrong, keep fixed-length output.
if len(selected) < self.max_frames:
selected += [selected[-1] if selected else all_indices[-1]] * (
self.max_frames - len(selected)
)
elif len(selected) > self.max_frames:
selected = selected[: self.max_frames]
# Keep temporal order for GRU.
selected = sorted(selected)
return [frame_paths[i] for i in selected]
if self.max_frames == 1:
return [frame_paths[-1]]
indices = _uniform_pick_indices(list(range(len(frame_paths))), self.max_frames)
return [frame_paths[i] for i in indices]
def _pad_to_max_hw(self, frames: list[torch.Tensor]) -> torch.Tensor:
max_h = max(frame.shape[-2] for frame in frames)
max_w = max(frame.shape[-1] for frame in frames)
padded = []
for frame in frames:
dh = max_h - frame.shape[-2]
dw = max_w - frame.shape[-1]
top = dh // 2
bottom = dh - top
left = dw // 2
right = dw - left
padded.append(
torch.nn.functional.pad(
frame,
(left, right, top, bottom),
mode="constant",
value=0.0,
)
)
return torch.stack(padded)
def _elapsed_for_frame(
self,
file_name: str,
frame_stamps: list[FrameStamp] | None,
) -> float:
if frame_stamps is not None:
return elapsed_seconds_by_file(frame_stamps).get(file_name, -1.0)
return -1.0
def _phase_for_frame(
self,
file_name: str,
frame_stamps: list[FrameStamp] | None,
) -> int:
if frame_stamps is not None:
seconds = elapsed_seconds_by_file(frame_stamps).get(file_name)
if seconds is not None:
return _phase_id_from_seconds(seconds)
return -1
[docs]
def sample(self, sample: tuple[str, typing.Any, list[FrameStamp] | None]) -> Sample:
exam_folder_rel, _, frame_stamps = sample
exam_dir = self.datadir / exam_folder_rel
if frame_stamps is not None:
names = frame_names(frame_stamps)
frame_paths = [
exam_dir / name for name in names if (exam_dir / name).is_file()
]
frame_paths = sorted(frame_paths, key=self._frame_sort_key)
else:
frame_paths = sorted(exam_dir.glob("*.jpg"), key=self._frame_sort_key)
if not frame_paths:
raise FileNotFoundError(f"No .jpg frames found in `{exam_dir}`.")
if self.max_frames_span == "one_per_phase":
slots = self._select_one_per_phase_slots(frame_paths, frame_stamps)
frames_tensors: list[torch.Tensor] = []
phase_ids: list[int] = []
elapsed_seconds: list[float] = []
for path, phase in slots:
phase_ids.append(phase)
if path is None:
elapsed_seconds.append(-1.0)
continue
frame = torch.as_tensor(self._load_one_frame(path))
frames_tensors.append(frame)
elapsed_seconds.append(self._elapsed_for_frame(path.name, frame_stamps))
if not frames_tensors:
raise FileNotFoundError(f"No usable frames for exam `{exam_dir}`.")
template = frames_tensors[0]
padded_frames: list[torch.Tensor] = []
fi = 0
for path, phase in slots:
if path is None:
padded_frames.append(torch.zeros_like(template))
else:
padded_frames.append(frames_tensors[fi])
fi += 1
images = tv_tensors.Image(self._pad_to_max_hw(padded_frames))
slot_names = [path.name if path is not None else "" for path, _ in slots]
return dict(
image=images,
target=self.target(sample),
# Fixed early→mid→late length; padding slots use ``PHASE_ID_PADDING``.
lengths=NUM_PHASES,
name=exam_folder_rel,
phase_id=torch.tensor(phase_ids, dtype=torch.long),
elapsed_seconds=torch.tensor(elapsed_seconds, dtype=torch.float32),
frame_names=slot_names,
)
frame_paths = self._select_frames(frame_paths, frame_stamps)
frames = [self._load_one_frame(p) for p in frame_paths]
images = tv_tensors.Image(
self._pad_to_max_hw([torch.as_tensor(f) for f in frames])
)
phase_ids = [
self._phase_for_frame(path.name, frame_stamps) for path in frame_paths
]
elapsed_seconds = [
self._elapsed_for_frame(path.name, frame_stamps) for path in frame_paths
]
return dict(
image=images,
target=self.target(sample),
lengths=len(frame_paths),
name=exam_folder_rel,
phase_id=torch.tensor(phase_ids, dtype=torch.long),
elapsed_seconds=torch.tensor(elapsed_seconds, dtype=torch.float32),
frame_names=[p.name for p in frame_paths],
)
[docs]
def target(
self, sample: tuple[str, typing.Any, list[FrameStamp] | None]
) -> torch.Tensor:
_, label, _ = sample
if self.problem_type == "binary":
return torch.FloatTensor([label])
if self.problem_type == "multilabel":
return binarize_findings(label)
return torch.LongTensor([label]).squeeze()
[docs]
def collate_sequence(batch: list[Sample]) -> dict[str, typing.Any]:
"""Collate variable-length sequences with zero-padding on the time axis.
Parameters
----------
batch
List of sample dicts from :class:`SequenceRawDataLoader`.
Returns
-------
dict[str, typing.Any]
Batched dict with keys ``image``, ``target``, ``lengths``, ``name``,
``phase_id``, and ``elapsed_seconds``.
"""
images = [torch.as_tensor(sample["image"]) for sample in batch]
targets = torch.utils.data.default_collate([sample["target"] for sample in batch])
lengths = torch.tensor([sample["lengths"] for sample in batch], dtype=torch.long)
names = [sample["name"] for sample in batch]
phase_ids = [
torch.as_tensor(sample["phase_id"], dtype=torch.long) for sample in batch
]
elapsed_seconds = [
torch.as_tensor(sample["elapsed_seconds"], dtype=torch.float32)
for sample in batch
]
padded = torch.nn.utils.rnn.pad_sequence(images, batch_first=True)
phase_padded = torch.nn.utils.rnn.pad_sequence(
phase_ids,
batch_first=True,
padding_value=-1,
)
elapsed_padded = torch.nn.utils.rnn.pad_sequence(
elapsed_seconds,
batch_first=True,
padding_value=-1.0,
)
return dict(
image=padded,
target=targets,
lengths=lengths,
name=names,
phase_id=phase_padded,
elapsed_seconds=elapsed_padded,
)
[docs]
class DataModule(CachingDataModule):
"""Temporal AngioReport dataset with multiple frames per exam.
Parameters
----------
split_path
Path to the split JSON file.
num_classes
Number of output classes.
problem_type
Target formatting strategy (``"binary"``, ``"multiclass"``, or
``"multilabel"``).
modality
Modality to crop from side-by-side frames (``"FA"`` or ``"ICGA"``).
max_frames
Cap on the number of frames per sequence; ``None`` loads all frames.
max_frames_span
Strategy for down-sampling when the sequence exceeds *max_frames*.
"""
def __init__(
self,
split_path: pathlib.Path | resources_abc.Traversable,
num_classes: int,
problem_type: typing.Literal["binary", "multiclass", "multilabel"],
modality: typing.Literal["FA", "ICGA"] = "FA",
max_frames: int | None = 32,
max_frames_span: typing.Literal[
"uniform", "last", "phase_stratified", "one_per_phase"
] = "uniform",
) -> None:
super().__init__(
database_split=GroupedJSONSplit(split_path),
raw_data_loader=SequenceRawDataLoader(
problem_type=problem_type,
modality=modality,
max_frames=max_frames,
max_frames_span=max_frames_span,
),
database_name=DATABASE_SLUG,
split_name=split_path.name.rsplit(".", 2)[0] + "_seq",
task="classification",
num_classes=num_classes,
collate_fn=collate_sequence,
)