Part 1: Mammography Baseline — DICOM to ConvNeXt

!pip install timm albumentations torcheval scikit-learn opencv-python-headless tqdm pydicom pylibjpeg pylibjpeg-libjpeg

import os, gc, time, copy, random
from pathlib import Path
from collections import defaultdict
 
import numpy as np
import pandas as pd
import cv2
import matplotlib.pyplot as plt
 
import pydicom
from pydicom.pixel_data_handlers.util import apply_voi_lut
 
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import Dataset, DataLoader
 
import timm
import albumentations as A
from albumentations.pytorch import ToTensorV2
 
from sklearn.model_selection import StratifiedGroupKFold
from sklearn.metrics import roc_auc_score
from torcheval.metrics.functional import binary_auroc
from tqdm import tqdm
 
print("PyTorch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("GPU:", torch.cuda.get_device_name(0))

CONFIG = {
    'data_dir':         './data/rsna-breast-cancer-detection',
    'train_images_dir': './data/rsna-breast-cancer-detection/train_images',
    'csv_path':         './data/rsna-breast-cancer-detection/train.csv',
    'processed_dir':    './data/processed_pngs',   # pre-converted 8-bit PNGs
    'models_folder':    './saved_models',
 
    'model_name':       'convnext_small.fb_in22k_ft_in1k',
    'img_h':            2048,          # height (tall axis of mammogram)
    'img_w':            1024,          # width
    'num_classes':      2,             # softmax: 0=benign, 1=malignant
    'drop_rate':        0.0,
    'drop_path_rate':   0.0,
 
    'seed':             42,
    'epochs':           15,
    'train_batch_size': 4,             # large images require small batches
    'valid_batch_size': 8,
    'n_accumulate':     8,             # effective batch = 4 × 8 = 32
    'device':           'cuda' if torch.cuda.is_available() else 'cpu',
    'n_folds':          4,             # same as winning solution
    'group_col':        'patient_id',
 
    'learning_rate':    2e-5,
    'weight_decay':     1e-6,
 
    'scheduler':        'CosineAnnealingLR',
    'T_max':            500,
    'min_lr':           1e-7,
}
 
def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
 
set_seed(CONFIG['seed'])
os.makedirs(CONFIG['models_folder'], exist_ok=True)
os.makedirs(CONFIG['processed_dir'], exist_ok=True)
print("Device:", CONFIG['device'])

df = pd.read_csv(CONFIG['csv_path'])
df = df.rename(columns={'cancer': 'target'})
 
print(f"Total images     : {len(df)}")
print(f"Unique patients  : {df.patient_id.nunique()}")
print(f"Malignant (1)    : {df.target.sum()} ({100*df.target.mean():.2f}%)")
print(f"\nViews: {df.view.value_counts().to_dict()}")
print(f"Laterality: {df.laterality.value_counts().to_dict()}")
df.head()

# TODO
# Task A — Label granularity:
#   The label 'cancer' is per-patient, but images are per-view. Does every
#   view of a cancerous patient get label=1, even the healthy breast (R vs L)?
#   Check using the 'laterality' column. This matters for training signal quality.
#
# Task B — Patient-level vs image-level leakage:
#   Why is group_col='patient_id' critical for the CV split?
#   What would happen if you split by image_id instead?
#
# Task C — Explore metadata:
#   Plot cancer rate by (a) view (CC vs MLO), (b) laterality, (c) age group.
#   Does machine_id correlate with cancer rate? (hint: site-level confounds)
 
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
 
df.groupby('view')['target'].mean().plot(kind='bar', ax=axes[0], title='Cancer rate by view', color='steelblue')
df.groupby('laterality')['target'].mean().plot(kind='bar', ax=axes[1], title='Cancer rate by laterality', color='coral')
df['age_bin'] = pd.cut(df['age'], bins=[30, 40, 50, 60, 70, 80, 90])
df.groupby('age_bin')['target'].mean().plot(kind='bar', ax=axes[2], title='Cancer rate by age', color='mediumseagreen')
 
for ax in axes: ax.set_ylabel('Cancer rate'); ax.tick_params(axis='x', rotation=45)
plt.tight_layout(); plt.show()

def read_dicom(path: str, voi_lut: bool = True) -> np.ndarray:
    """Read a DICOM file and return a normalised uint8 numpy array."""
    dcm = pydicom.dcmread(path)
    
    if voi_lut:
        # Apply the VOI LUT (window/level) embedded in the DICOM header.
        # This maps the raw stored values to a display-meaningful range.
        data = apply_voi_lut(dcm.pixel_array, dcm)
    else:
        data = dcm.pixel_array
 
    # MONOCHROME1: pixel value 0 = white (dense tissue), high = black (air)
    # We want the standard radiological convention: bright tissue, dark background.
    if dcm.PhotometricInterpretation == 'MONOCHROME1':
        data = np.max(data) - data   # invert
 
    # Normalise to uint8
    data = data.astype(np.float32)
    data -= data.min()
    if data.max() > 0:
        data /= data.max()
    data = (data * 255).astype(np.uint8)
    return data

# TODO
# Task A — Inspect a raw DICOM:
#   Load one DICOM and print dcm.PhotometricInterpretation, dcm.BitsStored,
#   dcm.PixelRepresentation, and dcm.pixel_array.shape.
#   What is the raw pixel value range before normalisation?
#
# Task B — VOI LUT effect:
#   Read the same DICOM with voi_lut=True and voi_lut=False.
#   Plot both histograms. When does the VOI LUT make a visible difference?
#
# Task C — MONOCHROME1 vs MONOCHROME2:
#   Find one example of each in the dataset. Plot them side by side,
#   before and after the photometric inversion step.
 
# Example: Load and display one mammogram
# sample_path = f"{CONFIG['train_images_dir']}/{df.patient_id[0]}/{df.image_id[0]}.dcm"
# img = read_dicom(sample_path)
# plt.figure(figsize=(4, 8))
# plt.imshow(img, cmap='gray'); plt.axis('off'); plt.title('Raw mammogram'); plt.show()
# print(f"Shape: {img.shape}, dtype: {img.dtype}, range: [{img.min()}, {img.max()}]")

def crop_breast_roi_threshold(img: np.ndarray, threshold: int = 10) -> np.ndarray:
    """
    Simple threshold-based breast ROI extraction.
    Finds the bounding box of pixels brighter than `threshold` and crops.
    Works well for clean backgrounds but can fail on noisy scanners.
    """
    # Binarise: breast tissue is bright, background is ~0
    mask = (img > threshold).astype(np.uint8)
    
    # Find the largest connected component (the breast)
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
    
    if num_labels < 2:
        return img   # no component found, return original
    
    # Component 0 is background; find largest foreground component
    largest_label = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA])
    
    x  = stats[largest_label, cv2.CC_STAT_LEFT]
    y  = stats[largest_label, cv2.CC_STAT_TOP]
    w  = stats[largest_label, cv2.CC_STAT_WIDTH]
    h  = stats[largest_label, cv2.CC_STAT_HEIGHT]
    
    return img[y:y+h, x:x+w]

