""" Load axial slices from a directory of NIfTI CT volumes. Each ``.nii.gz`` file is treated as one patient (the filename stem is the ground-truth patient ID). A handful of evenly-spaced axial slices are sampled per volume, HU-windowed, and returned both as grayscale thumbnails (for the manuscript's clustering method) and as full-resolution slices (for CNN feature extraction). This is used to validate the patient-clustering method on a dataset where the true patient identity IS known — e.g. the Medical Segmentation Decathlon Task06_Lung set — unlike IQ-OTH/NCCD, where patient IDs are unavailable. """ import os import glob import numpy as np class NiftiSliceDataset: """Sample HU-windowed axial slices from a folder of NIfTI CT volumes. Parameters ---------- volume_dir : str Directory containing ``*.nii.gz`` volumes, one per patient. n_slices_per_patient : int Number of evenly-spaced axial slices to sample per volume. thumbnail_size : tuple of int Grayscale thumbnail size for clustering (matches the manuscript's 64x64). hu_window : tuple of float (low, high) Hounsfield-unit clip range applied before scaling to 8-bit. The default (-1000, 400) keeps the body outline, soft tissue, and lung parenchyma — the cues that identify a patient. central_fraction : float Fraction of the volume (centred on the middle slice) to sample from, avoiding empty/partial slices at the superior/inferior ends. slice_axis : int Axis along which axial slices are indexed (-1 = last, correct for the (H, W, Z) layout used by the Decathlon volumes). """ def __init__(self, volume_dir, n_slices_per_patient=10, thumbnail_size=(64, 64), hu_window=(-1000.0, 400.0), central_fraction=0.6, slice_axis=-1, random_slices=True, seed=42, flip_vertical=False, rotate_deg=0): self.volume_dir = volume_dir self.n_slices_per_patient = n_slices_per_patient self.thumbnail_size = thumbnail_size self.hu_window = hu_window self.central_fraction = central_fraction self.slice_axis = slice_axis self.random_slices = random_slices self.seed = seed self.flip_vertical = flip_vertical self.rotate_deg = rotate_deg # 0, 90, 180, or 270 # Populated by load() self.thumbnails = None # (n_slices, H*W) float32 in [0, 1] self.slices_rgb = None # list of (H, W) uint8 arrays self.patient_labels = None # (n_slices,) ground-truth patient IDs self.filenames = None # (n_slices,) per-slice names def load(self, verbose=True): """Load and slice every volume in ``volume_dir``.""" import nibabel as nib from PIL import Image # Prefer decompressed .nii (fast memmap random-slice access) and fall # back to .nii.gz, skipping any .gz that has an uncompressed twin. nii = glob.glob(os.path.join(self.volume_dir, "*.nii")) stems = {os.path.basename(p)[:-4] for p in nii} gz = [p for p in glob.glob(os.path.join(self.volume_dir, "*.nii.gz")) if os.path.basename(p)[:-7] not in stems] paths = sorted(nii + gz) if not paths: raise FileNotFoundError( f"No .nii/.nii.gz volumes found in {self.volume_dir}") thumbs, slices_rgb, patient_labels, filenames = [], [], [], [] lo, hi = self.hu_window for path in paths: patient_id = os.path.basename(path).replace(".nii.gz", "").replace(".nii", "") img = nib.load(path) n_z = img.shape[self.slice_axis] axis = self.slice_axis % img.ndim for k in self._slice_indices(n_z): # Lazy slice: pull only this plane off disk instead of # materialising the whole float64 volume with get_fdata(). slicer = [slice(None)] * img.ndim slicer[axis] = k sl = np.asarray(img.dataobj[tuple(slicer)], dtype=np.float32) # HU window -> 8-bit grayscale sl = np.clip(sl, lo, hi) sl = (255.0 * (sl - lo) / (hi - lo)).astype(np.uint8) if self.flip_vertical: sl = np.flipud(sl) if self.rotate_deg: sl = np.rot90(sl, k=self.rotate_deg // 90) thumb = Image.fromarray(sl).convert("L").resize( self.thumbnail_size) thumbs.append(np.asarray(thumb, dtype=np.float32).flatten() / 255.0) slices_rgb.append(sl) patient_labels.append(patient_id) filenames.append(f"{patient_id}_slice{k:03d}") if verbose: print(f" {patient_id}: {n_z} slices -> " f"sampled {self.n_slices_per_patient}") self.thumbnails = np.array(thumbs, dtype=np.float32) self.slices_rgb = slices_rgb self.patient_labels = np.array(patient_labels) self.filenames = np.array(filenames) if verbose: print(f"\n Loaded {len(self.filenames)} slices from " f"{len(paths)} patients") return self def load_all_slices(self, stride=3, verbose=True): """Load ALL slices from the central fraction (not just a sample). Every ``stride``-th slice is taken to keep feature extraction manageable while still giving K-means enough data to cluster per patient. Returns self (populates thumbnails, slices_rgb, patient_labels, filenames). """ import nibabel as nib from PIL import Image nii = glob.glob(os.path.join(self.volume_dir, "*.nii")) stems = {os.path.basename(p)[:-4] for p in nii} gz = [p for p in glob.glob(os.path.join(self.volume_dir, "*.nii.gz")) if os.path.basename(p)[:-7] not in stems] paths = sorted(nii + gz) if not paths: raise FileNotFoundError( f"No .nii/.nii.gz volumes found in {self.volume_dir}") thumbs, slices_rgb, patient_labels, filenames = [], [], [], [] lo, hi = self.hu_window total_slices = 0 for path in paths: patient_id = os.path.basename(path).replace(".nii.gz", "").replace(".nii", "") img = nib.load(path) n_z = img.shape[self.slice_axis] axis = self.slice_axis % img.ndim # All slices in central fraction, stepped half = self.central_fraction / 2.0 start = int(n_z * (0.5 - half)) end = int(n_z * (0.5 + half)) indices = list(range(start, end, stride)) for k in indices: slicer = [slice(None)] * img.ndim slicer[axis] = k sl = np.asarray(img.dataobj[tuple(slicer)], dtype=np.float32) sl = np.clip(sl, lo, hi) sl = (255.0 * (sl - lo) / (hi - lo)).astype(np.uint8) if self.flip_vertical: sl = np.flipud(sl) if self.rotate_deg: sl = np.rot90(sl, k=self.rotate_deg // 90) thumb = Image.fromarray(sl).convert("L").resize( self.thumbnail_size) thumbs.append(np.asarray(thumb, dtype=np.float32).flatten() / 255.0) slices_rgb.append(sl) patient_labels.append(patient_id) filenames.append(f"{patient_id}_slice{k:03d}") total_slices += 1 if verbose: print(f" {patient_id}: {n_z} slices -> {len(indices)} sampled") self.thumbnails = np.array(thumbs, dtype=np.float32) self.slices_rgb = slices_rgb self.patient_labels = np.array(patient_labels) self.filenames = np.array(filenames) if verbose: print(f"\n Loaded {total_slices} slices from " f"{len(paths)} patients (stride={stride})") return self def export_pngs(self, cache_dir): """Export loaded slices as PNG files for fast reloading. Creates ``cache_dir/`` with one PNG per slice and a manifest.json mapping paths → patient IDs and z-indices. Returns (paths, patient_ids, z_indices). """ import json from PIL import Image as PILImage os.makedirs(cache_dir, exist_ok=True) manifest_path = os.path.join(cache_dir, "manifest.json") if os.path.exists(manifest_path): print(f" Loading cached PNGs from {cache_dir}") with open(manifest_path) as f: manifest = json.load(f) return (manifest["paths"], np.array(manifest["patient_ids"]), np.array(manifest["z_indices"])) if self.slices_rgb is None or self.patient_labels is None: raise RuntimeError("Call load() or load_all_slices() first.") print(f" Exporting {len(self.slices_rgb)} PNGs to {cache_dir} ...") paths, pids, zs = [], [], [] for i, (sl, pid, fname) in enumerate( zip(self.slices_rgb, self.patient_labels, self.filenames)): out_path = os.path.join(cache_dir, f"{fname}.png") PILImage.fromarray(sl).save(out_path) paths.append(out_path) pids.append(pid) zs.append(int(fname.split("_slice")[-1])) if (i + 1) % 500 == 0: print(f" {i + 1}/{len(self.slices_rgb)} ...", flush=True) manifest = {"paths": paths, "patient_ids": pids, "z_indices": zs} with open(manifest_path, "w") as f: json.dump(manifest, f) print(f" Exported {len(paths)} PNGs") return paths, np.array(pids), np.array(zs) @property def n_patients(self): return len(np.unique(self.patient_labels)) def _slice_indices(self, n_z): """Slice indices within the central fraction (even or random).""" half = self.central_fraction / 2.0 start = int(n_z * (0.5 - half)) end = int(n_z * (0.5 + half)) if self.random_slices: rng = np.random.default_rng(self.seed) return np.sort(rng.integers(start, end, size=self.n_slices_per_patient)) else: return np.linspace(start, end - 1, self.n_slices_per_patient).astype(int)