v4 update
This commit is contained in:
rpotter6298
@@ -0,0 +1,248 @@
|
||||
"""dataset — data packaging for v4: shells, DataBundle, HTDataset.
|
||||
|
||||
ShellEntry / LoaderShell
|
||||
Minimal, data-free structures representing *who* to sample and in what
|
||||
order. entity_id is opaque to the orchestrator; towers interpret it.
|
||||
|
||||
DataBundle
|
||||
Accumulates per-eye DataFrames and derives scalar/categorical stats used
|
||||
by ClinicalDataView. No kfold, no vectorization — those live in the
|
||||
profile and orchestrator respectively.
|
||||
|
||||
HTDataset / ht_collate
|
||||
PyTorch Dataset that delegates sample retrieval to towers via get_sample.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Optional
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shell
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@dataclass
|
||||
class ShellEntry:
|
||||
"""One sample slot in a LoaderShell.
|
||||
|
||||
entity_id : opaque — defined by the profile, interpreted by towers.
|
||||
label : integer class label.
|
||||
meta : per-entry context a profile wants to pass through.
|
||||
"""
|
||||
entity_id: Any
|
||||
label: int
|
||||
meta: dict = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class LoaderShell:
|
||||
"""An ordered sequence of ShellEntry objects for one split/fold."""
|
||||
entries: list[ShellEntry]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.entries)
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self.entries)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DataBundle
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class DataBundle:
|
||||
"""Accumulates per-eye DataFrames and derives feature metadata.
|
||||
|
||||
Responsibilities: column type inference, scalar stats (min/max/median),
|
||||
categorical index maps, and feature dim. Everything else (splits,
|
||||
vectorization, image paths) lives in the profile that uses this bundle.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
image_dir: str,
|
||||
clinical_dir: Optional[str] = None,
|
||||
label_col: str,
|
||||
patient_col: str = "Patient ID",
|
||||
cat_cols: Optional[Iterable[str]] = None,
|
||||
max_unique_for_cat: int = 4,
|
||||
n_splits: int = 5,
|
||||
random_seed: int = 42,
|
||||
filename_template: str = "RET{pid:03d}{eye}.jpg",
|
||||
) -> None:
|
||||
self.image_dir = Path(image_dir)
|
||||
self.label_col = label_col
|
||||
self.patient_col = patient_col
|
||||
self.max_unique_for_cat = max_unique_for_cat
|
||||
self.filename_template = filename_template
|
||||
self.clinical_dir = Path(clinical_dir) if clinical_dir else None
|
||||
|
||||
self.frames: List[pd.DataFrame] = []
|
||||
self.df: pd.DataFrame = pd.DataFrame()
|
||||
self.scalar_cols: List[str] = []
|
||||
self.cat_cols: List[str] = list(cat_cols) if cat_cols else []
|
||||
self.scalar_stats: Dict[str, Dict[str, float]] = {}
|
||||
self.cat_maps: Dict[str, Dict[object, int]] = {}
|
||||
self.feature_dim: int = 0
|
||||
|
||||
def add_df(
|
||||
self,
|
||||
df: pd.DataFrame,
|
||||
*,
|
||||
id_column: Optional[str] = None,
|
||||
exclude_cols: Optional[Iterable[str]] = None,
|
||||
) -> None:
|
||||
df = df.copy()
|
||||
self._ensure_patient_id(df, id_column)
|
||||
if self.label_col not in df.columns:
|
||||
raise ValueError(f"label_col '{self.label_col}' not found in added dataframe")
|
||||
self.frames.append(df)
|
||||
self._refresh_master_df(exclude_cols=exclude_cols)
|
||||
self._infer_or_validate_feature_types(exclude_cols=exclude_cols)
|
||||
self._compute_numeric_stats()
|
||||
self._build_cat_maps()
|
||||
self._compute_feature_dim()
|
||||
|
||||
# ── Internal ─────────────────────────────────────────────────────────────
|
||||
|
||||
def _ensure_patient_id(self, df: pd.DataFrame, id_column: Optional[str]) -> None:
|
||||
if self.patient_col in df.columns:
|
||||
return
|
||||
if id_column and id_column in df.columns:
|
||||
df.rename(columns={id_column: self.patient_col}, inplace=True)
|
||||
return
|
||||
candidates = [
|
||||
c for c in df.columns
|
||||
if c.lower().replace(" ", "") in {"patientid", "patient", "pid"}
|
||||
]
|
||||
if len(candidates) == 1:
|
||||
df.rename(columns={candidates[0]: self.patient_col}, inplace=True)
|
||||
return
|
||||
raise ValueError(
|
||||
f"A '{self.patient_col}' column is required; "
|
||||
f"provide id_column=... if it has a different name."
|
||||
)
|
||||
|
||||
def _refresh_master_df(self, exclude_cols: Optional[Iterable[str]] = None) -> None:
|
||||
self.df = pd.concat(self.frames, axis=0, ignore_index=True)
|
||||
if exclude_cols:
|
||||
self.df = self.df.drop(columns=[c for c in exclude_cols if c in self.df.columns])
|
||||
|
||||
def _infer_or_validate_feature_types(self, exclude_cols: Optional[Iterable[str]] = None) -> None:
|
||||
excluded = set(exclude_cols or []) | {self.label_col, self.patient_col}
|
||||
candidates = [c for c in self.df.columns if c not in excluded]
|
||||
cats = set(self.cat_cols)
|
||||
scalars: set[str] = set()
|
||||
for c in candidates:
|
||||
if c in cats:
|
||||
continue
|
||||
s = self.df[c]
|
||||
as_num = pd.to_numeric(s, errors="coerce")
|
||||
n_uniq = s.dropna().nunique()
|
||||
if as_num.notna().any() and as_num.isna().mean() < 1.0 and n_uniq > self.max_unique_for_cat:
|
||||
scalars.add(c)
|
||||
else:
|
||||
cats.add(c)
|
||||
self.cat_cols = sorted(cats)
|
||||
self.scalar_cols = sorted(scalars)
|
||||
|
||||
def _compute_numeric_stats(self) -> None:
|
||||
self.scalar_stats.clear()
|
||||
for col in self.scalar_cols:
|
||||
vals = pd.to_numeric(self.df[col], errors="coerce").dropna().astype(float).values
|
||||
if vals.size == 0:
|
||||
lo, hi, med = 0.0, 1.0, 0.0
|
||||
else:
|
||||
lo, hi = float(np.min(vals)), float(np.max(vals))
|
||||
med = float(np.median(vals))
|
||||
if hi <= lo:
|
||||
hi = lo + 1.0
|
||||
self.scalar_stats[col] = {"min": lo, "max": hi, "median": med}
|
||||
|
||||
def _build_cat_maps(self) -> None:
|
||||
self.cat_maps.clear()
|
||||
for col in self.cat_cols:
|
||||
cats = [v for v in self.df[col].dropna().unique().tolist()]
|
||||
try:
|
||||
cats = sorted(cats)
|
||||
except Exception:
|
||||
pass
|
||||
mapping: Dict[object, int] = {"<UNK>": 0}
|
||||
for i, v in enumerate(cats, start=1):
|
||||
mapping[v] = i
|
||||
self.cat_maps[col] = mapping
|
||||
|
||||
def _compute_feature_dim(self) -> None:
|
||||
self.feature_dim = (
|
||||
len(self.scalar_cols)
|
||||
+ sum(len(m) for m in self.cat_maps.values())
|
||||
+ len(self.scalar_cols)
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HTDataset / ht_collate
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class HTDataset(Dataset):
|
||||
"""PyTorch Dataset backed by a LoaderShell.
|
||||
|
||||
Delegates sample retrieval to each tower's get_sample(entry).
|
||||
Batch keys are tower names plus "label".
|
||||
"""
|
||||
|
||||
def __init__(self, shell: LoaderShell, towers: dict) -> None:
|
||||
self.entries = shell.entries
|
||||
self.towers = towers
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.entries)
|
||||
|
||||
def __getitem__(self, idx: int) -> dict[str, Any]:
|
||||
entry = self.entries[idx]
|
||||
sample = {
|
||||
"label": torch.tensor(entry.label, dtype=torch.long),
|
||||
"entity_id": entry.entity_id,
|
||||
}
|
||||
for name, tower in self.towers.items():
|
||||
sample[name] = tower.get_sample(entry)
|
||||
return sample
|
||||
|
||||
|
||||
def to_label_tensor(labels, device: torch.device) -> torch.Tensor:
|
||||
"""Normalise a batch of labels (tensor or list) to a long tensor on device."""
|
||||
if torch.is_tensor(labels):
|
||||
return labels.to(device=device, dtype=torch.long)
|
||||
return torch.as_tensor(labels, dtype=torch.long, device=device)
|
||||
|
||||
|
||||
def ht_collate(batch: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
"""Collate HTDataset samples.
|
||||
|
||||
Tensor values are stacked; dict values (side dicts from patient-level
|
||||
shells) are stacked per inner key; everything else becomes a list.
|
||||
"""
|
||||
if not batch:
|
||||
return {}
|
||||
result: dict[str, Any] = {}
|
||||
for key in batch[0]:
|
||||
vals = [s[key] for s in batch]
|
||||
first = vals[0]
|
||||
if isinstance(first, torch.Tensor):
|
||||
result[key] = torch.stack(vals, dim=0)
|
||||
elif isinstance(first, dict):
|
||||
result[key] = {
|
||||
side: torch.stack([v[side] for v in vals], dim=0)
|
||||
for side in first
|
||||
}
|
||||
else:
|
||||
result[key] = vals
|
||||
return result
|
||||
Reference in New Issue
Block a user