# TODO
# Task A — Visualise cropping quality:
#   Apply crop_breast_roi_threshold to 6 different images.
#   Show original vs cropped side by side.
#   Cases to check: normal scan, noisy scanner, implant, dense breast.
#
# Task B — Morphological cleanup:
#   Add a cv2.morphologyEx OPEN step before finding components to remove
#   small bright artifacts (scanner labels, rulers). Does it help?
#
# Task C — Compare to YOLOX (preview for Part 2):
#   Note any cases where threshold cropping produces a poor crop.
#   These are exactly the failure cases YOLOX is trained to handle.
 
# Quick test
# raw = read_dicom(sample_path)
# cropped = crop_breast_roi_threshold(raw)
# fig, axes = plt.subplots(1, 2, figsize=(10, 8))
# axes[0].imshow(raw, cmap='gray'); axes[0].set_title(f'Original {raw.shape}')
# axes[1].imshow(cropped, cmap='gray'); axes[1].set_title(f'Cropped {cropped.shape}')
# for ax in axes: ax.axis('off')
# plt.tight_layout(); plt.show()

def convert_dicom_to_png(row, src_dir: str, dst_dir: str, apply_crop: bool = True):
    """Convert a single DICOM to a normalised, optionally-cropped PNG."""
    src_path = os.path.join(src_dir, str(row.patient_id), f'{row.image_id}.dcm')
    dst_path = os.path.join(dst_dir, f'{row.patient_id}_{row.image_id}.png')
    
    if os.path.exists(dst_path):
        return dst_path   # already converted
    
    img = read_dicom(src_path)
    if apply_crop:
        img = crop_breast_roi_threshold(img)
    cv2.imwrite(dst_path, img)
    return dst_path
 
