""" NIH ChestX-ray14 PyTorch Dataset. Handles: - Multi-label binary encoding for 14 pathology classes - Patient-level train/val/test split (no patient leakage) - On-the-fly augmentation - Demographic metadata for fairness auditing """ from __future__ import annotations import ast from pathlib import Path from typing import Optional import numpy as np import pandas as pd import torch from PIL import Image from torch.utils.data import Dataset from data.transforms import build_transforms # Canonical class order — never reorder; indices are baked into checkpoints. CLASSES = [ "Atelectasis", "Cardiomegaly", "Effusion", "Infiltration", "Mass", "Nodule", "Pneumonia", "Pneumothorax", "Consolidation", "Edema", "Emphysema", "Fibrosis", "Pleural_Thickening", "Hernia", ] CLASS_TO_IDX = {c: i for i, c in enumerate(CLASSES)} NUM_CLASSES = len(CLASSES) class ChestXrayDataset(Dataset): """ NIH ChestX-ray14 dataset with multi-label targets. Args: root_dir: Directory containing image files. labels_csv: Path to Data_Entry_2017.csv. split_list: Text file listing image filenames for this split. transform: Torchvision transform applied to each image. return_metadata: If True, also return (age, gender) for fairness analysis. """ def __init__( self, root_dir: str | Path, labels_csv: str | Path, split_list: str | Path, transform=None, return_metadata: bool = False, ) -> None: self.root_dir = Path(root_dir) self.transform = transform self.return_metadata = return_metadata # Load and filter the label dataframe to only this split's files. df = pd.read_csv(labels_csv) with open(split_list) as f: split_files = set(line.strip() for line in f if line.strip()) self.df = df[df["Image Index"].isin(split_files)].reset_index(drop=True) # Pre-encode multi-hot label vectors once (faster than per-item encoding). self._labels = self._encode_labels() # Normalize age: clip outliers and convert to float. self.df["Patient Age"] = self.df["Patient Age"].clip(0, 100).astype(float) # ------------------------------------------------------------------ # Label encoding # ------------------------------------------------------------------ def _encode_labels(self) -> np.ndarray: """Return (N, 14) float32 multi-hot array.""" labels = np.zeros((len(self.df), NUM_CLASSES), dtype=np.float32) for i, finding_str in enumerate(self.df["Finding Labels"]): for finding in finding_str.split("|"): finding = finding.strip() if finding in CLASS_TO_IDX: labels[i, CLASS_TO_IDX[finding]] = 1.0 return labels # ------------------------------------------------------------------ # Class-level positive weights for loss weighting (call on train split) # ------------------------------------------------------------------ def get_pos_weights(self) -> torch.Tensor: """ Compute per-class positive weights = (neg_count / pos_count). Pass to nn.BCEWithLogitsLoss(pos_weight=...) to handle imbalance. """ pos = self._labels.sum(axis=0) neg = len(self._labels) - pos weights = neg / np.maximum(pos, 1) # avoid div-by-zero for Hernia return torch.tensor(weights, dtype=torch.float32) # ------------------------------------------------------------------ # Dataset interface # ------------------------------------------------------------------ def __len__(self) -> int: return len(self.df) def __getitem__(self, idx: int): row = self.df.iloc[idx] img_path = self.root_dir / row["Image Index"] # Load as RGB (some NIH images are grayscale PNGs). image = Image.open(img_path).convert("RGB") if self.transform: image = self.transform(image) label = torch.from_numpy(self._labels[idx]) if self.return_metadata: meta = { "age": float(row["Patient Age"]), "gender": row["Patient Gender"].strip(), # 'M' or 'F' "filename": row["Image Index"], } return image, label, meta return image, label # ------------------------------------------------------------------ # Utility # ------------------------------------------------------------------ @property def class_names(self) -> list[str]: return CLASSES def label_stats(self) -> pd.DataFrame: """Return a DataFrame with per-class prevalence for EDA.""" counts = self._labels.sum(axis=0) prev = counts / len(self._labels) return pd.DataFrame({ "class": CLASSES, "positive_count": counts.astype(int), "prevalence": prev, }).sort_values("prevalence", ascending=False) def build_datasets(cfg: dict) -> dict[str, ChestXrayDataset]: """ Build train / val / test datasets from config dict. Returns: {'train': ..., 'val': ..., 'test': ...} """ train_tf = build_transforms("train", cfg["data"]["image_size"]) val_tf = build_transforms("val", cfg["data"]["image_size"]) common = dict( root_dir=cfg["data"]["root_dir"], labels_csv=cfg["data"]["labels_csv"], ) # The official NIH split file covers both train AND val; we sub-split below. full_train = ChestXrayDataset( **common, split_list=cfg["data"]["train_val_list"], transform=train_tf, ) full_val = ChestXrayDataset( **common, split_list=cfg["data"]["train_val_list"], transform=val_tf, return_metadata=True, ) test_ds = ChestXrayDataset( **common, split_list=cfg["data"]["test_list"], transform=val_tf, return_metadata=True, ) # Patient-level 90/10 train/val split to prevent patient leakage. patient_ids = full_train.df["Patient ID"].unique() rng = np.random.default_rng(42) rng.shuffle(patient_ids) split_pt = int(0.9 * len(patient_ids)) train_patients = set(patient_ids[:split_pt]) val_patients = set(patient_ids[split_pt:]) train_mask = full_train.df["Patient ID"].isin(train_patients).values val_mask = full_val.df["Patient ID"].isin(val_patients).values # Apply masks to create proper subsets without data leakage. train_ds = _mask_dataset(full_train, train_mask) val_ds = _mask_dataset(full_val, val_mask) return {"train": train_ds, "val": val_ds, "test": test_ds} def _mask_dataset(ds: ChestXrayDataset, mask: np.ndarray) -> ChestXrayDataset: """Return a view of ds restricted to rows where mask is True.""" ds_copy = object.__new__(ChestXrayDataset) ds_copy.__dict__.update(ds.__dict__) ds_copy.df = ds.df[mask].reset_index(drop=True) ds_copy._labels = ds._labels[mask] return ds_copy