""" data/dataset.py MultiHeadOCTDataset — PyTorch Dataset for the Multi-Head ConvNeXt architecture. Parses `config/hierarchy.yaml` and maps files to structured multi-head tensors. """ import logging import os from pathlib import Path from typing import Dict, List, Optional, Tuple import numpy as np import pandas as pd import yaml import torch from torch.utils.data import Dataset, DataLoader from sklearn.model_selection import StratifiedKFold logger = logging.getLogger(__name__) _VALID_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"} class MultiHeadOCTDataset(Dataset): """ Dataset for the Multi-Head OCT pipeline. Args: config_path: Path to hierarchy.yaml. data_root: Optional override for the root data directory. fold_indices: Numpy integer array of row indices for this fold. transform: MONAI transform pipeline (must accept file paths, as LoadImage is used). """ def __init__( self, config_path: str, data_root: Optional[str] = None, fold_indices: Optional[np.ndarray] = None, transform = None, ) -> None: super().__init__() self.transform = transform with open(config_path, "r") as f: self._cfg = yaml.safe_load(f) default_root = os.environ.get("OCT_DATA_ROOT", self._cfg.get("data_root", "")) self._data_root = Path(data_root) if data_root else Path(default_root) self._l1_labels = self._cfg["l1_labels"] self._l2_labels = self._cfg["l2_labels"] self._l3_specs = self._cfg["l3_specialists"] self._manifest: pd.DataFrame = self._build_manifest() if fold_indices is not None: self._manifest = self._manifest.iloc[fold_indices].reset_index(drop=True) logger.info( "MultiHeadOCTDataset initialized: %d samples.", len(self._manifest), ) def _build_manifest(self) -> pd.DataFrame: records = [] for entry in self._cfg["class_map"]: dir_path = self._data_root / entry["path"] if not dir_path.exists(): logger.debug(f"Skipping {dir_path} as it does not exist.") continue spec_key = entry.get("l3_specialist") l3_class = entry.get("l3_class") for img_path in sorted(dir_path.iterdir()): if img_path.suffix.lower() not in _VALID_EXTENSIONS: continue records.append({ "image_path": str(img_path), "l1_idx": self._l1_labels.get(entry["l1"], 0), "l2_idx": self._l2_labels.get(entry.get("l2"), -1) if entry.get("l2") else -1, "spec_key": spec_key, "l3_class": l3_class }) return pd.DataFrame(records) def compute_class_weights(self, target="l2") -> torch.Tensor: """ Computes inverse-frequency class weights for a given target level. target: 'l1', 'l2', or a specialist key (e.g., 'Macular') """ if target == "l1": counts = self._manifest['l1_idx'].value_counts().sort_index() weights = 1.0 / counts weights = weights / weights.sum() * len(counts) return torch.tensor(weights.values, dtype=torch.float32) elif target == "l2": # Exclude Normal (-1) for L2 weights df_abnormal = self._manifest[self._manifest['l2_idx'] != -1] if df_abnormal.empty: return torch.ones(5) counts = df_abnormal['l2_idx'].value_counts().sort_index() # Ensure all 5 classes are represented in counts for i in range(5): if i not in counts: counts[i] = 1 # avoid inf counts = counts.sort_index() weights = 1.0 / counts weights = weights / weights.sum() * len(counts) return torch.tensor(weights.values, dtype=torch.float32) else: df_spec = self._manifest[self._manifest['spec_key'] == target] if df_spec.empty: return torch.ones(1) # Map l3_class string to index class_map = self._l3_specs[target]["classes"] class_indices = df_spec['l3_class'].map(class_map) counts = class_indices.value_counts().sort_index() num_classes = len(class_map) for i in range(num_classes): if i not in counts: counts[i] = 1 counts = counts.sort_index() weights = 1.0 / counts weights = weights / weights.sum() * len(counts) return torch.tensor(weights.values, dtype=torch.float32) def __len__(self) -> int: return len(self._manifest) def __getitem__(self, idx: int) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: row = self._manifest.iloc[idx] image_path = row["image_path"] # MONAI pipelines usually start with LoadImage which takes the file path if self.transform is not None: try: # Apply pipeline image = self.transform(image_path) # If MONAI returns a MetaTensor, convert to standard torch.Tensor if hasattr(image, "as_tensor"): image = image.as_tensor() elif isinstance(image, np.ndarray): image = torch.from_numpy(image) except Exception as exc: logger.error("Failed to process %s: %s", image_path, exc) # Return black placeholder for stability image = torch.zeros(3, 384, 384, dtype=torch.float32) else: image = image_path # ── Target H1: Binary (Normal=0, Abnormal=1) ── # Cast to float32 for BCEWithLogitsLoss h1 = torch.tensor([row["l1_idx"]], dtype=torch.float32) # ── Target H2: Router (Multi-class 0-4, or -1 for Normal) ── h2 = torch.tensor(row["l2_idx"], dtype=torch.long) # ── Target H3: Severity (5 Sub-tensors) ── # Initialized as zeros (negative for BCE/Multi-Label if activated) macular = torch.zeros(3, dtype=torch.float32) diabetic = torch.zeros(2, dtype=torch.float32) vascular = torch.zeros(3, dtype=torch.float32) fluid = torch.zeros(1, dtype=torch.float32) structural = torch.zeros(2, dtype=torch.float32) spec = row["spec_key"] l3_cls = row["l3_class"] if pd.notna(spec) and pd.notna(l3_cls) and spec and l3_cls: class_idx = self._l3_specs[spec]["classes"].get(l3_cls, -1) if class_idx != -1: if spec == "Macular": macular[class_idx] = 1.0 elif spec == "Diabetic": diabetic[class_idx] = 1.0 elif spec == "Vascular": vascular[class_idx] = 1.0 elif spec == "Fluid": fluid[class_idx] = 1.0 elif spec == "Structural": structural[class_idx] = 1.0 targets = { "normal_abnormal": h1, "pathology": h2, "severity": { "macular": macular, "diabetic": diabetic, "vascular": vascular, "fluid": fluid, "structural": structural } } return image, targets def build_dataloader( config_path: str, data_root: str, batch_size: int = 16, num_workers: int = 4, transform = None, shuffle: bool = True ) -> DataLoader: """ Utility to quickly build a DataLoader (e.g. for the micro_dataset sanity test). """ dataset = MultiHeadOCTDataset( config_path=config_path, data_root=data_root, transform=transform ) _pin = torch.cuda.is_available() loader = DataLoader( dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers, pin_memory=_pin, persistent_workers=(num_workers > 0) ) return loader def build_kfold_dataloaders( config_path: str, mode: str, n_splits: int = 5, batch_size: int = 32, num_workers: int = 4, train_transform = None, val_transform = None, use_weighted_sampler: bool = False, seed: int = 42 ) -> List[Tuple[DataLoader, DataLoader]]: dataset = MultiHeadOCTDataset(config_path=config_path) # Stratify by l1 for level1, l2 for level2, etc. (simplifying here to l1) labels = [row["l1_idx"] for _, row in dataset._manifest.iterrows()] skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed) fold_loaders = [] for train_idx, val_idx in skf.split(np.zeros(len(labels)), labels): train_ds = MultiHeadOCTDataset(config_path=config_path, fold_indices=train_idx, transform=train_transform) val_ds = MultiHeadOCTDataset(config_path=config_path, fold_indices=val_idx, transform=val_transform) train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers) val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers) fold_loaders.append((train_loader, val_loader)) return fold_loaders