# TODO
# Task: Run this conversion on the full training set.
#   Parallelise with concurrent.futures.ThreadPoolExecutor (I/O bound)
#   or multiprocessing.Pool (CPU bound) for speed.
#   Estimated time: ~2-4 hours for 54,000 images on a single CPU core.
 
from concurrent.futures import ThreadPoolExecutor
from functools import partial
 
def batch_convert(df, src_dir, dst_dir, n_workers=8):
    convert_fn = partial(convert_dicom_to_png, src_dir=src_dir, dst_dir=dst_dir)
    with ThreadPoolExecutor(max_workers=n_workers) as executor:
        paths = list(tqdm(
            executor.map(convert_fn, [row for _, row in df.iterrows()]),
            total=len(df), desc='Converting DICOMs'
        ))
    return paths
 
# Uncomment to run:
# paths = batch_convert(df, CONFIG['train_images_dir'], CONFIG['processed_dir'])
# df['path'] = paths
 
# OR: point to pre-processed paths if conversion already done
df['path'] = df.apply(
    lambda r: os.path.join(CONFIG['processed_dir'], f"{r.patient_id}_{r.image_id}.png"), axis=1
)
print("Paths added to dataframe.")

def get_mammography_augmentations(CONFIG):
    img_h, img_w = CONFIG['img_h'], CONFIG['img_w']
    
    train_transform = A.Compose([
        # Geometry
        A.HorizontalFlip(p=0.5),
        A.VerticalFlip(p=0.5),
        A.ShiftScaleRotate(
            shift_limit=0.05, scale_limit=0.05,
            rotate_limit=10, border_mode=cv2.BORDER_CONSTANT,
            value=0, p=0.5
        ),
        
        # Pixel-level — grayscale-safe
        A.RandomBrightnessContrast(brightness_limit=0.2, contrast_limit=0.2, p=0.5),
        A.OneOf([
            A.GaussianBlur(blur_limit=(3, 5)),
            A.MotionBlur(blur_limit=5),
            A.MedianBlur(blur_limit=5),
        ], p=0.3),
        A.GaussNoise(var_limit=(5.0, 20.0), p=0.3),
        
        # Distortions — subtle, preserve tissue structure
        A.OneOf([
            A.ElasticTransform(alpha=1, sigma=20, p=0.5),
            A.GridDistortion(num_steps=5, distort_limit=0.3, p=0.5),
        ], p=0.3),
 
        # Simulate lower-resolution acquisitions
        A.Downscale(scale_range=(0.5, 0.9), p=0.3),
 
        # Resize to model input
        A.Resize(img_h, img_w),
 
        # Regularisation
        A.CoarseDropout(
            max_holes=1,
            max_height=int(img_h * 0.2),
            max_width=int(img_w * 0.2),
            num_holes_range=(1, 1),
            p=0.5
        ),
 
        # Normalise with ImageNet stats (ConvNeXt pretrained on ImageNet)
        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225],
                    max_pixel_value=255.0, p=1.0),
        ToTensorV2(),
    ])
 
    valid_transform = A.Compose([
        A.Resize(img_h, img_w),
        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225],
                    max_pixel_value=255.0, p=1.0),
        ToTensorV2(),
    ])
 
    return {'train': train_transform, 'valid': valid_transform}
 
