rpotter6298
243 lines
7.8 KiB
Python
243 lines
7.8 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any, Callable, Optional
|
|
|
|
import torch
|
|
from torch.utils.data import DataLoader, WeightedRandomSampler
|
|
|
|
from .network_manager import LoaderBundle, PatientSplit
|
|
from .slot_dataset import SlotDataset, slot_collate
|
|
from .profiles.base import SlotDescriptor, SimpleDatasetProfile
|
|
|
|
|
|
def _default_slot_descriptors(patient_col: str, label_col: str) -> dict[str, SlotDescriptor]:
|
|
return {
|
|
"id_1": SlotDescriptor(
|
|
key="id_1",
|
|
kind="id",
|
|
description=f"Patient identifier column ({patient_col})",
|
|
required=True,
|
|
shape_hint="scalar",
|
|
),
|
|
"eye_id_1": SlotDescriptor(
|
|
key="eye_id_1",
|
|
kind="id",
|
|
description="Eye side identifier (OD/OS)",
|
|
required=False,
|
|
shape_hint="scalar",
|
|
),
|
|
"label_1": SlotDescriptor(
|
|
key="label_1",
|
|
kind="label",
|
|
description=f"Label column ({label_col})",
|
|
required=True,
|
|
shape_hint="scalar",
|
|
),
|
|
"image_1": SlotDescriptor(
|
|
key="image_1",
|
|
kind="image",
|
|
description="Primary image slot",
|
|
required=False,
|
|
shape_hint="HWC or CHW",
|
|
),
|
|
"matrix_1": SlotDescriptor(
|
|
key="matrix_1",
|
|
kind="matrix",
|
|
description="Primary matrix slot",
|
|
required=False,
|
|
shape_hint="[feature_dim]",
|
|
),
|
|
}
|
|
|
|
|
|
def _row_to_sample(
|
|
row: Any,
|
|
*,
|
|
clinical: Any,
|
|
patient_col: str,
|
|
label_col: str,
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"id_1": row[patient_col],
|
|
"eye_id_1": str(row.get("eyeID", "")),
|
|
"label_1": row[label_col],
|
|
"image_1": clinical.get_image_path(row) if hasattr(clinical, "get_image_path") else None,
|
|
"matrix_1": clinical.vectorize_row(row) if hasattr(clinical, "vectorize_row") else None,
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class SlotLoaderFactory:
|
|
"""
|
|
Generic loader factory that emits dict batches keyed by slot names.
|
|
"""
|
|
|
|
image_transform: Optional[Callable] = None
|
|
matrix_transform: Optional[Callable] = None
|
|
num_workers: int = 0
|
|
|
|
def build(
|
|
self,
|
|
*,
|
|
clinical: Any,
|
|
split: PatientSplit,
|
|
args: Any,
|
|
fold: int,
|
|
profile: Optional[Any] = None,
|
|
) -> LoaderBundle:
|
|
batch_size = int(getattr(args, "batch_size", 8))
|
|
slot_desc = self._resolve_slot_descriptors(clinical=clinical, profile=profile)
|
|
|
|
train_samples = self._build_samples(split.train, clinical, profile, slot_desc)
|
|
val_samples = self._build_samples(split.val, clinical, profile, slot_desc)
|
|
holdout_samples = (
|
|
self._build_samples(split.holdout, clinical, profile, slot_desc)
|
|
if split.holdout is not None
|
|
else None
|
|
)
|
|
|
|
train_loader = DataLoader(
|
|
SlotDataset(
|
|
train_samples,
|
|
slot_desc,
|
|
image_transform=self.image_transform,
|
|
matrix_transform=self.matrix_transform,
|
|
),
|
|
batch_size=batch_size,
|
|
shuffle=True,
|
|
num_workers=self.num_workers,
|
|
collate_fn=slot_collate,
|
|
)
|
|
val_loader = DataLoader(
|
|
SlotDataset(
|
|
val_samples,
|
|
slot_desc,
|
|
image_transform=self.image_transform,
|
|
matrix_transform=self.matrix_transform,
|
|
),
|
|
batch_size=batch_size,
|
|
shuffle=False,
|
|
num_workers=self.num_workers,
|
|
collate_fn=slot_collate,
|
|
)
|
|
holdout_loader = None
|
|
if holdout_samples is not None:
|
|
holdout_loader = DataLoader(
|
|
SlotDataset(
|
|
holdout_samples,
|
|
slot_desc,
|
|
image_transform=self.image_transform,
|
|
matrix_transform=self.matrix_transform,
|
|
),
|
|
batch_size=batch_size,
|
|
shuffle=False,
|
|
num_workers=self.num_workers,
|
|
collate_fn=slot_collate,
|
|
)
|
|
return LoaderBundle(train=train_loader, val=val_loader, holdout=holdout_loader)
|
|
|
|
@staticmethod
|
|
def _resolve_slot_descriptors(
|
|
*,
|
|
clinical: Any,
|
|
profile: Optional[Any],
|
|
) -> dict[str, SlotDescriptor]:
|
|
if profile is not None and hasattr(profile, "slot_descriptors"):
|
|
return profile.slot_descriptors()
|
|
patient_col = getattr(clinical, "patient_col", "Patient ID")
|
|
label_col = getattr(clinical, "label_col", "Diagnosis")
|
|
return _default_slot_descriptors(patient_col, label_col)
|
|
|
|
@staticmethod
|
|
def _build_samples(
|
|
df,
|
|
clinical: Any,
|
|
profile: Optional[Any],
|
|
slot_desc: dict[str, SlotDescriptor],
|
|
) -> list[dict[str, Any]]:
|
|
if df is None or df.empty:
|
|
return []
|
|
if profile is not None and hasattr(profile, "build_samples"):
|
|
return profile.build_samples(df=df, clinical=clinical)
|
|
|
|
patient_col = getattr(profile, "patient_col", None) if profile is not None else None
|
|
label_col = getattr(profile, "label_col", None) if profile is not None else None
|
|
pcol = patient_col or "Patient ID"
|
|
lcol = label_col or getattr(clinical, "label_col", "Diagnosis")
|
|
samples = []
|
|
for _, row in df.iterrows():
|
|
sample = _row_to_sample(row, clinical=clinical, patient_col=pcol, label_col=lcol)
|
|
for key in slot_desc.keys():
|
|
sample.setdefault(key, None)
|
|
samples.append(sample)
|
|
return samples
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# V2 filter / loader helpers (used by V2HyperTower._run_fold)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def filter_eye_samples(samples: list[dict]) -> list[dict]:
|
|
"""Keep any single-eye sample with a valid image, matrix, and label."""
|
|
return [
|
|
s for s in samples
|
|
if s.get("image_1") is not None
|
|
and s.get("matrix_1") is not None
|
|
and s.get("label_1") is not None
|
|
]
|
|
|
|
|
|
def filter_bilateral_samples(samples: list[dict]) -> list[dict]:
|
|
"""Keep only patient-level samples where both eyes are fully present."""
|
|
return [
|
|
s for s in samples
|
|
if s.get("image_1") is not None
|
|
and s.get("matrix_1") is not None
|
|
and s.get("image_2") is not None
|
|
and s.get("matrix_2") is not None
|
|
and s.get("label_1") is not None
|
|
]
|
|
|
|
|
|
def make_loader(
|
|
samples: list[dict],
|
|
slots: dict,
|
|
*,
|
|
image_transform,
|
|
image_preprocessor=None,
|
|
batch_size: int,
|
|
shuffle: bool,
|
|
num_workers: int,
|
|
sampler: Optional[WeightedRandomSampler] = None,
|
|
) -> DataLoader:
|
|
ds = SlotDataset(
|
|
samples,
|
|
slots,
|
|
image_transform=image_transform,
|
|
image_preprocessor=image_preprocessor,
|
|
)
|
|
return DataLoader(
|
|
ds,
|
|
batch_size=batch_size,
|
|
shuffle=(shuffle if sampler is None else False),
|
|
sampler=sampler,
|
|
num_workers=num_workers,
|
|
collate_fn=slot_collate,
|
|
)
|
|
|
|
|
|
def build_balanced_sampler(samples: list[dict], label_key: str = "label_1") -> WeightedRandomSampler:
|
|
"""Return a WeightedRandomSampler that equalises class frequency for training."""
|
|
from collections import Counter
|
|
labels = [s[label_key] for s in samples]
|
|
counts = Counter(labels)
|
|
weights = [1.0 / counts[lbl] for lbl in labels]
|
|
return WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)
|
|
|
|
|
|
def to_label_tensor(labels, device: torch.device) -> torch.Tensor:
|
|
if torch.is_tensor(labels):
|
|
return labels.to(device=device, dtype=torch.long)
|
|
return torch.as_tensor(labels, dtype=torch.long, device=device)
|