diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e5f29ef --- /dev/null +++ b/.gitignore @@ -0,0 +1,21 @@ +# Python +__pycache__/ +*.py[cod] +*.pyo + +# Data +CROP1/ +CROP EXT/ +*.xlsx +*.nrrd +*.dcm + +# Outputs / models +try2/outputs/ +try3/results/ +*.joblib +*.pth +*.pt + +# System +.DS_Store diff --git a/data_utils.py b/data_utils.py new file mode 100644 index 0000000..be7e991 --- /dev/null +++ b/data_utils.py @@ -0,0 +1,534 @@ +#!/usr/bin/env python3 +""" +Data Utilities for Medical Image Segmentation +============================================= +Handles loading and aligning DICOM images with NRRD segmentation files, +accounting for different coordinate systems used by 3D Slicer. +""" + +import numpy as np +import pydicom +import nrrd +from pathlib import Path +from typing import Tuple, Optional, Dict, List +from tqdm import tqdm +import warnings + + +def load_dicom_series(dicom_dir: str, verbose: bool = True) -> Tuple[np.ndarray, Dict]: + """ + Load a DICOM series from a directory. + + Returns: + volume: 3D numpy array (H, W, D) + metadata: Dictionary with spatial information + """ + dicom_dir = Path(dicom_dir) + dcm_files = list(dicom_dir.glob("*.dcm")) + list(dicom_dir.glob("*.DCM")) + + if not dcm_files: + raise FileNotFoundError(f"No DICOM files found in {dicom_dir}") + + if verbose: + print(f"Found {len(dcm_files)} DICOM files") + + # Load all slices + slices = [] + for f in tqdm(dcm_files, desc="Loading DICOMs", disable=not verbose): + try: + ds = pydicom.dcmread(str(f)) + slices.append(ds) + except Exception as e: + if verbose: + print(f" Warning: Failed to load {f.name}: {e}") + + if not slices: + raise ValueError("No valid DICOM files could be loaded") + + # Sort by slice position + try: + slices.sort(key=lambda s: float(s.ImagePositionPatient[2])) + sorted_by = "ImagePositionPatient" + except AttributeError: + try: + slices.sort(key=lambda s: float(s.SliceLocation)) + sorted_by = "SliceLocation" + except AttributeError: + slices.sort(key=lambda s: float(s.InstanceNumber)) + sorted_by = "InstanceNumber" + + # Stack into 3D volume, converting stored pixel values to Hounsfield + # units via the per-slice RescaleSlope / RescaleIntercept tags. HU is + # required for fixed CT windowing; min-max-normalised pipelines are + # unaffected since min-max is invariant to a positive affine rescale. + volume = np.stack([ + s.pixel_array.astype(np.float32) * float(getattr(s, "RescaleSlope", 1.0)) + + float(getattr(s, "RescaleIntercept", 0.0)) + for s in slices + ], axis=-1) + + # Extract metadata + ds = slices[0] + pixel_spacing = list(map(float, getattr(ds, "PixelSpacing", [1.0, 1.0]))) + slice_thickness = float(getattr(ds, "SliceThickness", 1.0)) + + # Calculate slice spacing from positions if available + if len(slices) > 1 and hasattr(slices[0], 'ImagePositionPatient'): + pos0 = np.array(slices[0].ImagePositionPatient) + pos1 = np.array(slices[1].ImagePositionPatient) + slice_spacing = np.linalg.norm(pos1 - pos0) + else: + slice_spacing = slice_thickness + + metadata = { + "pixel_spacing": pixel_spacing, + "slice_thickness": slice_thickness, + "slice_spacing": slice_spacing, + "rows": ds.Rows, + "columns": ds.Columns, + "num_slices": len(slices), + "modality": getattr(ds, "Modality", "Unknown"), + "patient_id": getattr(ds, "PatientID", "Unknown"), + "sorted_by": sorted_by, + "image_position": list(map(float, getattr(ds, "ImagePositionPatient", [0, 0, 0]))), + "image_orientation": list(map(float, getattr(ds, "ImageOrientationPatient", [1, 0, 0, 0, 1, 0]))), + "rescale_slope": float(getattr(ds, "RescaleSlope", 1.0)), + "rescale_intercept": float(getattr(ds, "RescaleIntercept", 0.0)), + } + + if verbose: + print(f"DICOM volume shape: {volume.shape}") + print(f"Pixel spacing: {pixel_spacing}, Slice spacing: {slice_spacing:.2f}") + + return volume, metadata + + +def load_nrrd_segmentation(nrrd_path: str, verbose: bool = True) -> Tuple[np.ndarray, Dict]: + """ + Load NRRD segmentation file with metadata. + + Returns: + segmentation: 3D numpy array with integer labels + metadata: Dictionary with spatial info and segment names + """ + if verbose: + print(f"Loading NRRD from {nrrd_path}") + + data, header = nrrd.read(nrrd_path) + + # Parse space directions (3x3 matrix) + space_directions = header.get("space directions", np.eye(3)) + if isinstance(space_directions, list): + space_directions = np.array(space_directions) + + # Parse space origin + space_origin = header.get("space origin", [0, 0, 0]) + if isinstance(space_origin, list): + space_origin = np.array(space_origin) + + # Parse segment information + segments = {} + for key, value in header.items(): + if key.startswith("Segment") and "_Name" in key: + seg_idx = key.split("_")[0].replace("Segment", "") + label_key = f"Segment{seg_idx}_LabelValue" + if label_key in header: + try: + label = int(header[label_key]) + segments[label] = value + except: + pass + + metadata = { + "space": header.get("space", "unknown"), + "space_directions": space_directions, + "space_origin": space_origin, + "segments": segments, + "header": header, + } + + if verbose: + print(f"NRRD shape: {data.shape}") + print(f"Unique labels: {np.unique(data)}") + print(f"Segments: {segments}") + + return data.astype(np.int64), metadata + + +def align_nrrd_to_dicom( + dicom_volume: np.ndarray, + nrrd_volume: np.ndarray, + nrrd_metadata: Dict, + verbose: bool = True, + dicom_metadata: Optional[Dict] = None, +) -> Tuple[np.ndarray, bool]: + """ + Align NRRD segmentation to DICOM volume using world coordinates. + + Uses DICOM spatial metadata (ImagePositionPatient, ImageOrientationPatient, + PixelSpacing, slice_spacing) and NRRD spatial metadata (space directions, + space origin) to compute a proper coordinate-based mapping. + + Returns: + aligned_nrrd: NRRD aligned to DICOM shape (H, W, D) + success: Whether alignment was successful + """ + if verbose: + print(f"\nAligning: DICOM {dicom_volume.shape} <-> NRRD {nrrd_volume.shape}") + + # Already aligned? + if dicom_volume.shape == nrrd_volume.shape: + if verbose: + print("Shapes already match!") + return nrrd_volume, True + + # --- Coordinate-based alignment (preferred) --- + if dicom_metadata is not None: + result = _align_by_coordinates( + dicom_volume, nrrd_volume, dicom_metadata, nrrd_metadata, verbose + ) + if result is not None: + return result, True + + # --- Fallback: simple permutation (shape-only) --- + result = _try_simple_alignment(dicom_volume, nrrd_volume, verbose) + if result is not None: + return result, True + + if verbose: + print("Could not find alignment") + return nrrd_volume, False + + +def _align_by_coordinates( + dicom_volume: np.ndarray, + nrrd_volume: np.ndarray, + dicom_metadata: Dict, + nrrd_metadata: Dict, + verbose: bool = True, +) -> Optional[np.ndarray]: + """Align using world-coordinate mapping between DICOM and NRRD. + + Builds the affine transforms for both volumes, computes the NRRD voxel + index corresponding to each corner of the DICOM grid, extracts the + sub-volume, and applies any necessary axis flips. + """ + # --- DICOM affine components --- + ipp = np.array(dicom_metadata.get("image_position", [0, 0, 0]), dtype=np.float64) + iop = np.array( + dicom_metadata.get("image_orientation", [1, 0, 0, 0, 1, 0]), dtype=np.float64 + ) + ps = dicom_metadata.get("pixel_spacing", [1.0, 1.0]) + row_spacing, col_spacing = float(ps[0]), float(ps[1]) + slice_spacing = float(dicom_metadata.get("slice_spacing", 1.0)) + + # DICOM ImageOrientationPatient (PS3.3 C.7.6.2.1.1): + # iop[0:3] = "direction cosines of the first row" + # = direction ALONG the row = direction of COLUMN increase + # iop[3:6] = "direction cosines of the first column" + # = direction ALONG the column = direction of ROW increase + col_dir = iop[:3] # direction of increasing column index + row_dir = iop[3:] # direction of increasing row index + slice_dir = np.cross(col_dir, row_dir) # normal to image plane + + # DICOM PixelSpacing: + # ps[0] = row spacing (distance between adjacent rows) + # ps[1] = column spacing (distance between adjacent columns) + # + # DICOM pixel (row=r, col=c, slice=s) -> world: + # world = ipp + c * col_dir * col_spacing + # + r * row_dir * row_spacing + # + s * slice_dir * slice_spacing + + # --- NRRD affine components --- + space_dirs = nrrd_metadata.get("space_directions", np.eye(3)) + if isinstance(space_dirs, list): + space_dirs = np.array(space_dirs, dtype=np.float64) + space_origin = np.array( + nrrd_metadata.get("space_origin", [0, 0, 0]), dtype=np.float64 + ) + + # NRRD voxel (i, j, k) -> world: world = space_origin + space_dirs^T @ [i, j, k] + # We invert: voxel = inv(space_dirs) @ (world - space_origin) + try: + inv_space_dirs = np.linalg.inv(space_dirs) + except np.linalg.LinAlgError: + if verbose: + print("NRRD space_directions matrix is singular, cannot invert") + return None + + H, W, D = dicom_volume.shape + + # Compute NRRD voxel index for DICOM voxel (0,0,0) and the three unit steps + world_origin = ipp # DICOM voxel (0,0,0) + nrrd_origin = inv_space_dirs @ (world_origin - space_origin) + + # Step vectors in NRRD voxel space for one DICOM voxel step + step_row = inv_space_dirs @ (row_dir * row_spacing) + step_col = inv_space_dirs @ (col_dir * col_spacing) + step_slice = inv_space_dirs @ (slice_dir * slice_spacing) + + if verbose: + print(f" NRRD origin (for DICOM 0,0,0): {nrrd_origin}") + print(f" NRRD step per DICOM row: {step_row}") + print(f" NRRD step per DICOM col: {step_col}") + print(f" NRRD step per DICOM slice: {step_slice}") + + # Build the full NRRD index array for every DICOM voxel would be huge. + # Instead, since DICOM->NRRD is an affine mapping of axis-aligned grids, + # each DICOM axis maps to exactly one NRRD axis (with possible sign flip). + # Detect which NRRD axis each DICOM axis maps to. + + steps = np.array([step_row, step_col, step_slice]) # (3, 3) + # For each DICOM axis, find the dominant NRRD axis + axis_map = {} # dicom_axis -> nrrd_axis + axis_sign = {} # dicom_axis -> +1 or -1 + for d_ax in range(3): + abs_step = np.abs(steps[d_ax]) + n_ax = int(np.argmax(abs_step)) + # Verify this is truly axis-aligned (dominant component >> others) + if abs_step[n_ax] < 1e-6: + if verbose: + print(f" DICOM axis {d_ax} has zero step in NRRD space") + return None + off_axis = np.delete(abs_step, n_ax) + if np.any(off_axis > 0.1 * abs_step[n_ax]): + if verbose: + print(f" DICOM axis {d_ax} is not axis-aligned in NRRD space: {steps[d_ax]}") + return None + axis_map[d_ax] = n_ax + axis_sign[d_ax] = 1 if steps[d_ax][n_ax] > 0 else -1 + + # Check we have a valid 1-to-1 mapping + if len(set(axis_map.values())) != 3: + if verbose: + print(f" Axis mapping is not 1-to-1: {axis_map}") + return None + + if verbose: + labels = ["row", "col", "slice"] + for d_ax in range(3): + sign = "+" if axis_sign[d_ax] > 0 else "-" + print(f" DICOM {labels[d_ax]} -> NRRD axis {axis_map[d_ax]} ({sign})") + + # Compute the NRRD index range for each DICOM axis. + # If the DICOM volume extends slightly beyond the NRRD canvas, clamp to + # the valid range and zero-pad afterwards (those edge voxels are background). + dicom_sizes = [H, W, D] + slices_per_axis = [None, None, None] # slice objects for NRRD extraction + pad_before = [0, 0, 0] # padding needed before the extracted region (per NRRD axis) + pad_after = [0, 0, 0] # padding needed after + + for d_ax in range(3): + n_ax = axis_map[d_ax] + start_nrrd = nrrd_origin[n_ax] + step = steps[d_ax][n_ax] + end_nrrd = start_nrrd + step * (dicom_sizes[d_ax] - 1) + + lo = min(start_nrrd, end_nrrd) + hi = max(start_nrrd, end_nrrd) + lo_int = int(round(lo)) + hi_int = int(round(hi)) + expected_size = hi_int - lo_int + 1 + + # Clamp to valid NRRD range, track how much padding is needed + clamped_lo = max(0, lo_int) + clamped_hi = min(nrrd_volume.shape[n_ax] - 1, hi_int) + + pb = clamped_lo - lo_int # voxels clipped at the low end + pa = hi_int - clamped_hi # voxels clipped at the high end + + # Reject if more than 30% of the axis is out of bounds + if pb + pa > 0.3 * expected_size: + if verbose: + print( + f" NRRD axis {n_ax}: range [{lo_int}, {hi_int}] has " + f"{pb + pa}/{expected_size} voxels out of bounds " + f"[0, {nrrd_volume.shape[n_ax] - 1}] (>30%)" + ) + return None + + pad_before[n_ax] = pb + pad_after[n_ax] = pa + slices_per_axis[n_ax] = slice(clamped_lo, clamped_hi + 1) + + # Extract the sub-volume from NRRD + extracted = nrrd_volume[slices_per_axis[0], slices_per_axis[1], slices_per_axis[2]] + + # Zero-pad if any edges were clipped + if any(p > 0 for p in pad_before) or any(p > 0 for p in pad_after): + pad_widths = [(pad_before[ax], pad_after[ax]) for ax in range(3)] + extracted = np.pad(extracted, pad_widths, mode="constant", constant_values=0) + if verbose: + print(f" Padded {pad_widths} to compensate for edge clipping") + + if verbose: + print(f" Extracted NRRD region: {[str(s) for s in slices_per_axis]}, shape {extracted.shape}") + + # Permute axes: we need NRRD axes in the order [axis_map[0], axis_map[1], axis_map[2]] + # so that the result is (H, W, D) matching DICOM + perm = [axis_map[0], axis_map[1], axis_map[2]] + if perm != [0, 1, 2]: + extracted = np.transpose(extracted, perm) + + # Apply flips for negative axis signs + for d_ax in range(3): + if axis_sign[d_ax] < 0: + extracted = np.flip(extracted, axis=d_ax) + + # Verify shape + if extracted.shape != dicom_volume.shape: + if verbose: + print( + f" Shape mismatch after extraction: {extracted.shape} vs {dicom_volume.shape}" + ) + return None + + # Make contiguous copy (np.flip returns a view) + extracted = np.ascontiguousarray(extracted) + + if verbose: + n_labels = np.sum(extracted > 0) + print(f" Coordinate-based alignment successful, {n_labels} label voxels") + + return extracted + + +def _try_simple_alignment( + dicom_volume: np.ndarray, + nrrd_volume: np.ndarray, + verbose: bool = True, +) -> Optional[np.ndarray]: + """Try simple axis permutations to align volumes (shape-only fallback).""" + + dicom_shape = dicom_volume.shape + nrrd_shape = nrrd_volume.shape + + permutations = [ + (0, 1, 2), + (1, 0, 2), + (0, 2, 1), + (2, 1, 0), + (1, 2, 0), + (2, 0, 1), + ] + + for perm in permutations: + transformed = np.transpose(nrrd_volume, perm) + if transformed.shape == dicom_shape: + if verbose: + print(f" Fallback alignment: transpose{perm}") + return transformed + + return None + + +def get_labeled_slice_indices(segmentation: np.ndarray) -> List[int]: + """Get indices of slices containing labels.""" + labeled = [] + for i in range(segmentation.shape[-1]): + if np.any(segmentation[:, :, i] > 0): + labeled.append(i) + return labeled + + +def load_patient_data( + dicom_dir: str, + nrrd_path: str, + verbose: bool = True +) -> Tuple[np.ndarray, np.ndarray, Dict]: + """ + Load and align a patient's DICOM and NRRD data. + + Returns: + dicom_volume: 3D DICOM volume + aligned_segmentation: Aligned NRRD segmentation + metadata: Combined metadata + """ + # Load both + dicom_vol, dicom_meta = load_dicom_series(dicom_dir, verbose) + nrrd_vol, nrrd_meta = load_nrrd_segmentation(nrrd_path, verbose) + + # Align + aligned_seg, success = align_nrrd_to_dicom( + dicom_vol, nrrd_vol, nrrd_meta, verbose, dicom_metadata=dicom_meta + ) + + if not success: + warnings.warn(f"Alignment failed for {dicom_dir}") + + # Combined metadata + metadata = { + **dicom_meta, + "segments": nrrd_meta.get("segments", {}), + "alignment_success": success, + } + + return dicom_vol, aligned_seg, metadata + + +def discover_patients(base_dir: str) -> List[Dict]: + """ + Discover all patient folders in a dataset directory. + + Returns list of dicts with 'dicom_dir' and 'nrrd_path' keys. + """ + base = Path(base_dir) + patients = [] + + for patient_dir in sorted(base.iterdir()): + if not patient_dir.is_dir() or patient_dir.name.startswith('.'): + continue + + # Find NRRD file + nrrd_files = list(patient_dir.glob("*.nrrd")) + if not nrrd_files: + continue + + # Find DICOM subdirectory — pick the one with the most .dcm files + # (e.g. patient 001 has NL001/ with DICOMs and NL001_previews/ without) + dicom_dirs = [d for d in patient_dir.iterdir() if d.is_dir()] + if not dicom_dirs: + continue + + best_dicom_dir = None + best_dcm_count = 0 + for d in dicom_dirs: + n = len(list(d.glob("*.dcm")) + list(d.glob("*.DCM"))) + if n > best_dcm_count: + best_dcm_count = n + best_dicom_dir = d + + if best_dicom_dir is None or best_dcm_count == 0: + continue + dicom_dir = best_dicom_dir + dcm_files = list(dicom_dir.glob("*.dcm")) + list(dicom_dir.glob("*.DCM")) + + patients.append({ + "patient_id": patient_dir.name, + "dicom_dir": str(dicom_dir), + "nrrd_path": str(nrrd_files[0]), + }) + + return patients + + +if __name__ == "__main__": + # Test with sample patient + import sys + + if len(sys.argv) >= 3: + dicom_dir = sys.argv[1] + nrrd_path = sys.argv[2] + + dicom, seg, meta = load_patient_data(dicom_dir, nrrd_path) + + print(f"\n=== Results ===") + print(f"DICOM shape: {dicom.shape}") + print(f"Segmentation shape: {seg.shape}") + print(f"Alignment success: {meta['alignment_success']}") + print(f"Labeled slices: {len(get_labeled_slice_indices(seg))}") + else: + print("Usage: python data_utils.py ") + + diff --git a/try2/experiments/exp_classification.py b/try2/experiments/exp_classification.py new file mode 100644 index 0000000..a9e3f8d --- /dev/null +++ b/try2/experiments/exp_classification.py @@ -0,0 +1,349 @@ +""" +Train and evaluate AEA classification from segmentation masks. + +Usage: + # Train + LOO-CV on CROP1, then optionally validate on CROP EXT: + python exp_classification.py [--data-dir DATA_DIR] [--xlsx PATH] [--output-dir DIR] + [--ext-dir EXT_DIR] + +CROP1 folder format : "016. VA016 VINTELER ANA-MARIA" → code extracted as "VA016" +CROP EXT structure : + EXT_DIR/ + 01/ + EXT01/ ← DICOM directory + EXT01.nrrd ← segmentation mask + 02/ + EXT02/ + EXT02.nrrd + ... + +EXT patients are matched to the Excel by their EXTxx code. If the Excel has no +EXT codes, predictions are still written to output_dir/ext_predictions.csv. +""" + +import sys +import os +import argparse +import re +import csv + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..')) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +import numpy as np +from tqdm import tqdm +from sklearn.model_selection import LeaveOneOut +from sklearn.metrics import ( + accuracy_score, roc_auc_score, classification_report, balanced_accuracy_score +) + +from data_utils import load_patient_data, discover_patients, get_labeled_slice_indices +from shared.classification import ( + load_classification_labels, + build_feature_matrix, + train_classifiers, + save_classifiers, + predict_classifications, + TASK_TO_FEATURES, + _make_clf, + extract_features, +) +from shared.config import ExperimentConfig + + +XLSX_PATH = os.path.join( + os.path.dirname(__file__), '..', '..', 'AEA_classification_sinuzite.xlsx' +) + + +# --------------------------------------------------------------------------- +# Patient code extraction +# --------------------------------------------------------------------------- + +def _extract_patient_code(dicom_dir: str) -> str: + """Folder format: "016. VA016 VINTELER ANA-MARIA" — search anywhere in name.""" + basename = os.path.basename(os.path.normpath(dicom_dir)) + m = re.search(r'([A-Z]{2}\d{3})', basename.upper()) + if m: + return m.group(1) + return basename.upper().strip() + + +# --------------------------------------------------------------------------- +# CROP1 data loading +# --------------------------------------------------------------------------- + +def load_data(data_dir: str): + patients_meta = discover_patients(data_dir) + volumes, masks, codes = [], [], [] + for p in tqdm(patients_meta, desc="Loading CROP1"): + try: + vol, seg, meta = load_patient_data( + p['dicom_dir'], p['nrrd_path'], verbose=False + ) + except Exception as e: + print(f" [skip] {p.get('dicom_dir','?')}: {e}") + continue + if not meta.get('alignment_success', False): + continue + if len(get_labeled_slice_indices(seg)) < 2: + continue + volumes.append(vol) + masks.append(seg) + codes.append(_extract_patient_code(p['dicom_dir'])) + return volumes, masks, codes + + +# --------------------------------------------------------------------------- +# CROP EXT data loading +# --------------------------------------------------------------------------- + +def discover_ext_patients(ext_dir: str) -> list: + """ + Walk CROP EXT structure: + ext_dir/01/EXT01/ (DICOM) + ext_dir/01/EXT01.nrrd + ext_dir/02/EXT02/ + ext_dir/02/EXT02.nrrd + ... + Returns list of dicts: {code, dicom_dir, nrrd_path} + """ + patients = [] + for index_folder in sorted(os.listdir(ext_dir)): + index_path = os.path.join(ext_dir, index_folder) + if not os.path.isdir(index_path): + continue + + # Find the EXTxx subfolder and matching nrrd + ext_code = None + dicom_dir = None + nrrd_path = None + + for entry in os.listdir(index_path): + entry_path = os.path.join(index_path, entry) + if os.path.isdir(entry_path) and re.match(r'^EXT\d+$', entry.upper()): + dicom_dir = entry_path + ext_code = entry.upper() + elif entry.lower().endswith('.nrrd') and re.match(r'^EXT\d+', entry.upper()): + nrrd_path = entry_path + + if dicom_dir and nrrd_path and ext_code: + patients.append({ + 'code': ext_code, + 'dicom_dir': dicom_dir, + 'nrrd_path': nrrd_path, + }) + else: + print(f" [skip EXT] {index_path}: could not find EXTxx dir + nrrd") + + return patients + + +def load_ext_data(ext_dir: str): + patients_meta = discover_ext_patients(ext_dir) + volumes, masks, codes = [], [], [] + for p in tqdm(patients_meta, desc="Loading CROP EXT"): + try: + vol, seg, meta = load_patient_data( + p['dicom_dir'], p['nrrd_path'], verbose=False + ) + except Exception as e: + print(f" [skip] {p['code']}: {e}") + continue + if not meta.get('alignment_success', False): + continue + volumes.append(vol) + masks.append(seg) + codes.append(p['code']) + return volumes, masks, codes + + +# --------------------------------------------------------------------------- +# LOO cross-validation +# --------------------------------------------------------------------------- + +def loo_evaluate( + X: np.ndarray, y: np.ndarray, patient_idx: np.ndarray, task: str +) -> dict: + """Patient-level LOO-CV: hold out both sides of one patient per fold.""" + patients = np.unique(patient_idx) + if len(patients) < 4: + print(f" [skip] {task}: only {len(patients)} patients with matched labels") + return {} + + y_true, y_pred, y_prob = [], [], [] + + for p in patients: + test_mask = patient_idx == p + train_mask = ~test_mask + clf = _make_clf(task) + clf.fit(X[train_mask], y[train_mask]) + probs = clf.predict_proba(X[test_mask])[:, 1] + preds = clf.predict(X[test_mask]) + y_true.extend(y[test_mask].tolist()) + y_pred.extend(preds.tolist()) + y_prob.extend(probs.tolist()) + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_prob = np.array(y_prob) + + auc = roc_auc_score(y_true, y_prob) if len(np.unique(y_true)) > 1 else float('nan') + return { + "n": len(y_true), + "accuracy": accuracy_score(y_true, y_pred), + "balanced_accuracy": balanced_accuracy_score(y_true, y_pred), + "auc": auc, + "report": classification_report(y_true, y_pred, zero_division=0), + } + + +# --------------------------------------------------------------------------- +# External validation +# --------------------------------------------------------------------------- + +def run_ext_validation( + ext_dir: str, + clfs: dict, + labels: dict, + output_dir: str, +) -> None: + print(f"\n{'=' * 60}") + print("External Validation (CROP EXT)") + print(f"{'=' * 60}") + + volumes, masks, codes = load_ext_data(ext_dir) + print(f"Loaded {len(volumes)} EXT patients") + + if len(volumes) == 0: + print("No EXT patients loaded — check ext_dir structure.") + return + + # Run predictions — each patient produces left + right rows + side_labels = { + "left": ("aeal_roof_contact", "left_ethmoid_filled"), + "right": ("aear_roof_contact", "right_ethmoid_filled"), + } + rows = [] + for code, vol, msk in zip(codes, volumes, masks): + preds = predict_classifications(vol, msk, clfs) + for side, (roof_lbl, eth_lbl) in side_labels.items(): + row = {"patient_code": code, "side": side} + row["roof_contact"] = preds[side].get("roof_contact", "") + row["ethmoid_filled"] = preds[side].get("ethmoid_filled", "") + if code in labels: + row["roof_contact_gt"] = labels[code][roof_lbl] + row["ethmoid_filled_gt"] = labels[code][eth_lbl] + rows.append(row) + + # Save predictions CSV + os.makedirs(output_dir, exist_ok=True) + csv_path = os.path.join(output_dir, "ext_predictions.csv") + fieldnames = ["patient_code", "side", "roof_contact", "ethmoid_filled", + "roof_contact_gt", "ethmoid_filled_gt"] + with open(csv_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore") + writer.writeheader() + writer.writerows(rows) + print(f"Predictions saved → {csv_path}") + + # Compute metrics for patients that have Excel labels + labeled_rows = [r for r in rows if "roof_contact_gt" in r] + if not labeled_rows: + print("No EXT patients found in Excel — predictions saved, no metrics computed.") + return + + print(f"\nMetrics on {len(labeled_rows)} EXT instances with Excel labels:") + for task in ("roof_contact", "ethmoid_filled"): + y_true = np.array([r[f"{task}_gt"] for r in labeled_rows]) + y_pred = np.array([r[task] for r in labeled_rows]) + if len(np.unique(y_true)) < 2: + print(f"\n--- {task}: only one class in EXT labels, skipping AUC ---") + continue + print(f"\n--- {task} ---") + print(f" Balanced accuracy: {balanced_accuracy_score(y_true, y_pred):.3f}") + print(f" {classification_report(y_true, y_pred, zero_division=0)}") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--data-dir", default=ExperimentConfig.DATA_DIR) + parser.add_argument("--xlsx", default=XLSX_PATH) + parser.add_argument("--output-dir", default=None) + parser.add_argument("--ext-dir", default=None, + help="Path to CROP EXT folder for external validation") + args = parser.parse_args() + + output_dir = args.output_dir or os.path.join( + os.path.dirname(__file__), '..', 'outputs', 'classification' + ) + + print("=== AEA Classification Experiment ===\n") + print(f"Data dir : {args.data_dir}") + print(f"XLSX : {args.xlsx}") + print(f"Output : {output_dir}") + if args.ext_dir: + print(f"EXT dir : {args.ext_dir}") + print() + + labels = load_classification_labels(args.xlsx) + print(f"Loaded {len(labels)} patient labels from XLSX") + + print("\nLoading imaging data...") + volumes, masks, codes = load_data(args.data_dir) + print(f"Loaded {len(volumes)} valid CROP1 patients") + + matched = [c for c in codes if c in labels] + print(f"Matched to XLSX labels: {len(matched)}/{len(codes)}") + unmatched = [c for c in codes if c not in labels] + if unmatched: + print(f" Unmatched codes: {unmatched[:10]}{'...' if len(unmatched) > 10 else ''}") + print(" → Adjust _extract_patient_code() if codes look wrong") + + if len(matched) == 0: + print("\nERROR: No patients matched. Check patient code extraction.") + return + + X_dict, y_dict, patient_idx = build_feature_matrix(codes, volumes, masks, labels) + n_patients = len(np.unique(patient_idx)) + n_samples = len(y_dict["roof_contact"]) + print(f"\nFeature matrix built: {n_patients} patients → {n_samples} samples (L+R pooled)") + + # LOO cross-validation (patient-level) + print(f"\n{'=' * 60}") + print("Patient-level Leave-One-Out Cross-Validation (CROP1)") + print(f"{'=' * 60}") + + for task, feat_key in TASK_TO_FEATURES.items(): + X = X_dict[feat_key] + y = y_dict[task] + print(f"\n--- {task} ---") + print(f" Features : {feat_key} shape={X.shape}") + print(f" Positive : {int(y.sum())}/{len(y)} ({100*y.mean():.1f}%)") + res = loo_evaluate(X, y, patient_idx, task) + if res: + print(f" Accuracy : {res['accuracy']:.3f}") + print(f" Balanced accuracy : {res['balanced_accuracy']:.3f}") + print(f" AUC : {res['auc']:.3f}") + print(f" {res['report']}") + + # Train final classifiers on all CROP1 data + print(f"\n{'=' * 60}") + print("Training final classifiers on all CROP1 data...") + clfs = train_classifiers(X_dict, y_dict) + save_classifiers(clfs, output_dir) + print(f"Saved {len(clfs)} classifiers to {output_dir}/") + for task in clfs: + print(f" → {task}.joblib") + + # External validation + if args.ext_dir: + run_ext_validation(args.ext_dir, clfs, labels, output_dir) + + print("\nDone.") + + +if __name__ == "__main__": + main() diff --git a/try2/experiments/exp_ethmoid_analysis.py b/try2/experiments/exp_ethmoid_analysis.py new file mode 100644 index 0000000..629e5e0 --- /dev/null +++ b/try2/experiments/exp_ethmoid_analysis.py @@ -0,0 +1,403 @@ +""" +Systematic analysis of ethmoid sinus filling prediction. +Tests HU thresholds, slab widths, distribution features, sagittal position, +classifiers, and combinations via LOO-CV. +""" + +import sys +import os + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..')) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +import numpy as np +from tqdm import tqdm +from scipy.stats import kurtosis as scipy_kurtosis +from sklearn.model_selection import LeaveOneOut +from sklearn.metrics import roc_auc_score, balanced_accuracy_score +from sklearn.linear_model import LogisticRegression +from sklearn.ensemble import RandomForestClassifier +from sklearn.preprocessing import StandardScaler +from sklearn.pipeline import Pipeline + +from exp_classification import load_data +from shared.classification import load_classification_labels +from shared.config import ExperimentConfig + +XLSX_PATH = os.path.join(os.path.dirname(__file__), '..', '..', 'AEA_classification_sinuzite.xlsx') +DATA_DIR = ExperimentConfig.DATA_DIR + + +# --------------------------------------------------------------------------- +# Feature extractors +# --------------------------------------------------------------------------- + +def get_mask_voxels(mask, class_id): + return np.argwhere(mask == class_id) + + +def get_slab_intensities(volume, mask, class_id, w_lo, w_hi): + voxels = get_mask_voxels(mask, class_id) + if len(voxels) == 0: + return None + sel = (voxels[:, 1] >= w_lo) & (voxels[:, 1] <= w_hi) + sv = voxels[sel] + if len(sv) == 0: + return None + return volume[sv[:, 0], sv[:, 1], sv[:, 2]].astype(np.float32) + + +def get_slab_bounds(voxels, volume_shape, half_width, position='center'): + """Return (w_lo, w_hi) for a slab. + position: 'center', 'anterior', 'middle', 'posterior', or 'full' + half_width: ±half_width voxels (ignored if position='full') + """ + W = volume_shape[1] + if position == 'full': + return 0, W - 1 + + w_vals = voxels[:, 1] + w_min, w_max = int(w_vals.min()), int(w_vals.max()) + w_range = w_max - w_min + third = max(1, w_range // 3) + + if position == 'anterior': + # anterior = low W index (remember H=A-P, W=L-R ... actually W is sagittal/L-R) + # "anterior" in the A-P sense is H, not W — but task says "anterior third of W range" + # so we interpret as the first third of the mask's W extent + anchor = w_min + third // 2 + elif position == 'middle': + anchor = w_min + third + third // 2 + elif position == 'posterior': + anchor = w_min + 2 * third + third // 2 + else: # center = W centroid + anchor = int(round(float(w_vals.mean()))) + + if half_width == 0: + return anchor, anchor + return max(0, anchor - half_width), min(W - 1, anchor + half_width) + + +def base_features(intensities, threshold=-500): + """6 features matching original: mean, median, std, p10, p90, filled_ratio.""" + if intensities is None or len(intensities) == 0: + return np.zeros(6, dtype=np.float32) + return np.array([ + float(np.mean(intensities)), + float(np.median(intensities)), + float(np.std(intensities)), + float(np.percentile(intensities, 10)), + float(np.percentile(intensities, 90)), + float(np.mean(intensities > threshold)), + ], dtype=np.float32) + + +def distribution_features(intensities): + """Extra distribution shape features: kurtosis, fluid_frac, soft_frac, entropy.""" + if intensities is None or len(intensities) == 0: + return np.zeros(4, dtype=np.float32) + kurt = float(scipy_kurtosis(intensities, fisher=True)) + fluid_frac = float(np.mean((intensities >= 0) & (intensities <= 100))) + soft_frac = float(np.mean((intensities >= -100) & (intensities <= 100))) + # entropy of 20-bin histogram from -1000 to 500 HU + counts, _ = np.histogram(intensities, bins=20, range=(-1000, 500)) + counts = counts.astype(np.float64) + counts += 1e-10 # avoid log(0) + probs = counts / counts.sum() + entropy = float(-np.sum(probs * np.log(probs))) + return np.array([kurt, fluid_frac, soft_frac, entropy], dtype=np.float32) + + +# --------------------------------------------------------------------------- +# LOO-CV runner +# --------------------------------------------------------------------------- + +def make_lr(): + return Pipeline([ + ('scaler', StandardScaler()), + ('clf', LogisticRegression(class_weight='balanced', max_iter=1000, C=1.0)), + ]) + + +def make_rf(): + return Pipeline([ + ('scaler', StandardScaler()), + ('clf', RandomForestClassifier(n_estimators=100, class_weight='balanced', random_state=42)), + ]) + + +def loo_cv(X, y, clf_factory=make_lr): + loo = LeaveOneOut() + y_true, y_prob = [], [] + for train_idx, test_idx in loo.split(X): + clf = clf_factory() + clf.fit(X[train_idx], y[train_idx]) + y_true.append(y[test_idx[0]]) + y_prob.append(clf.predict_proba(X[test_idx])[0, 1]) + y_true = np.array(y_true) + y_prob = np.array(y_prob) + y_pred = (y_prob >= 0.5).astype(int) + auc = roc_auc_score(y_true, y_prob) if len(np.unique(y_true)) > 1 else float('nan') + bal_acc = balanced_accuracy_score(y_true, y_pred) + return auc, bal_acc + + +# --------------------------------------------------------------------------- +# Build feature matrices for all experiments +# --------------------------------------------------------------------------- + +def build_features_for_experiment(volumes, masks, codes, labels, class_id, + half_width, position, threshold, + use_dist_features=False): + X_list, y_list = [], [] + for vol, msk, code in zip(volumes, masks, codes): + if code not in labels: + continue + voxels = get_mask_voxels(msk, class_id) + if len(voxels) == 0: + # still need to append a zero row to keep alignment - but skip patient + continue + w_lo, w_hi = get_slab_bounds(voxels, vol.shape, half_width, position) + intensities = get_slab_intensities(vol, msk, class_id, w_lo, w_hi) + feats = base_features(intensities, threshold) + if use_dist_features: + feats = np.concatenate([feats, distribution_features(intensities)]) + X_list.append(feats) + y_list.append(labels[code]) + if not X_list: + return np.empty((0, 6)), np.empty(0, dtype=int) + return np.stack(X_list), np.array(y_list, dtype=int) + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + print("=== Ethmoid Sinus Filling — Systematic Analysis ===\n") + print(f"Data dir : {DATA_DIR}") + print(f"XLSX : {XLSX_PATH}\n") + + # Load data + labels_all = load_classification_labels(XLSX_PATH) + + print("Loading imaging data (this may take a while)...") + volumes, masks, codes = load_data(DATA_DIR) + print(f"Loaded {len(volumes)} patients\n") + + results = [] # (experiment_name, side, auc, bal_acc, n_pos, n_total) + + sides = [ + ('Left', 1, 'left_ethmoid_filled'), + ('Right', 2, 'right_ethmoid_filled'), + ] + + def run_exp(name, half_width, position, threshold, use_dist, clf_factory=make_lr): + for side_name, class_id, label_key in sides: + # build label dict for this task + task_labels = {code: labels_all[code][label_key] + for code in codes if code in labels_all} + X, y = build_features_for_experiment( + volumes, masks, codes, task_labels, class_id, + half_width, position, threshold, use_dist + ) + if len(X) < 4: + continue + n_pos = int(y.sum()) + n_total = len(y) + auc, bal_acc = loo_cv(X, y, clf_factory) + results.append({ + 'experiment': name, + 'side': side_name, + 'auc': auc, + 'bal_acc': bal_acc, + 'n_pos': n_pos, + 'n_total': n_total, + }) + print(f" {name} | {side_name:5s} | AUC={auc:.3f} | BalAcc={bal_acc:.3f} | pos={n_pos}/{n_total}") + + # ----------------------------------------------------------------------- + # BASELINE (reproduce original) + print("\n--- BASELINE (slab=±2, threshold=-500, LR) ---") + run_exp("Baseline(slab±2,thr=-500,LR)", half_width=2, position='center', + threshold=-500, use_dist=False, clf_factory=make_lr) + + # ----------------------------------------------------------------------- + # A. HU threshold sweep + print("\n--- A. HU Threshold Sweep (slab=±2, center) ---") + for thr in [-800, -500, -200, 0, 100]: + run_exp(f"A_thr={thr}", half_width=2, position='center', + threshold=thr, use_dist=False, clf_factory=make_lr) + + # ----------------------------------------------------------------------- + # B. Sagittal slab width sweep + print("\n--- B. Slab Width Sweep (threshold=-500, center) ---") + slab_configs = [ + ('±0(single)', 0), + ('±1', 1), + ('±2', 2), + ('±5', 5), + ('full3D', 999), # we'll handle full3D via position='full' + ] + for label, hw in slab_configs: + pos = 'full' if hw == 999 else 'center' + hw_eff = 0 if hw == 999 else hw + run_exp(f"B_slab={label}", half_width=hw_eff, position=pos, + threshold=-500, use_dist=False, clf_factory=make_lr) + + # ----------------------------------------------------------------------- + # C. Distribution shape features (slab=±2, threshold=-500) + print("\n--- C. Distribution Shape Features (slab=±2, thr=-500) ---") + run_exp("C_base_only", half_width=2, position='center', + threshold=-500, use_dist=False, clf_factory=make_lr) + run_exp("C_base+dist", half_width=2, position='center', + threshold=-500, use_dist=True, clf_factory=make_lr) + + # ----------------------------------------------------------------------- + # D. Sagittal position (slab=±2, threshold=-500) + # Also find which position gives highest filled_ratio variance + print("\n--- D. Sagittal Position (slab=±2, thr=-500) ---") + + # Compute variance of filled_ratio across patients for each position to pick best + positions_to_test = ['anterior', 'middle', 'posterior', 'center'] + best_pos = {} + for side_name, class_id, label_key in sides: + task_labels = {code: labels_all[code][label_key] + for code in codes if code in labels_all} + best_var = -1 + best_p = 'center' + for pos in positions_to_test: + ratios = [] + for vol, msk, code in zip(volumes, masks, codes): + if code not in task_labels: + continue + voxels = get_mask_voxels(msk, class_id) + if len(voxels) == 0: + continue + w_lo, w_hi = get_slab_bounds(voxels, vol.shape, 2, pos) + intensities = get_slab_intensities(vol, msk, class_id, w_lo, w_hi) + if intensities is not None and len(intensities) > 0: + ratios.append(float(np.mean(intensities > -500))) + if ratios: + v = float(np.var(ratios)) + if v > best_var: + best_var = v + best_p = pos + best_pos[side_name] = best_p + print(f" Best position for {side_name}: {best_p} (filled_ratio variance={best_var:.4f})") + + for pos in positions_to_test: + run_exp(f"D_pos={pos}", half_width=2, position=pos, + threshold=-500, use_dist=False, clf_factory=make_lr) + + # ----------------------------------------------------------------------- + # E. Classifier comparison (slab=±2, threshold=-500, center, no dist) + print("\n--- E. Classifier Comparison (slab=±2, thr=-500, center, no dist) ---") + run_exp("E_LR", half_width=2, position='center', + threshold=-500, use_dist=False, clf_factory=make_lr) + run_exp("E_RF", half_width=2, position='center', + threshold=-500, use_dist=False, clf_factory=make_rf) + + # ----------------------------------------------------------------------- + # F. Best combination + # From the results so far, pick best slab width, threshold, position + # Then test both classifiers with dist features + print("\n--- F. Best Combination ---") + + # Gather AUCs from experiments B (slab width) and A (threshold) for each side + # to determine best settings programmatically + def best_setting(prefix, param_key): + subset = [r for r in results if r['experiment'].startswith(prefix)] + if not subset: + return None + best = max(subset, key=lambda r: r['auc'] if not np.isnan(r['auc']) else -1) + return best['experiment'] + + # Pick best slab from B experiments (average over sides) + b_exps = {} + for r in results: + if r['experiment'].startswith('B_'): + key = r['experiment'] + b_exps.setdefault(key, []).append(r['auc']) + b_avg = {k: np.nanmean(v) for k, v in b_exps.items()} + best_slab_exp = max(b_avg, key=lambda k: b_avg[k]) if b_avg else 'B_slab=±2' + # extract half_width from best_slab_exp name + slab_map = {'±0(single)': (0, 'center'), '±1': (1, 'center'), '±2': (2, 'center'), + '±5': (5, 'center'), 'full3D': (0, 'full')} + best_slab_label = best_slab_exp.replace('B_slab=', '') + best_hw, best_pos_f = slab_map.get(best_slab_label, (2, 'center')) + print(f" Best slab from B: {best_slab_exp} (hw={best_hw}, pos={best_pos_f})") + + # Pick best threshold from A experiments + a_exps = {} + for r in results: + if r['experiment'].startswith('A_'): + key = r['experiment'] + a_exps.setdefault(key, []).append(r['auc']) + a_avg = {k: np.nanmean(v) for k, v in a_exps.items()} + best_thr_exp = max(a_avg, key=lambda k: a_avg[k]) if a_avg else 'A_thr=-500' + best_thr = int(best_thr_exp.replace('A_thr=', '')) + print(f" Best threshold from A: {best_thr_exp} (thr={best_thr})") + + # Best position from D + d_exps = {} + for r in results: + if r['experiment'].startswith('D_'): + key = r['experiment'] + d_exps.setdefault(key, []).append(r['auc']) + d_avg = {k: np.nanmean(v) for k, v in d_exps.items()} + best_pos_exp = max(d_avg, key=lambda k: d_avg[k]) if d_avg else 'D_pos=center' + best_pos_d = best_pos_exp.replace('D_pos=', '') + print(f" Best position from D: {best_pos_exp} (pos={best_pos_d})") + + # Combination experiments + for use_dist in [False, True]: + for clf_name, clf_f in [('LR', make_lr), ('RF', make_rf)]: + dist_tag = '+dist' if use_dist else '' + name = f"F_slab={best_slab_label},thr={best_thr},pos={best_pos_d}{dist_tag},{clf_name}" + run_exp(name, half_width=best_hw, position=best_pos_d, + threshold=best_thr, use_dist=use_dist, clf_factory=clf_f) + + # ----------------------------------------------------------------------- + # Summary table + print("\n" + "=" * 100) + print("SUMMARY TABLE — All Experiments Ranked by Mean AUC (averaged over Left+Right)") + print("=" * 100) + + # Aggregate per experiment (mean over sides) + exp_summary = {} + for r in results: + key = r['experiment'] + exp_summary.setdefault(key, []).append(r) + + rows_summary = [] + for exp_name, exp_results in exp_summary.items(): + aucs = [r['auc'] for r in exp_results] + bals = [r['bal_acc'] for r in exp_results] + mean_auc = float(np.nanmean(aucs)) + mean_bal = float(np.nanmean(bals)) + # individual side info + sides_info = {r['side']: r for r in exp_results} + rows_summary.append((mean_auc, mean_bal, exp_name, sides_info)) + + rows_summary.sort(key=lambda x: x[0], reverse=True) + + header = f"{'Rank':>4} {'Experiment':<55} {'MeanAUC':>7} {'MeanBalAcc':>10} {'Left_AUC':>8} {'Right_AUC':>9} {'pos/n':>7}" + print(header) + print("-" * len(header)) + for rank, (mean_auc, mean_bal, exp_name, sides_info) in enumerate(rows_summary, 1): + left_r = sides_info.get('Left', {}) + right_r = sides_info.get('Right', {}) + left_auc = f"{left_r.get('auc', float('nan')):.3f}" if left_r else ' N/A ' + right_auc = f"{right_r.get('auc', float('nan')):.3f}" if right_r else ' N/A ' + # use left side for pos/n (representative) + n_pos = left_r.get('n_pos', '?') + n_total = left_r.get('n_total', '?') + pos_n = f"{n_pos}/{n_total}" + print(f"{rank:>4} {exp_name:<55} {mean_auc:>7.3f} {mean_bal:>10.3f} {left_auc:>8} {right_auc:>9} {pos_n:>7}") + + print("\nDone.") + + +if __name__ == '__main__': + main() diff --git a/try2/experiments/exp_ethmoid_round2.py b/try2/experiments/exp_ethmoid_round2.py new file mode 100644 index 0000000..3a08ca1 --- /dev/null +++ b/try2/experiments/exp_ethmoid_round2.py @@ -0,0 +1,361 @@ +""" +Round 2 ethmoid sinus filling prediction experiments. + +Base config: middle-W slab ±1, threshold -200 HU, LogisticRegression (balanced). + +Experiments: + G. I-S stratification — superior vs inferior half of D axis + H. Texture / distribution shape — kurtosis, skew, fluid/mucus fractions, bimodality + I. HU histogram as features — 10-bin histogram + PCA(3) before LogReg + J. Sinus air volume proxy — fractions below -800, -500, -200 HU + K. Left-right symmetry — joint classifier with cross-side symmetry feature + L. Combination — best features from G+H+J together +""" + +import sys +import os + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..')) +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +import numpy as np +from scipy import stats as scipy_stats +from sklearn.model_selection import LeaveOneOut +from sklearn.metrics import roc_auc_score, balanced_accuracy_score +from sklearn.linear_model import LogisticRegression +from sklearn.preprocessing import StandardScaler +from sklearn.pipeline import Pipeline +from sklearn.decomposition import PCA + +from exp_classification import load_data +from shared.classification import load_classification_labels +from shared.config import ExperimentConfig + +XLSX_PATH = os.path.join(os.path.dirname(__file__), '..', '..', 'AEA_classification_sinuzite.xlsx') +DATA_DIR = ExperimentConfig.DATA_DIR + +# Base config constants +BASE_HALF_WIDTH = 1 +BASE_THRESHOLD = -200 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def get_slab_voxels_and_intensities(volume, mask, class_id): + """Return (slab_voxels, intensities) using base config: middle-W slab ±1.""" + voxels = np.argwhere(mask == class_id) + if len(voxels) == 0: + return None, None + w_min, w_max = int(voxels[:, 1].min()), int(voxels[:, 1].max()) + w_mid = (w_min + w_max) // 2 + w_lo = max(0, w_mid - BASE_HALF_WIDTH) + w_hi = min(volume.shape[1] - 1, w_mid + BASE_HALF_WIDTH) + sel = (voxels[:, 1] >= w_lo) & (voxels[:, 1] <= w_hi) + sv = voxels[sel] + if len(sv) == 0: + return None, None + intensities = volume[sv[:, 0], sv[:, 1], sv[:, 2]].astype(np.float32) + return sv, intensities + + +def make_lr(): + return Pipeline([ + ('scaler', StandardScaler()), + ('clf', LogisticRegression(class_weight='balanced', max_iter=1000, C=1.0)), + ]) + + +def loo_cv(X, y, clf_factory=make_lr): + loo = LeaveOneOut() + y_true, y_prob = [], [] + for train_idx, test_idx in loo.split(X): + clf = clf_factory() + clf.fit(X[train_idx], y[train_idx]) + y_true.append(y[test_idx[0]]) + y_prob.append(clf.predict_proba(X[test_idx])[0, 1]) + y_true = np.array(y_true) + y_prob = np.array(y_prob) + y_pred = (y_prob >= 0.5).astype(int) + auc = roc_auc_score(y_true, y_prob) if len(np.unique(y_true)) > 1 else float('nan') + bal_acc = balanced_accuracy_score(y_true, y_pred) + return auc, bal_acc + + +# --------------------------------------------------------------------------- +# Feature extractors +# --------------------------------------------------------------------------- + +def features_G(volume, mask, class_id): + """G. I-S stratification: superior vs inferior half of D axis + gradient.""" + sv, intensities = get_slab_voxels_and_intensities(volume, mask, class_id) + if sv is None: + return np.zeros(5, dtype=np.float32) + d_coords = sv[:, 2] + d_min, d_max = int(d_coords.min()), int(d_coords.max()) + d_mid = (d_min + d_max) // 2 + inf_mask = d_coords <= d_mid + sup_mask = d_coords > d_mid + inf_ints = intensities[inf_mask] + sup_ints = intensities[sup_mask] + filled_ratio_inf = float(np.mean(inf_ints > BASE_THRESHOLD)) if len(inf_ints) > 0 else 0.0 + filled_ratio_sup = float(np.mean(sup_ints > BASE_THRESHOLD)) if len(sup_ints) > 0 else 0.0 + filled_ratio_all = float(np.mean(intensities > BASE_THRESHOLD)) + gradient = filled_ratio_inf - filled_ratio_sup # positive = more filling inferiorly + # fraction of voxels in inferior half + inf_frac = float(len(inf_ints)) / max(1, len(intensities)) + return np.array([filled_ratio_inf, filled_ratio_sup, filled_ratio_all, gradient, inf_frac], + dtype=np.float32) + + +def features_H(volume, mask, class_id): + """H. Texture / distribution shape features.""" + _, intensities = get_slab_voxels_and_intensities(volume, mask, class_id) + if intensities is None or len(intensities) < 3: + return np.zeros(8, dtype=np.float32) + mean_hu = float(np.mean(intensities)) + std_hu = float(np.std(intensities)) + filled_ratio = float(np.mean(intensities > BASE_THRESHOLD)) + skew = float(scipy_stats.skew(intensities)) + kurt = float(scipy_stats.kurtosis(intensities, fisher=True)) + fluid_frac = float(np.mean((intensities >= 0) & (intensities <= 100))) + mucus_frac = float(np.mean((intensities >= -100) & (intensities < 0))) + # bimodality coefficient + n = len(intensities) + if kurt > -1 and n > 3: + bimodality = (skew ** 2 + 1) / (kurt + 3 * (n - 1) ** 2 / ((n - 2) * (n - 3) + 1e-9)) + else: + bimodality = 0.0 + return np.array([mean_hu, std_hu, filled_ratio, skew, kurt, fluid_frac, mucus_frac, bimodality], + dtype=np.float32) + + +def features_I_raw(volume, mask, class_id): + """I. 10-bin HU histogram features (raw, before PCA).""" + _, intensities = get_slab_voxels_and_intensities(volume, mask, class_id) + if intensities is None or len(intensities) == 0: + return np.zeros(10, dtype=np.float32) + counts, _ = np.histogram(intensities, bins=10, range=(-1000, 500)) + total = max(1, counts.sum()) + return (counts / total).astype(np.float32) + + +def features_J(volume, mask, class_id): + """J. Sinus air volume proxy: fractions below -800, -500, -200 HU.""" + _, intensities = get_slab_voxels_and_intensities(volume, mask, class_id) + if intensities is None or len(intensities) == 0: + return np.zeros(3, dtype=np.float32) + frac_air = float(np.mean(intensities < -800)) + frac_below500 = float(np.mean(intensities < -500)) + frac_below200 = float(np.mean(intensities < -200)) + return np.array([frac_air, frac_below500, frac_below200], dtype=np.float32) + + +def features_GHJ(volume, mask, class_id): + """L. Combined G + H + J features.""" + g = features_G(volume, mask, class_id) + h = features_H(volume, mask, class_id) + j = features_J(volume, mask, class_id) + return np.concatenate([g, h, j]) + + +# --------------------------------------------------------------------------- +# Build feature matrices +# --------------------------------------------------------------------------- + +def build_X_y(volumes, masks, codes, labels_all, class_id, label_key, feat_fn): + X_list, y_list = [], [] + for vol, msk, code in zip(volumes, masks, codes): + if code not in labels_all: + continue + feats = feat_fn(vol, msk, class_id) + X_list.append(feats) + y_list.append(labels_all[code][label_key]) + if not X_list: + return np.empty((0, 1)), np.empty(0, dtype=int) + return np.stack(X_list), np.array(y_list, dtype=int) + + +def build_X_y_joint(volumes, masks, codes, labels_all): + """K. Build joint feature matrix with both sides + symmetry features.""" + X_list, y_L_list, y_R_list = [], [], [] + for vol, msk, code in zip(volumes, masks, codes): + if code not in labels_all: + continue + # Base features for both sides + _, ints_L = get_slab_voxels_and_intensities(vol, msk, class_id=1) + _, ints_R = get_slab_voxels_and_intensities(vol, msk, class_id=2) + + fr_L = float(np.mean(ints_L > BASE_THRESHOLD)) if ints_L is not None and len(ints_L) > 0 else 0.0 + fr_R = float(np.mean(ints_R > BASE_THRESHOLD)) if ints_R is not None and len(ints_R) > 0 else 0.0 + + mean_L = float(np.mean(ints_L)) if ints_L is not None and len(ints_L) > 0 else 0.0 + mean_R = float(np.mean(ints_R)) if ints_R is not None and len(ints_R) > 0 else 0.0 + + # Symmetry features + ratio = fr_L / (fr_R + 1e-6) + diff = fr_L - fr_R + mean_diff = mean_L - mean_R + + feats = np.array([fr_L, fr_R, mean_L, mean_R, ratio, diff, mean_diff], dtype=np.float32) + X_list.append(feats) + y_L_list.append(labels_all[code]['left_ethmoid_filled']) + y_R_list.append(labels_all[code]['right_ethmoid_filled']) + + if not X_list: + return np.empty((0, 7)), np.empty(0, dtype=int), np.empty(0, dtype=int) + return np.stack(X_list), np.array(y_L_list, dtype=int), np.array(y_R_list, dtype=int) + + +# --------------------------------------------------------------------------- +# PCA-LR factory for experiment I +# --------------------------------------------------------------------------- + +def make_pca_lr(n_components=3): + def factory(): + return Pipeline([ + ('scaler', StandardScaler()), + ('pca', PCA(n_components=n_components)), + ('clf', LogisticRegression(class_weight='balanced', max_iter=1000, C=1.0)), + ]) + return factory + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + print("=== Ethmoid Sinus Filling — Round 2 Experiments ===\n") + print(f"Base config: middle-W slab ±{BASE_HALF_WIDTH}, threshold {BASE_THRESHOLD} HU, LogReg balanced\n") + print(f"Data dir : {DATA_DIR}") + print(f"XLSX : {XLSX_PATH}\n") + + labels_all = load_classification_labels(XLSX_PATH) + + print("Loading imaging data...") + volumes, masks, codes = load_data(DATA_DIR) + print(f"Loaded {len(volumes)} patients\n") + + results = [] + + def report(name, side, auc, bal_acc, n_pos, n_total): + results.append({'experiment': name, 'side': side, 'auc': auc, + 'bal_acc': bal_acc, 'n_pos': n_pos, 'n_total': n_total}) + tag = f" {name} | {side:5s} | AUC={auc:.3f} | BalAcc={bal_acc:.3f} | pos={n_pos}/{n_total}" + print(tag) + + def run_both_sides(name, feat_fn, clf_factory=make_lr): + for side_name, class_id, label_key in [ + ('Left', 1, 'left_ethmoid_filled'), + ('Right', 2, 'right_ethmoid_filled'), + ]: + X, y = build_X_y(volumes, masks, codes, labels_all, class_id, label_key, feat_fn) + if len(X) < 4: + print(f" {name} | {side_name} | SKIPPED (n={len(X)})") + continue + n_pos = int(y.sum()) + n_total = len(y) + auc, bal_acc = loo_cv(X, y, clf_factory) + report(name, side_name, auc, bal_acc, n_pos, n_total) + + # ----------------------------------------------------------------------- + # G. I-S stratification + print("\n--- G. I-S Stratification (superior vs inferior half of D axis) ---") + run_both_sides("G_IS_stratification", features_G) + + # ----------------------------------------------------------------------- + # H. Texture / distribution shape + print("\n--- H. Texture / Distribution Shape Features ---") + run_both_sides("H_texture_dist", features_H) + + # ----------------------------------------------------------------------- + # I. HU histogram + PCA + print("\n--- I. HU Histogram (10 bins) + PCA(3) + LogReg ---") + # Check n_components doesn't exceed n_features or n_samples-1 + # We have 10 features; use PCA(3) + run_both_sides("I_histogram_PCA3", features_I_raw, clf_factory=make_pca_lr(n_components=3)) + + # ----------------------------------------------------------------------- + # J. Sinus air volume proxy + print("\n--- J. Sinus Air Volume Proxy (fractions <-800, <-500, <-200 HU) ---") + run_both_sides("J_air_proxy", features_J) + + # ----------------------------------------------------------------------- + # K. Left-right symmetry (joint classifier) + print("\n--- K. Left-Right Symmetry (joint features, separate LOO-CV per side) ---") + X_joint, y_L, y_R = build_X_y_joint(volumes, masks, codes, labels_all) + if len(X_joint) >= 4: + auc_L, bal_L = loo_cv(X_joint, y_L) + report("K_LR_symmetry_joint", 'Left', auc_L, bal_L, int(y_L.sum()), len(y_L)) + auc_R, bal_R = loo_cv(X_joint, y_R) + report("K_LR_symmetry_joint", 'Right', auc_R, bal_R, int(y_R.sum()), len(y_R)) + else: + print(f" K_LR_symmetry_joint | SKIPPED (n={len(X_joint)})") + + # ----------------------------------------------------------------------- + # L. Combination G + H + J + print("\n--- L. Combination: G + H + J features ---") + run_both_sides("L_GHJ_combined", features_GHJ) + + # ----------------------------------------------------------------------- + # Summary + print("\n" + "=" * 110) + print("ROUND 2 SUMMARY — Ranked by Mean AUC (Left+Right average)") + print("=" * 110) + + exp_summary = {} + for r in results: + exp_summary.setdefault(r['experiment'], []).append(r) + + rows = [] + for exp_name, exp_results in exp_summary.items(): + aucs = [r['auc'] for r in exp_results] + bals = [r['bal_acc'] for r in exp_results] + mean_auc = float(np.nanmean(aucs)) + mean_bal = float(np.nanmean(bals)) + sides_info = {r['side']: r for r in exp_results} + rows.append((mean_auc, mean_bal, exp_name, sides_info)) + + rows.sort(key=lambda x: x[0], reverse=True) + + ROUND1_BASELINE = {'Left': 0.508, 'Right': 0.688, 'Mean': 0.598} + + header = (f"{'Rank':>4} {'Experiment':<30} {'MeanAUC':>7} {'BalAcc':>7}" + f" {'Left_AUC':>8} {'Right_AUC':>9} {'Left_vs_R1':>10} {'n':>6}") + print(header) + print("-" * len(header)) + + for rank, (mean_auc, mean_bal, exp_name, sides_info) in enumerate(rows, 1): + left_r = sides_info.get('Left', {}) + right_r = sides_info.get('Right', {}) + left_auc = left_r.get('auc', float('nan')) + right_auc = right_r.get('auc', float('nan')) + left_diff = left_auc - ROUND1_BASELINE['Left'] + diff_tag = f"{left_diff:+.3f}" + n_pos = left_r.get('n_pos', '?') + n_total = left_r.get('n_total', '?') + print(f"{rank:>4} {exp_name:<30} {mean_auc:>7.3f} {mean_bal:>7.3f}" + f" {left_auc:>8.3f} {right_auc:>9.3f} {diff_tag:>10} {n_pos}/{n_total}") + + print() + print("Round 1 best: Mean=0.598, Left=0.508, Right=0.688 (RF, slab±1, thr=-200, middle-W)") + print() + + # Highlight experiments that improved Left AUC + improved_left = [(mean_auc, exp_name, sides_info['Left'].get('auc', float('nan'))) + for mean_auc, _, exp_name, sides_info in rows + if 'Left' in sides_info and sides_info['Left'].get('auc', 0) > ROUND1_BASELINE['Left']] + if improved_left: + print("Experiments that IMPROVED Left AUC vs Round 1 (baseline=0.508):") + for mean_auc, exp_name, left_auc in sorted(improved_left, key=lambda x: -x[2]): + print(f" {exp_name:<30} Left AUC={left_auc:.3f} (+{left_auc - ROUND1_BASELINE['Left']:.3f})") + else: + print("No experiment exceeded the Round 1 Left AUC baseline of 0.508.") + + print("\nDone.") + + +if __name__ == '__main__': + main() diff --git a/try2/outputs/classification/aeal_roof_contact.joblib b/try2/outputs/classification/aeal_roof_contact.joblib new file mode 100644 index 0000000..3bb2772 Binary files /dev/null and b/try2/outputs/classification/aeal_roof_contact.joblib differ diff --git a/try2/outputs/classification/aear_roof_contact.joblib b/try2/outputs/classification/aear_roof_contact.joblib new file mode 100644 index 0000000..b039be4 Binary files /dev/null and b/try2/outputs/classification/aear_roof_contact.joblib differ diff --git a/try2/outputs/classification/ext_predictions.csv b/try2/outputs/classification/ext_predictions.csv new file mode 100644 index 0000000..321b3d5 --- /dev/null +++ b/try2/outputs/classification/ext_predictions.csv @@ -0,0 +1,39 @@ +patient_code,side,roof_contact,ethmoid_filled,roof_contact_gt,ethmoid_filled_gt +EXT01,left,1,1,, +EXT01,right,0,1,, +EXT03,left,0,0,, +EXT03,right,0,1,, +EXT04,left,0,1,, +EXT04,right,0,1,, +EXT05,left,0,0,, +EXT05,right,0,1,, +EXT06,left,1,1,, +EXT06,right,0,1,, +EXT07,left,0,1,, +EXT07,right,0,1,, +EXT08,left,0,0,, +EXT08,right,1,0,, +EXT09,left,0,0,, +EXT09,right,0,0,, +EXT10,left,0,1,, +EXT10,right,0,1,, +EXT11,left,0,1,, +EXT11,right,1,1,, +EXT12,left,0,0,, +EXT12,right,0,1,, +EXT13,left,1,1,, +EXT13,right,0,1,, +EXT14,left,0,0,, +EXT14,right,0,1,, +EXT15,left,1,1,, +EXT15,right,1,0,, +EXT17,left,0,0,, +EXT17,right,0,1,, +EXT18,left,1,1,, +EXT18,right,1,1,, +EXT19,left,1,1,, +EXT19,right,1,0,, +EXT24,left,1,1,, +EXT24,right,1,1,, +EXT25,left,0,1,, +EXT25,right,0,1,, diff --git a/try2/outputs/classification/left_ethmoid_filled.joblib b/try2/outputs/classification/left_ethmoid_filled.joblib new file mode 100644 index 0000000..73b2189 Binary files /dev/null and b/try2/outputs/classification/left_ethmoid_filled.joblib differ diff --git a/try2/outputs/classification/right_ethmoid_filled.joblib b/try2/outputs/classification/right_ethmoid_filled.joblib new file mode 100644 index 0000000..83fbd90 Binary files /dev/null and b/try2/outputs/classification/right_ethmoid_filled.joblib differ diff --git a/try2/shared/__pycache__/__init__.cpython-39.pyc b/try2/shared/__pycache__/__init__.cpython-39.pyc new file mode 100644 index 0000000..c3c480d Binary files /dev/null and b/try2/shared/__pycache__/__init__.cpython-39.pyc differ diff --git a/try2/shared/__pycache__/classification.cpython-39.pyc b/try2/shared/__pycache__/classification.cpython-39.pyc new file mode 100644 index 0000000..47347a9 Binary files /dev/null and b/try2/shared/__pycache__/classification.cpython-39.pyc differ diff --git a/try2/shared/__pycache__/config.cpython-39.pyc b/try2/shared/__pycache__/config.cpython-39.pyc new file mode 100644 index 0000000..ddb1cd3 Binary files /dev/null and b/try2/shared/__pycache__/config.cpython-39.pyc differ diff --git a/try2/shared/__pycache__/metrics.cpython-39.pyc b/try2/shared/__pycache__/metrics.cpython-39.pyc new file mode 100644 index 0000000..6eb37ff Binary files /dev/null and b/try2/shared/__pycache__/metrics.cpython-39.pyc differ diff --git a/try2/shared/classification.py b/try2/shared/classification.py new file mode 100644 index 0000000..40cf1d4 --- /dev/null +++ b/try2/shared/classification.py @@ -0,0 +1,308 @@ +""" +Post-hoc classification of AEA segmentation masks. + +Given a 3D volume (H×W×D) and a 3-class segmentation mask (BG=0, AEAL=1, AEAR=2), +predicts 2 binary labels per AEA instance (side-agnostic): + + roof_contact : AEA in contact with skull roof (1) vs distant (0) + ethmoid_filled : ethmoid sinus adjacent to AEA is filled/sinusitis (1) vs clear (0) + +Each patient contributes 2 samples (left + right), pooled into a single classifier +per task. LOO-CV is patient-level: both sides of a patient are held out together. +""" + +import os +from typing import Dict, List, Optional, Tuple + +import numpy as np +import joblib +import openpyxl +from sklearn.linear_model import LogisticRegression +from sklearn.ensemble import RandomForestClassifier +from sklearn.preprocessing import StandardScaler +from sklearn.pipeline import Pipeline + + +def load_classification_labels(xlsx_path: str) -> Dict[str, Dict[str, int]]: + """Return dict keyed by patient_code. + + Each entry has per-side labels (still stored separately so we can pool + them into side-agnostic samples in build_feature_matrix): + aeal_roof_contact, aear_roof_contact, + left_ethmoid_filled, right_ethmoid_filled + """ + wb = openpyxl.load_workbook(xlsx_path) + ws = wb.active + labels: Dict[str, Dict[str, int]] = {} + for row in list(ws.iter_rows(values_only=True))[1:]: + code = row[0] + if code is None: + continue + code = str(code).strip() + labels[code] = { + "aeal_roof_contact": 1 if row[4] == "X" else 0, + "aear_roof_contact": 1 if row[6] == "X" else 0, + "left_ethmoid_filled": 1 if row[9] == "X" else 0, + "right_ethmoid_filled": 1 if row[11] == "X" else 0, + } + return labels + +def extract_roof_features(volume: np.ndarray, mask: np.ndarray, class_id: int) -> np.ndarray: + """ + Features for skull-roof contact detection. + + Slices are sorted ascending by ImagePositionPatient[2] (standard DICOM), so + high z-index = superior. "Roof contact" means the anatomy abuts the skull + base / roof, detected by sampling CT intensities just *above* (z+1..z+3) the + topmost mask voxels — bone HU >> soft tissue. + + Returns zeros if class_id is absent. + """ + voxels = np.argwhere(mask == class_id) # (N, 3): rows are (h, w, d) + D = mask.shape[2] + + if len(voxels) == 0: + return np.zeros(6, dtype=np.float32) + + z_coords = voxels[:, 2].astype(np.float32) + z_max = int(z_coords.max()) + top_z_norm = z_max / (D - 1) + centroid_z_norm = float(z_coords.mean()) / (D - 1) + z_extent_norm = (z_coords.max() - z_coords.min()) / (D - 1) + log_count = float(np.log1p(len(voxels))) + + # Sample HU values in a shell of slices just above the top of the mask. + # bone ~400–1900 HU; air ≈ -1000 HU; soft tissue -100..+100 HU. + shell_hu: List[float] = [] + top_voxels = voxels[voxels[:, 2] == z_max] # (M, 3) + for dz in (1, 2, 3): + z_above = z_max + dz + if z_above >= D: + break + hu_vals = volume[top_voxels[:, 0], top_voxels[:, 1], z_above].astype(np.float32) + shell_hu.extend(hu_vals.tolist()) + + if shell_hu: + shell_arr = np.array(shell_hu, dtype=np.float32) + bone_contact_ratio = float(np.mean(shell_arr > 400)) + mean_above_hu = float(np.mean(shell_arr)) + else: + bone_contact_ratio = 0.0 + mean_above_hu = 0.0 + + return np.array( + [top_z_norm, centroid_z_norm, z_extent_norm, log_count, + bone_contact_ratio, mean_above_hu], + dtype=np.float32, + ) + + +def extract_ethmoid_features( + volume: np.ndarray, mask: np.ndarray, class_id: int +) -> np.ndarray: + """ + Ethmoid filling features derived from the middle sagittal slab. + + Best config from sweep: slab ±1 voxel around the middle of the W range + (not the centroid), threshold -200 HU for filled_ratio. + + volume shape: (H, W, D) — W = columns = left-right axis. + Sagittal plane = H×D at fixed W. + Returns zeros if class_id is absent. + """ + voxels = np.argwhere(mask == class_id) # (N, 3): h, w, d + + if len(voxels) == 0: + return np.zeros(6, dtype=np.float32) + + # Middle of the W extent of the mask (not centroid — midpoint of range) + w_min, w_max = int(voxels[:, 1].min()), int(voxels[:, 1].max()) + w_mid = (w_min + w_max) // 2 + w_lo = max(0, w_mid - 1) + w_hi = min(volume.shape[1] - 1, w_mid + 1) + + slab_mask = (voxels[:, 1] >= w_lo) & (voxels[:, 1] <= w_hi) + slab_voxels = voxels[slab_mask] + + if len(slab_voxels) == 0: + return np.zeros(6, dtype=np.float32) + + intensities = volume[ + slab_voxels[:, 0], slab_voxels[:, 1], slab_voxels[:, 2] + ].astype(np.float32) + + mean_hu = float(np.mean(intensities)) + median_hu = float(np.median(intensities)) + std_hu = float(np.std(intensities)) + p10_hu = float(np.percentile(intensities, 10)) + p90_hu = float(np.percentile(intensities, 90)) + # -200 HU threshold: better air/soft-tissue boundary than -500 + filled_ratio = float(np.mean(intensities > -200)) + + return np.array( + [mean_hu, median_hu, std_hu, p10_hu, p90_hu, filled_ratio], + dtype=np.float32, + ) + + +def extract_features( + volume: np.ndarray, + mask: np.ndarray, +) -> Dict[str, np.ndarray]: + """Return feature vectors for both sides, for both tasks.""" + return { + "roof_L": extract_roof_features(volume, mask, class_id=1), + "roof_R": extract_roof_features(volume, mask, class_id=2), + "ethmoid_L": extract_ethmoid_features(volume, mask, class_id=1), + "ethmoid_R": extract_ethmoid_features(volume, mask, class_id=2), + } + + +def build_feature_matrix( + patient_codes: List[str], + volumes: List[np.ndarray], + masks: List[np.ndarray], + labels: Dict[str, Dict[str, int]], +) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], np.ndarray]: + """ + Pool left and right sides into side-agnostic feature matrices. + + Each patient contributes 2 rows (left side first, then right side). + Returns: + X_dict — {"roof": (2N, 6), "ethmoid": (2N, 6)} + y_dict — {"roof_contact": (2N,), "ethmoid_filled": (2N,)} + patient_idx — (2N,) integer array mapping each row to its patient index. + Used for patient-level LOO-CV (hold out both sides together). + """ + feat_roof, feat_eth = [], [] + y_roof, y_eth = [], [] + patient_idx = [] + + for p_idx, (code, vol, msk) in enumerate(zip(patient_codes, volumes, masks)): + if code not in labels: + continue + feats = extract_features(vol, msk) + lbl = labels[code] + # Left side (class_id=1) + feat_roof.append(feats["roof_L"]) + feat_eth.append(feats["ethmoid_L"]) + y_roof.append(lbl["aeal_roof_contact"]) + y_eth.append(lbl["left_ethmoid_filled"]) + patient_idx.append(p_idx) + # Right side (class_id=2) + feat_roof.append(feats["roof_R"]) + feat_eth.append(feats["ethmoid_R"]) + y_roof.append(lbl["aear_roof_contact"]) + y_eth.append(lbl["right_ethmoid_filled"]) + patient_idx.append(p_idx) + + X_dict = { + "roof": np.stack(feat_roof) if feat_roof else np.empty((0, 6)), + "ethmoid": np.stack(feat_eth) if feat_eth else np.empty((0, 6)), + } + y_dict = { + "roof_contact": np.array(y_roof, dtype=np.int32), + "ethmoid_filled": np.array(y_eth, dtype=np.int32), + } + return X_dict, y_dict, np.array(patient_idx, dtype=np.int32) + + +def _make_clf(task: str) -> Pipeline: + """Return the best classifier pipeline for the given task. + + All tasks use LogisticRegression — stable under LOO-CV with small + positive class sizes (23–26 positives out of 133). RandomForest showed + higher variance and collapsed recall in LOO-CV despite appearing better + in a fixed-split sweep. + """ + return Pipeline([ + ("scaler", StandardScaler()), + ("clf", LogisticRegression(class_weight="balanced", max_iter=1000, C=1.0)), + ]) + + +TASK_TO_FEATURES = { + "roof_contact": "roof", + "ethmoid_filled": "ethmoid", +} + + +def train_classifiers( + X_dict: Dict[str, np.ndarray], + y_dict: Dict[str, np.ndarray], +) -> Dict[str, Pipeline]: + """Train one pipeline per task. Returns dict keyed by task name.""" + clfs: Dict[str, Pipeline] = {} + for task, feat_key in TASK_TO_FEATURES.items(): + X = X_dict[feat_key] + y = y_dict[task] + if len(X) == 0: + continue + clf = _make_clf(task) + clf.fit(X, y) + clfs[task] = clf + return clfs + + +def save_classifiers(clfs: Dict[str, Pipeline], output_dir: str) -> None: + os.makedirs(output_dir, exist_ok=True) + for task, clf in clfs.items(): + joblib.dump(clf, os.path.join(output_dir, f"{task}.joblib")) + + +def load_classifiers(model_dir: str) -> Dict[str, Pipeline]: + clfs: Dict[str, Pipeline] = {} + for task in TASK_TO_FEATURES: + path = os.path.join(model_dir, f"{task}.joblib") + if os.path.exists(path): + clfs[task] = joblib.load(path) + return clfs + + +def predict_classifications( + volume: np.ndarray, + mask: np.ndarray, + clfs: Dict[str, Pipeline], +) -> Dict[str, Dict[str, int]]: + """ + Given a single patient's volume and mask, return side-agnostic predictions. + + Returns: {"left": {"roof_contact": 0/1, "ethmoid_filled": 0/1}, + "right": {"roof_contact": 0/1, "ethmoid_filled": 0/1}} + """ + feats = extract_features(volume, mask) + side_feat = {"left": ("roof_L", "ethmoid_L"), "right": ("roof_R", "ethmoid_R")} + results: Dict[str, Dict[str, int]] = {} + for side, (rf_key, eth_key) in side_feat.items(): + results[side] = {} + if "roof_contact" in clfs: + results[side]["roof_contact"] = int( + clfs["roof_contact"].predict(feats[rf_key].reshape(1, -1))[0] + ) + if "ethmoid_filled" in clfs: + results[side]["ethmoid_filled"] = int( + clfs["ethmoid_filled"].predict(feats[eth_key].reshape(1, -1))[0] + ) + return results + + +def predict_classifications_proba( + volume: np.ndarray, + mask: np.ndarray, + clfs: Dict[str, Pipeline], +) -> Dict[str, Dict[str, float]]: + """Same as predict_classifications but returns probability of positive class.""" + feats = extract_features(volume, mask) + side_feat = {"left": ("roof_L", "ethmoid_L"), "right": ("roof_R", "ethmoid_R")} + results: Dict[str, Dict[str, float]] = {} + for side, (rf_key, eth_key) in side_feat.items(): + results[side] = {} + if "roof_contact" in clfs: + results[side]["roof_contact"] = float( + clfs["roof_contact"].predict_proba(feats[rf_key].reshape(1, -1))[0, 1] + ) + if "ethmoid_filled" in clfs: + results[side]["ethmoid_filled"] = float( + clfs["ethmoid_filled"].predict_proba(feats[eth_key].reshape(1, -1))[0, 1] + ) + return results