data_transforms = get_mammography_augmentations(CONFIG)
print("Train transforms:\n", data_transforms['train'])

# TODO
# Task A — Grayscale → RGB conversion:
#   PNGs saved above are grayscale. Albumentations and timm expect 3-channel
#   input. Verify that the Dataset class below handles the cv2.COLOR_GRAY2RGB
#   conversion. What would happen if you passed a single-channel tensor to
#   a model expecting 3 channels?
#
# Task B — Anatomy-aware flipping:
#   In mammography, left (L) and right (R) breasts are mirror images.
#   A common strategy is to always flip R images to face left (normalise
#   laterality) before augmentation. Implement this as a preprocessing step.
#
# Task C — CLAHE for mammography:
#   CLAHE (Contrast Limited Adaptive Histogram Equalisation) is widely used
#   in medical imaging. Add A.CLAHE(clip_limit=2.0, p=0.5) to the pipeline
#   and compare training curves vs without.

class RSNADatasetSimple(Dataset):
    """Sequential dataset — used for validation and inference."""
    def __init__(self, meta_df, transforms=None, do_augmentations=True):
        self.meta_df = meta_df.reset_index(drop=True)
        self.transforms = transforms
        self.do_augmentations = do_augmentations
 
    def __len__(self):
        return len(self.meta_df)
 
    def __getitem__(self, idx):
        row = self.meta_df.iloc[idx]
        target = int(row.target)
 
        img = cv2.imread(row.path, cv2.IMREAD_GRAYSCALE)
        if img is None:
            raise FileNotFoundError(f"Image not found: {row.path}")
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)   # HxWx3
 
        if self.transforms and self.do_augmentations:
            img = self.transforms(image=img)['image']
 
        # One-hot encode for softmax training
        label = torch.zeros(2, dtype=torch.float32)
        label[target] = 1.0
 
        return {'image': img, 'target': label, 'target_int': target}
 
 
class RSNADatasetSampler(Dataset):
    """50/50 positive/negative oversampling — used for training."""
    def __init__(self, meta_df, transforms=None, do_augmentations=True):
        self.df_pos = meta_df[meta_df.target == 1].reset_index(drop=True)
        self.df_neg = meta_df[meta_df.target == 0].reset_index(drop=True)
        self.transforms = transforms
        self.do_augmentations = do_augmentations
 
    def __len__(self):
        return len(self.df_pos) * 2
 
    def _load_img(self, path):
        img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
        if img is None:
            raise FileNotFoundError(f"Image not found: {path}")
        return cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
 
    def __getitem__(self, index):
        # Alternate between positive and negative samples
        if random.random() >= 0.5:
            row = self.df_pos.iloc[index % len(self.df_pos)]
        else:
            row = self.df_neg.iloc[random.randint(0, len(self.df_neg) - 1)]
 
        img = self._load_img(row.path)
        target = int(row.target)
 
        if self.transforms and self.do_augmentations:
            img = self.transforms(image=img)['image']
 
        label = torch.zeros(2, dtype=torch.float32)
        label[target] = 1.0
 
        return {'image': img, 'target': label, 'target_int': target}
 
 
def prepare_loaders(df_train, df_valid, CONFIG, data_transforms, num_workers=4):
    train_ds = RSNADatasetSampler(df_train, transforms=data_transforms['train'])
    valid_ds = RSNADatasetSimple(df_valid,  transforms=data_transforms['valid'])
 
    train_loader = DataLoader(train_ds, batch_size=CONFIG['train_batch_size'],
                              shuffle=True, num_workers=num_workers,
                              pin_memory=True, drop_last=True)
    valid_loader = DataLoader(valid_ds, batch_size=CONFIG['valid_batch_size'],
                              shuffle=False, num_workers=num_workers,
                              pin_memory=True)
    return train_loader, valid_loader

