"""v5 segmentation trainer: tumor-positive + healthy-brain joint training. Headline differences from v3: 1. Reads dataset_v5 (positive + negative samples; negatives have empty PNG masks). Built by `python prepare_v5_dataset.py`. 2. Balanced batch sampler: each batch contains ~50% positive (non-empty mask) and ~50% negative (empty mask). Without this the gradient is dominated by whichever class has more samples and the positive bias / no-detection bias persist. 3. Loss is Dice (positives only) + BCE (all samples). On empty-mask samples Dice is undefined; we contribute zero Dice loss for them and rely on BCE alone to penalise spurious positive predictions. 4. Optional modality-dropout augmentation: with probability `p_mod_drop`, randomly mask one or two of the three input channels (set to channel mean). Teaches the model to handle single-modality grayscale inputs without needing the v3 cascade T1c trick. 5. Reports false-positive-rate on validation (pixels predicted >= threshold inside scans whose ground-truth mask is empty) in addition to per-scan Dice / IoU. This is the metric that matters for the no-tumor-bias bug; v3's val_dice never measured it. Output: segmentation_artifacts/attention_unet_v5/ best_model.pt (val Dice + tiny FP-rate penalty -> single picked checkpoint) last.pt (most recent epoch) history.json per-epoch metrics training.log text log including FP rate training_curves.png Dice + FP-rate vs epoch evaluation_metrics.json final val + test report Usage: python src/train_segmentation_v5.py --data_dir dataset_v5 \\ --epochs 25 --batch_size 8 --image_size 256 Run prepare_v5_dataset.py first to build dataset_v5/. """ from __future__ import annotations import argparse import json import math import os import random import sys import time import warnings from pathlib import Path import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from PIL import Image from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler # ----------------------------------------------------------------------- # Reproducibility # ----------------------------------------------------------------------- def _set_seed(s: int = 42) -> None: random.seed(s) np.random.seed(s) torch.manual_seed(s) if torch.cuda.is_available(): torch.cuda.manual_seed_all(s) # ----------------------------------------------------------------------- # Dataset # ----------------------------------------------------------------------- class V5SegDataset(Dataset): """Dataset that exposes per-sample 'has_tumor' flag for balanced sampling. Supports optional in-RAM caching of the raw image / mask bytes via cache_in_ram=True. On Linux (Colab, most cloud GPUs) DataLoader workers fork from the main process with copy-on-write semantics, so the cache is physically shared across workers without N-times memory blow-up. For dataset_v8 (~860 MB on disk) the cached representation is ~900 MB of raw bytes -- trivial on a machine with 100+ GB RAM, transformational for I/O-bound training. """ def __init__(self, root: Path, image_size: int, p_mod_drop: float = 0.0, augment: bool = False, imagenet_normalize: bool = True, cache_in_ram: bool = False): self.root = Path(root) self.image_size = image_size self.p_mod_drop = p_mod_drop self.augment = augment self.imagenet_normalize = imagenet_normalize images = sorted((self.root / 'images').iterdir()) masks = sorted((self.root / 'masks').iterdir()) assert len(images) == len(masks), ( f'image/mask count mismatch in {root}: {len(images)} vs {len(masks)}' ) # Compute has_tumor by mask non-zero on a cheap thumbnail. self.samples = [] for img_p, msk_p in zip(images, masks): assert img_p.stem == msk_p.stem, f'pair mismatch: {img_p.name} vs {msk_p.name}' try: m = np.array(Image.open(msk_p).convert('L').resize((64, 64), Image.NEAREST)) has_tumor = bool((m > 127).any()) except Exception: has_tumor = False self.samples.append((img_p, msk_p, has_tumor)) # Optional RAM cache. Stored as raw bytes (Image.open(BytesIO) on # demand) - keeps it small AND allows PIL's lazy decode to work. self._cache = None if cache_in_ram: print(f'[V5SegDataset] caching {len(self.samples)} files in RAM from {root}...', flush=True) t0 = time.time() self._cache = {} for img_p, msk_p, _ in self.samples: self._cache[(str(img_p), 0)] = img_p.read_bytes() self._cache[(str(msk_p), 1)] = msk_p.read_bytes() total_mb = sum(len(v) for v in self._cache.values()) / 1e6 print(f'[V5SegDataset] cached {total_mb:.0f} MB in {time.time() - t0:.0f}s', flush=True) def _open_image(self, p): if self._cache is not None: import io as _io return Image.open(_io.BytesIO(self._cache[(str(p), 0)])) return Image.open(p) def _open_mask(self, p): if self._cache is not None: import io as _io return Image.open(_io.BytesIO(self._cache[(str(p), 1)])) return Image.open(p) def __len__(self) -> int: return len(self.samples) def has_tumor_flags(self) -> list: return [s[2] for s in self.samples] def __getitem__(self, i: int): img_p, msk_p, has_tumor = self.samples[i] img = self._open_image(img_p).convert('RGB').resize((self.image_size, self.image_size), Image.BILINEAR) msk = self._open_mask(msk_p).convert('L').resize((self.image_size, self.image_size), Image.NEAREST) x = np.asarray(img, dtype=np.float32) / 255.0 y = (np.asarray(msk, dtype=np.uint8) > 127).astype(np.float32) if self.augment: # Horizontal flip if random.random() < 0.5: x = x[:, ::-1, :].copy() y = y[:, ::-1].copy() # Vertical flip if random.random() < 0.2: x = x[::-1, :, :].copy() y = y[::-1, :].copy() # Brightness / contrast jitter if random.random() < 0.5: x = np.clip(x * (1.0 + (random.random() - 0.5) * 0.2), 0, 1) x = np.clip(x + (random.random() - 0.5) * 0.1, 0, 1) # Modality dropout. Zero (or grey) out 1 or 2 channels at random. # Forces the model to predict from any single-modality view. if self.augment and self.p_mod_drop > 0 and random.random() < self.p_mod_drop: n_drop = random.choice([1, 2]) chans = random.sample([0, 1, 2], n_drop) for c in chans: x[:, :, c] = x[:, :, c].mean() # grey replacement, not zero if self.imagenet_normalize: mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) std = np.array([0.229, 0.224, 0.225], dtype=np.float32) x = (x - mean) / std x_t = torch.from_numpy(x.transpose(2, 0, 1).copy()).float() y_t = torch.from_numpy(y[None].copy()).float() return x_t, y_t, float(has_tumor) # ----------------------------------------------------------------------- # Model # ----------------------------------------------------------------------- def _build_model() -> nn.Module: """SMP UNet + ResNet34 (matches v3) so we can compare cleanly.""" import segmentation_models_pytorch as smp return smp.Unet( encoder_name='resnet34', encoder_weights='imagenet', in_channels=3, classes=1, ) # ----------------------------------------------------------------------- # Loss # ----------------------------------------------------------------------- class DiceBceLoss(nn.Module): """Dice (positives only) + BCE (all samples). For samples with an entirely empty target mask, Dice is degenerate (target_sum=0). We zero out the Dice contribution for those and let BCE drive the gradient. BCE penalises any false-positive pixel in those samples - exactly the regularisation v3 lacked. """ def __init__(self, dice_w: float = 0.7, bce_w: float = 0.3, pos_weight: float = 1.0): super().__init__() self.dice_w = dice_w self.bce_w = bce_w self.register_buffer('pos_weight', torch.tensor(pos_weight)) def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor: eps = 1e-6 pred = torch.sigmoid(logits) # Per-sample target sum used to mask Dice for empty targets. target_sum = target.flatten(1).sum(dim=1) # (B,) pos_mask = (target_sum > 0).float() # 1 if positive sample, 0 if negative # Dice per sample inter = (pred * target).flatten(1).sum(dim=1) denom = pred.flatten(1).sum(dim=1) + target.flatten(1).sum(dim=1) dice = (2 * inter + eps) / (denom + eps) dice_loss_per_sample = (1.0 - dice) * pos_mask # Reduce safely. n_pos = pos_mask.sum().clamp(min=1.0) dice_loss = dice_loss_per_sample.sum() / n_pos # BCE per pixel, all samples. bce = F.binary_cross_entropy_with_logits( logits, target, pos_weight=self.pos_weight, ) return self.dice_w * dice_loss + self.bce_w * bce # ----------------------------------------------------------------------- # Metrics # ----------------------------------------------------------------------- @torch.no_grad() def _evaluate(model: nn.Module, loader: DataLoader, device: torch.device, threshold: float = 0.5) -> dict: """Compute Dice + IoU on positive samples and FP-rate on negatives.""" model.eval() dices = [] ious = [] fp_rates = [] # for each negative sample: fraction of pixels predicted >= threshold n_pos = 0 n_neg = 0 for x, y, has in loader: x = x.to(device) y = y.to(device) p = torch.sigmoid(model(x)) m = (p >= threshold).float() # Per-sample metrics for i in range(x.size(0)): yi = y[i] mi = m[i] if yi.sum() > 0: inter = (mi * yi).sum().item() pred_sum = mi.sum().item() tgt_sum = yi.sum().item() d = (2 * inter + 1e-6) / (pred_sum + tgt_sum + 1e-6) u = (inter + 1e-6) / (pred_sum + tgt_sum - inter + 1e-6) dices.append(d) ious.append(u) n_pos += 1 else: fp_rates.append(mi.mean().item()) n_neg += 1 return { 'n_positive': n_pos, 'n_negative': n_neg, 'dice_mean': float(np.mean(dices)) if dices else 0.0, 'iou_mean': float(np.mean(ious)) if ious else 0.0, 'fp_rate_mean': float(np.mean(fp_rates)) if fp_rates else 0.0, 'fp_rate_p95': float(np.percentile(fp_rates, 95)) if fp_rates else 0.0, } # ----------------------------------------------------------------------- # Training # ----------------------------------------------------------------------- def _make_balanced_loader(ds: V5SegDataset, batch_size: int, num_workers: int, shuffle_class_balance: bool = True) -> DataLoader: """50/50 batches of positives vs negatives via WeightedRandomSampler.""" flags = ds.has_tumor_flags() n_pos = sum(1 for f in flags if f) n_neg = sum(1 for f in flags if not f) if n_pos == 0 or n_neg == 0 or not shuffle_class_balance: return DataLoader(ds, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True, drop_last=True) w_pos = 1.0 / n_pos w_neg = 1.0 / n_neg weights = [w_pos if f else w_neg for f in flags] sampler = WeightedRandomSampler(weights, num_samples=len(ds), replacement=True) return DataLoader(ds, batch_size=batch_size, sampler=sampler, num_workers=num_workers, pin_memory=True, drop_last=True) def main(): ap = argparse.ArgumentParser() ap.add_argument('--data_dir', default='dataset_v5') ap.add_argument('--output_dir', default='segmentation_artifacts/attention_unet_v5') ap.add_argument('--epochs', type=int, default=25) ap.add_argument('--batch_size', type=int, default=8) ap.add_argument('--image_size', type=int, default=256) ap.add_argument('--lr', type=float, default=1e-4) ap.add_argument('--weight_decay', type=float, default=1e-5) ap.add_argument('--num_workers', type=int, default=2) ap.add_argument('--p_mod_drop', type=float, default=0.3, help='Probability of dropping 1-2 channels per train sample.') ap.add_argument('--bce_pos_weight', type=float, default=2.0, help='Multiplier on positive-pixel BCE; >1 reduces false-negative rate.') ap.add_argument('--seed', type=int, default=42) ap.add_argument('--resume', default=None, help='Path to a .pt to resume from.') args = ap.parse_args() _set_seed(args.seed) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f'[info] device={device}, dataset={args.data_dir}, output={args.output_dir}') train_ds = V5SegDataset(Path(args.data_dir) / 'train', args.image_size, p_mod_drop=args.p_mod_drop, augment=True) val_ds = V5SegDataset(Path(args.data_dir) / 'val', args.image_size, augment=False) test_ds = V5SegDataset(Path(args.data_dir) / 'test', args.image_size, augment=False) print(f'[info] train n={len(train_ds)} ' f'(pos={sum(train_ds.has_tumor_flags())}, ' f'neg={len(train_ds) - sum(train_ds.has_tumor_flags())})') print(f'[info] val n={len(val_ds)}') print(f'[info] test n={len(test_ds)}') train_loader = _make_balanced_loader(train_ds, args.batch_size, args.num_workers) val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True) test_loader = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True) model = _build_model().to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) loss_fn = DiceBceLoss(dice_w=0.7, bce_w=0.3, pos_weight=args.bce_pos_weight).to(device) scaler = torch.cuda.amp.GradScaler() if device.type == 'cuda' else None out_dir = Path(args.output_dir) out_dir.mkdir(parents=True, exist_ok=True) history: list = [] best_score = -1.0 start_epoch = 1 if args.resume: ckpt = torch.load(args.resume, map_location=device, weights_only=False) model.load_state_dict(ckpt['state_dict']) if 'optimizer' in ckpt: optimizer.load_state_dict(ckpt['optimizer']) start_epoch = int(ckpt.get('epoch', 0)) + 1 best_score = float(ckpt.get('best_score', -1.0)) print(f'[info] resumed from {args.resume} at epoch {start_epoch}') log_path = out_dir / 'training.log' with log_path.open('a', encoding='utf-8') as logf: for ep in range(start_epoch, args.epochs + 1): model.train() t0 = time.time() train_loss = 0.0 n_batches = 0 for x, y, _has in train_loader: x = x.to(device, non_blocking=True) y = y.to(device, non_blocking=True) optimizer.zero_grad(set_to_none=True) if scaler is not None: with torch.cuda.amp.autocast(): logits = model(x) loss = loss_fn(logits, y) if not torch.isfinite(loss): # NaN guard continue scaler.scale(loss).backward() scaler.unscale_(optimizer) grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) if not torch.isfinite(grad_norm): scaler.update() continue scaler.step(optimizer) scaler.update() else: logits = model(x) loss = loss_fn(logits, y) if not torch.isfinite(loss): continue loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() train_loss += float(loss) n_batches += 1 scheduler.step() avg_loss = train_loss / max(n_batches, 1) val_metrics = _evaluate(model, val_loader, device) # Composite score: Dice on positives MINUS a penalty proportional to # FP-rate on negatives. A model that gets 0.7 Dice but 30% pixel-FPR # on healthy brains scores 0.7 - 0.3*5 = -0.8 - worse than a model # that gets 0.6 Dice with 1% FPR (0.6 - 0.01*5 = 0.55). This is the # objective that actually matches the bug we're trying to fix. composite = val_metrics['dice_mean'] - 5.0 * val_metrics['fp_rate_mean'] elapsed = time.time() - t0 line = ( f'[epoch {ep:02d}/{args.epochs}] ' f'train_loss={avg_loss:.4f} ' f'val_dice={val_metrics["dice_mean"]:.4f} ' f'val_iou={val_metrics["iou_mean"]:.4f} ' f'val_fp_rate={val_metrics["fp_rate_mean"]:.4f} ' f'val_fp_p95={val_metrics["fp_rate_p95"]:.4f} ' f'composite={composite:.4f} ' f'lr={scheduler.get_last_lr()[0]:.2e} ' f'({elapsed:.1f}s)' ) print(line) logf.write(line + '\n') logf.flush() history.append({ 'epoch': ep, 'train_loss': avg_loss, 'val': val_metrics, 'composite': composite, }) torch.save({ 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': ep, 'best_score': best_score, 'architecture': 'Unet', 'encoder': 'resnet34', 'image_size': args.image_size, 'config': {'base_filters': 32, 'dropout': 0.2, 'image_size': args.image_size, '_v5': True}, }, out_dir / 'last.pt') if composite > best_score: best_score = composite torch.save({ 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': ep, 'best_score': best_score, 'architecture': 'Unet', 'encoder': 'resnet34', 'image_size': args.image_size, 'config': {'base_filters': 32, 'dropout': 0.2, 'image_size': args.image_size, '_v5': True}, }, out_dir / 'best_model.pt') print(f' -> new best composite={best_score:.4f}; weights saved.') # Final test evaluation on the best checkpoint. best = torch.load(out_dir / 'best_model.pt', map_location=device, weights_only=False) model.load_state_dict(best['state_dict']) test_metrics = _evaluate(model, test_loader, device) final = {'val_best': max(history, key=lambda h: h['composite'])['val'], 'test': test_metrics, 'best_composite': best_score} print(f'[done] best composite={best_score:.4f} ' f'test_dice={test_metrics["dice_mean"]:.4f} ' f'test_fp_rate={test_metrics["fp_rate_mean"]:.4f}') with (out_dir / 'evaluation_metrics.json').open('w', encoding='utf-8') as fh: json.dump(final, fh, indent=2) with (out_dir / 'history.json').open('w', encoding='utf-8') as fh: json.dump(history, fh, indent=2) if __name__ == '__main__': warnings.filterwarnings('ignore', category=UserWarning) main()