# TODO
# Task A — Verify grayscale → RGB:
#   Load one batch and check image.shape == [B, 3, H, W].
#   Are the three channels identical (since input is grayscale)?
#   This is fine — ImageNet-pretrained models expect RGB, and repeated
#   grayscale channels still carry the correct intensity information.
#
# Task B — Label distribution in sampler:
#   Iterate 100 batches from train_loader. Compute the mean of target[:, 1]
#   (fraction of positives). Does it converge to ~0.5 as expected?
#
# Task C — Laterality normalisation in the dataset:
#   Add a 'flip' flag to the dataframe rows where laterality == 'R',
#   and apply cv2.flip(img, 1) inside __getitem__ before augmentations.
#   This normalises all breasts to face left, reducing the domain shift.

class MammographyConvNeXt(nn.Module):
    def __init__(self, model_name: str, num_classes: int = 2,
                 drop_rate: float = 0.0, drop_path_rate: float = 0.0,
                 pretrained: bool = True):
        super().__init__()
        self.model = timm.create_model(
            model_name,
            pretrained=pretrained,
            drop_rate=drop_rate,
            drop_path_rate=drop_path_rate,
        )
        # Replace classification head
        in_features = self.model.head.fc.in_features
        self.model.head.fc = nn.Linear(in_features, num_classes)
        self.softmax = nn.Softmax(dim=1)
 
    def forward(self, images):
        return self.softmax(self.model(images))
 
    def get_cancer_probability(self, images):
        """Convenience method: returns only the malignant class probability."""
        return self.forward(images)[:, 1]
 
 
def setup_model(CONFIG):
    model = MammographyConvNeXt(
        model_name=CONFIG['model_name'],
        num_classes=CONFIG['num_classes'],
        drop_rate=CONFIG['drop_rate'],
        drop_path_rate=CONFIG['drop_path_rate'],
        pretrained=True,
    )
    return model.to(CONFIG['device'])
 
 
def print_trainable_parameters(model):
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"Trainable: {trainable:,} / Total: {total:,}  ({100*trainable/total:.1f}%)")

model = setup_model(CONFIG)
print_trainable_parameters(model)
 
# Verify forward pass shape
dummy = torch.zeros(2, 3, CONFIG['img_h'], CONFIG['img_w']).to(CONFIG['device'])
with torch.no_grad():
    out = model(dummy)
print(f"Output shape: {out.shape}")           # [2, 2]
print(f"Sum per sample (should be 1.0): {out.sum(dim=1)}")  # softmax sums to 1

# TODO
# Task A — Explore timm ConvNeXt variants:
#   Swap model_name to 'convnext_tiny.fb_in22k_ft_in1k' (smaller, faster)
#   or 'convnext_base.fb_in22k_ft_in1k' (larger, potentially higher accuracy).
#   Compare parameter counts and estimated GPU memory usage.
#
# Task B — Global pooling strategy:
#   The winning team notes MaxPool worked better than AvgPool ("AvgPool tends
#   to wash away the signal"). This makes clinical sense: cancer is a focal
#   finding — the maximum activation in any region matters more than the average.
#   Try modifying the head to use nn.AdaptiveMaxPool2d before the linear layer.
#
# Task C — Mixed precision:
#   At 2048×1024, memory is tight. Enable AMP (Automatic Mixed Precision)
#   using torch.cuda.amp.autocast() and GradScaler. This can 2× throughput.

pF1 = 2 * sum(p_i * y_i) / ( sum(p_i) + sum(y_i) )

def criterion(outputs, targets):
    """Cross-entropy loss for softmax output with one-hot targets."""
    return nn.CrossEntropyLoss()(outputs, targets)
 
 
def probabilistic_f1(y_pred_proba: np.ndarray, y_true: np.ndarray) -> float:
    """
    Probabilistic F1 score (RSNA competition metric).
    
    Args:
        y_pred_proba: predicted probabilities for class 1, shape (N,)
        y_true: binary ground truth labels, shape (N,)
    """
    tp_sum = np.sum(y_pred_proba * y_true)
    pred_sum = np.sum(y_pred_proba)
    true_sum = np.sum(y_true)
    if pred_sum + true_sum == 0:
        return 0.0
    return 2 * tp_sum / (pred_sum + true_sum)
 
 
# Demonstrate pF1 sensitivity
np.random.seed(42)
y_true_demo  = np.random.binomial(1, 0.02, 1000)
y_good       = np.clip(y_true_demo + np.random.normal(0, 0.1, 1000), 0, 1)
y_low_recall = np.clip(y_true_demo * np.random.uniform(0, 0.3, 1000), 0, 1)
 
print(f"Good model  pF1: {probabilistic_f1(y_good, y_true_demo):.4f}")
print(f"Low recall  pF1: {probabilistic_f1(y_low_recall, y_true_demo):.4f}")
print(f"All-zero    pF1: {probabilistic_f1(np.zeros_like(y_true_demo), y_true_demo):.4f}")

# TODO
# Task A — pF1 vs threshold F1:
#   For the same set of predictions, compute pF1 and hard-threshold F1
#   at thresholds in [0.1, 0.3, 0.5, 0.7, 0.9]. Plot all values.
#   Why does pF1 avoid the threshold selection problem?
#
# Task B — Clinical interpretation:
#   A false negative (missed cancer) has far worse consequences than a
#   false positive (unnecessary recall). How does pF1 account for this?
#   Compare to pAUC from the ISIC notebook — which metric is more
#   sensitive to the rare-positive problem?
#
# Task C — Class-weighted loss:
#   With ~2% positives, the model can score well on CE loss by predicting
#   all zeros. Add weight=torch.tensor([0.02, 0.98]) to CrossEntropyLoss
#   to penalise false negatives more heavily. Does it improve pF1?

def fetch_scheduler(optimizer, CONFIG):
    if CONFIG['scheduler'] == 'CosineAnnealingLR':
        return lr_scheduler.CosineAnnealingLR(
            optimizer, T_max=CONFIG['T_max'], eta_min=CONFIG['min_lr'])
    elif CONFIG['scheduler'] == 'CosineAnnealingWarmRestarts':
        return lr_scheduler.CosineAnnealingWarmRestarts(
            optimizer, T_0=25, eta_min=CONFIG['min_lr'])
    return None
 
 
def train_one_epoch(model, optimizer, scheduler, dataloader, device, epoch, CONFIG):
    model.train()
    running_loss, dataset_size = 0.0, 0
    scaler = torch.cuda.amp.GradScaler()   # AMP for memory efficiency
 
    bar = tqdm(enumerate(dataloader), total=len(dataloader))
    for step, data in bar:
        images  = data['image'].to(device, dtype=torch.float)
        targets = data['target'].to(device, dtype=torch.float)
        batch_size = images.size(0)
 
        with torch.cuda.amp.autocast():
            outputs = model(images)           # [B, 2]
            loss    = criterion(outputs, targets) / CONFIG['n_accumulate']
 
        scaler.scale(loss).backward()
 
        if (step + 1) % CONFIG['n_accumulate'] == 0:
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()
            if scheduler is not None:
                scheduler.step()
 
        running_loss += loss.item() * batch_size * CONFIG['n_accumulate']
        dataset_size += batch_size
        epoch_loss = running_loss / dataset_size
 
        bar.set_postfix(Epoch=epoch, Loss=f'{epoch_loss:.4f}',
                        LR=f'{optimizer.param_groups[0]["lr"]:.2e}')
 
    gc.collect()
    return epoch_loss
 
 
@torch.inference_mode()
def valid_one_epoch(model, dataloader, device, epoch, optimizer, return_preds=False):
    model.eval()
    running_loss, dataset_size = 0.0, 0
    all_preds, all_targets = [], []
 
    bar = tqdm(enumerate(dataloader), total=len(dataloader))
    for step, data in bar:
        images  = data['image'].to(device, dtype=torch.float)
        targets = data['target'].to(device, dtype=torch.float)
        t_int   = data['target_int'].numpy()
        batch_size = images.size(0)
 
        outputs = model(images)                    # [B, 2]
        loss    = criterion(outputs, targets)
 
        cancer_prob = outputs[:, 1].cpu().numpy()  # malignant probability
        all_preds.append(cancer_prob)
        all_targets.append(t_int)
 
        running_loss += loss.item() * batch_size
        dataset_size += batch_size
        epoch_loss = running_loss / dataset_size
 
        bar.set_postfix(Epoch=epoch, Val_Loss=f'{epoch_loss:.4f}',
                        LR=f'{optimizer.param_groups[0]["lr"]:.2e}')
 
    gc.collect()
    all_preds   = np.concatenate(all_preds)
    all_targets = np.concatenate(all_targets)
 
    pf1  = probabilistic_f1(all_preds, all_targets)
    auroc = roc_auc_score(all_targets, all_preds)
 
    if return_preds:
        return epoch_loss, pf1, auroc, all_preds, all_targets
    return epoch_loss, pf1, auroc

def run_training(train_loader, valid_loader, model, optimizer, scheduler,
                 CONFIG, model_name='best_model.pth', tolerance_max=8, seed=42):
    set_seed(seed)
    best_pf1 = -np.inf
    best_weights = copy.deepcopy(model.state_dict())
    history = defaultdict(list)
    tolerance = 0
    start = time.time()
 
    for epoch in range(1, CONFIG['epochs'] + 1):
        if tolerance > tolerance_max:
            print(f"Early stopping at epoch {epoch}")
            break
 
        train_loss = train_one_epoch(
            model, optimizer, scheduler,
            train_loader, CONFIG['device'], epoch, CONFIG)
 
        val_loss, val_pf1, val_auroc = valid_one_epoch(
            model, valid_loader, CONFIG['device'], epoch, optimizer)
 
        history['train_loss'].append(train_loss)
        history['val_loss'].append(val_loss)
        history['val_pf1'].append(val_pf1)
        history['val_auroc'].append(val_auroc)
        history['lr'].append(scheduler.get_last_lr()[0] if scheduler else CONFIG['learning_rate'])
 
        print(f"Epoch {epoch:02d} | "
              f"Train Loss: {train_loss:.4f} | "
              f"Val Loss: {val_loss:.4f} | "
              f"Val pF1: {val_pf1:.4f} | "
              f"Val AUC: {val_auroc:.4f}")
 
        if val_pf1 > best_pf1:
            tolerance = 0
            best_pf1 = val_pf1
            best_weights = copy.deepcopy(model.state_dict())
            save_path = os.path.join(CONFIG['models_folder'], model_name)
            torch.save(model.state_dict(), save_path)
            print(f"  ✓ New best pF1: {best_pf1:.4f} — saved to {save_path}")
        else:
            tolerance += 1
 
    elapsed = time.time() - start
    print(f"\nTraining complete in {elapsed//3600:.0f}h {(elapsed%3600)//60:.0f}m")
    print(f"Best pF1: {best_pf1:.4f}")
    model.load_state_dict(best_weights)
    return model, history

sgkf = StratifiedGroupKFold(n_splits=CONFIG['n_folds'], shuffle=True, random_state=CONFIG['seed'])
 
fold_results = []
oof_df_list  = []
 
for fold_n, (train_idx, val_idx) in enumerate(sgkf.split(df, y=df.target, groups=df[CONFIG['group_col']])):
    print(f"\n{'='*60}")
    print(f"FOLD {fold_n + 1} / {CONFIG['n_folds']}")
    print(f"{'='*60}")
 
    fold_train = df.iloc[train_idx].reset_index(drop=True)
    fold_valid = df.iloc[val_idx].reset_index(drop=True)
 
    print(f"  Train: {len(fold_train)} images | Positive rate: {fold_train.target.mean():.3f}")
    print(f"  Valid: {len(fold_valid)} images | Positive rate: {fold_valid.target.mean():.3f}")
 
    set_seed(CONFIG['seed'])
    model     = setup_model(CONFIG)
    optimizer = optim.AdamW(model.parameters(),
                            lr=CONFIG['learning_rate'],
                            weight_decay=CONFIG['weight_decay'])
    scheduler = fetch_scheduler(optimizer, CONFIG)
 
    train_loader, valid_loader = prepare_loaders(
        fold_train, fold_valid, CONFIG, data_transforms, num_workers=4)
 
    model, history = run_training(
        train_loader, valid_loader, model, optimizer, scheduler,
        CONFIG=CONFIG,
        model_name=f'convnext_fold{fold_n}.pth',
        tolerance_max=5,
        seed=CONFIG['seed'],
    )
 
    # Get out-of-fold predictions
    _, pf1, auroc, oof_preds, oof_targets = valid_one_epoch(
        model, valid_loader, CONFIG['device'], epoch=0,
        optimizer=optimizer, return_preds=True
    )
 
    fold_valid['oof_pred'] = oof_preds
    fold_valid['fold_n']   = fold_n
    oof_df_list.append(fold_valid)
    fold_results.append({'fold': fold_n, 'pf1': pf1, 'auroc': auroc})
    print(f"  Fold {fold_n+1} — pF1: {pf1:.4f} | AUC: {auroc:.4f}")
 
    torch.cuda.empty_cache(); gc.collect()
 
print("\n=== Cross-Validation Summary ===")
results_df = pd.DataFrame(fold_results)
print(results_df)
print(f"\nMean pF1:  {results_df.pf1.mean():.4f} ± {results_df.pf1.std():.4f}")
print(f"Mean AUC:  {results_df.auroc.mean():.4f} ± {results_df.auroc.std():.4f}")

oof_df = pd.concat(oof_df_list).reset_index(drop=True)
 
oof_pf1  = probabilistic_f1(oof_df.oof_pred.values, oof_df.target.values)
oof_auroc = roc_auc_score(oof_df.target.values, oof_df.oof_pred.values)
print(f"OOF pF1 (all folds combined): {oof_pf1:.4f}")
print(f"OOF AUC (all folds combined): {oof_auroc:.4f}")
 
# Score breakdown by fold
for fn, g in oof_df.groupby('fold_n'):
    f = probabilistic_f1(g.oof_pred.values, g.target.values)
    a = roc_auc_score(g.target.values, g.oof_pred.values)
    print(f"  Fold {fn}: pF1={f:.4f}  AUC={a:.4f}")

from sklearn.metrics import roc_curve, precision_recall_curve
 
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
 
# ROC curve
fpr, tpr, _ = roc_curve(oof_df.target.values, oof_df.oof_pred.values)
axes[0].plot(fpr, tpr, label=f'AUC={oof_auroc:.3f}')
axes[0].plot([0,1],[0,1],'--', color='gray')
axes[0].set_xlabel('FPR'); axes[0].set_ylabel('TPR')
axes[0].set_title('ROC Curve'); axes[0].legend()
 
# Precision-Recall curve
prec, rec, _ = precision_recall_curve(oof_df.target.values, oof_df.oof_pred.values)
axes[1].plot(rec, prec, color='coral')
axes[1].axhline(oof_df.target.mean(), linestyle='--', color='gray', label=f'Baseline ({oof_df.target.mean():.3f})')
axes[1].set_xlabel('Recall'); axes[1].set_ylabel('Precision')
axes[1].set_title('Precision-Recall Curve'); axes[1].legend()
 
# Prediction distribution
axes[2].hist(oof_df[oof_df.target==0].oof_pred, bins=50, alpha=0.6, label='Benign', color='steelblue')
axes[2].hist(oof_df[oof_df.target==1].oof_pred, bins=50, alpha=0.6, label='Malignant', color='red')
axes[2].set_xlabel('Predicted cancer probability')
axes[2].set_title('Score Distribution'); axes[2].legend()
 
plt.tight_layout(); plt.show()

# TODO
# Task A — Subgroup analysis:
#   Compute pF1 and AUC separately for:
#   (a) CC view vs MLO view
#   (b) Left vs Right laterality  
#   (c) Age < 55 vs Age ≥ 55
#   Are there systematic performance gaps across subgroups?
#
# Task B — Threshold optimisation:
#   Find the threshold that maximises hard-threshold F1 on the OOF predictions.
#   Is it close to 0.5 or significantly different?
#
# Task C — Ensemble the 4 folds:
#   Load all 4 saved checkpoints, run inference on the validation set,
#   and average the predictions. Does the ensemble improve over any single fold?