{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":71549,"databundleVersionId":8561470,"sourceType":"competition"}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# RSNA 2024 Lumbar Spine — Version 6\n## Ordinal-Aware Pipeline with Anti-Overfitting\n\n### v5 Post-Mortem:\n- v5 achieved 74.5% BA (same as v4 at 74.9%)\n- Model peaked at **epoch 3**, then overfitted\n- 27.7% of Severe confused with Moderate (ordinal neighbor)\n- TTA hurt (-0.9%) — horizontal flip wrong for sagittal spine\n\n### v6 Key Changes:\n1. **CORAL Ordinal Loss** — encodes class ordering (Normal < Moderate < Severe)\n2. **Lower LR (1e-4)** + **4 epoch warmup** — v5 peaked too early\n3. **Frame Dropout** — randomly masks frames to prevent sequence overfitting\n4. **Stronger Augmentation** — CoarseDropout, wider rotation, more scale variation\n5. **Progressive Backbone Unfreezing** — freeze backbone for first 3 epochs\n6. **Fixed TTA** — no horizontal flip\n","metadata":{}},{"cell_type":"code","source":"import os\nimport copy\nimport cv2\nimport glob\nimport pydicom\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport random\nfrom tqdm import tqdm\nfrom sklearn.model_selection import StratifiedGroupKFold\nfrom sklearn.metrics import confusion_matrix, classification_report, balanced_accuracy_score\nfrom collections import Counter\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nimport torchvision.models as models\nfrom torch.utils.data import Dataset, DataLoader, WeightedRandomSampler\nfrom torch.amp import autocast, GradScaler\nfrom torch.optim.swa_utils import AveragedModel, SWALR\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"CONFIG = {\n    'seed': 42,\n    'img_size': 256,\n    'seq_length': 7,\n    'batch_size': 8,\n    'epochs': 30,\n    \n    # Lower LR — v5 peaked at epoch 3 meaning LR was too high\n    'learning_rate': 1e-4,         # Reduced from 3e-4\n    'backbone_lr': 1e-5,           # Reduced from 3e-5\n    'weight_decay': 0.03,          # Slightly reduced\n    'patience': 15,\n    'num_folds': 5,\n    'train_folds': [0],\n    \n    # Loss — ordinal, no label smoothing\n    'coral_lambda': 1.0,           # CORAL ordinal loss weight\n    'ce_weight': 0.6,              # Small CE component for stability\n    'label_smoothing': 0.0,        # REMOVED: was fighting ordinal signal\n    \n    # Training stability\n    'clip_grad_norm': 1.0,\n    'use_swa': True,\n    'swa_start_epoch': 20,\n    'swa_lr': 5e-6,\n    \n    # Architecture\n    'hidden_dim': 256,\n    'dropout': 0.4,\n    'frame_dropout': 0.15,         # NEW: randomly mask frames\n    'stochastic_depth_rate': 0.1,\n    \n    # Scheduler\n    'warmup_epochs': 3,            # Increased from 2\n    'freeze_backbone_epochs': 0,   # NEW: progressive unfreezing\n    \n    # Mixup\n    'use_mixup': True,\n    'mixup_alpha': 0.2,            # Reduced from 0.3\n    \n    'device': torch.device('cuda' if torch.cuda.is_available() else 'cpu'),\n    'target_condition': 'spinal_canal_stenosis',\n    'target_series': 'Sagittal T2/STIR'\n}\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n\nseed_everything(CONFIG['seed'])\nprint(f\"✅ Device: {CONFIG['device']}\")\nprint(f\"   Version: 6 (Ordinal-Aware + Anti-Overfitting)\")\nprint(f\"   LR: {CONFIG['learning_rate']} (head) / {CONFIG['backbone_lr']} (backbone)\")\nprint(f\"   Warmup: {CONFIG['warmup_epochs']} epochs, Backbone freeze: {CONFIG['freeze_backbone_epochs']} epochs\")\nprint(f\"   Frame dropout: {CONFIG['frame_dropout']}\")\nprint(f\"   CORAL ordinal loss + {CONFIG['ce_weight']} CE\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 1. Data Loading","metadata":{}},{"cell_type":"code","source":"DATA_ROOT = \"/kaggle/input/rsna-2024-lumbar-spine-degenerative-classification/\"\nTRAIN_IMAGES = os.path.join(DATA_ROOT, \"train_images\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_train = pd.read_csv(f\"{DATA_ROOT}/train.csv\")\ndf_coords = pd.read_csv(f\"{DATA_ROOT}/train_label_coordinates.csv\")\ndf_desc = pd.read_csv(f\"{DATA_ROOT}/train_series_descriptions.csv\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_train.columns = [col.lower().replace('/', '_') for col in df_train.columns]\ncondition_cols = [c for c in df_train.columns if c != 'study_id']\ndf_labels = pd.melt(df_train, id_vars=['study_id'], value_vars=condition_cols,\n                    var_name='condition_level', value_name='severity')\ndf_labels = df_labels.dropna(subset=['severity'])\ndf_labels['severity'] = df_labels['severity'].astype(str).str.lower().str.replace('/', '_')\n\ndef extract_meta(val):\n    parts = val.split('_')\n    level = parts[-2] + '_' + parts[-1]\n    condition = '_'.join(parts[:-2])\n    return condition, level\n\ndf_labels[['base_condition', 'level_str']] = df_labels['condition_level'].apply(lambda x: pd.Series(extract_meta(x)))\nseverity_map = {'normal_mild': 0, 'moderate': 1, 'severe': 2}\ndf_labels['label'] = df_labels['severity'].map(severity_map)\ndf_labels = df_labels.dropna(subset=['label'])\ndf_labels['label'] = df_labels['label'].astype(int)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_coords = df_coords.merge(df_desc, on=['study_id', 'series_id'], how='left')\ndf_coords['condition'] = df_coords['condition'].str.lower().str.replace(' ', '_')\ndf_coords['level'] = df_coords['level'].str.lower().str.replace('/', '_')\ndf_coords['condition_level'] = df_coords['condition'] + '_' + df_coords['level']\n\ndf_model = df_labels[df_labels['base_condition'] == CONFIG['target_condition']].copy()\ndf_coords_filt = df_coords[(df_coords['condition'] == CONFIG['target_condition']) & \n                           (df_coords['series_description'] == CONFIG['target_series'])]\n\ndf_final = df_model.merge(df_coords_filt[['study_id', 'condition_level', 'series_id', 'instance_number', 'x', 'y']],\n                          on=['study_id', 'condition_level'], how='inner')\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Filter valid files\nvalid_rows = []\nfor index, row in tqdm(df_final.iterrows(), total=len(df_final), desc=\"Checking Files\"):\n    path = f\"{TRAIN_IMAGES}/{row['study_id']}/{row['series_id']}/{int(row['instance_number'])}.dcm\"\n    if os.path.exists(path):\n        valid_rows.append(row)\n\ndf_final = pd.DataFrame(valid_rows).reset_index(drop=True)\nlevel_map = {'l1_l2': 0, 'l2_l3': 1, 'l3_l4': 2, 'l4_l5': 3, 'l5_s1': 4}\ndf_final['level_idx'] = df_final['level_str'].map(level_map)\n\nprint(f\"\\n✅ Data Ready: {len(df_final)} samples\")\nclass_counts = df_final['label'].value_counts().sort_index()\nfor i, count in enumerate(class_counts):\n    pct = count / len(df_final) * 100\n    print(f\"   Class {i}: {count} samples ({pct:.1f}%)\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. Weighted Sampler","metadata":{}},{"cell_type":"code","source":"def create_weighted_sampler(df):\n    class_counts = np.bincount(df['label'].values, minlength=3).astype(float)\n    class_weights = 1.0 / class_counts\n    sample_weights = class_weights[df['label'].values]\n    sampler = WeightedRandomSampler(\n        weights=sample_weights, num_samples=len(df), replacement=True\n    )\n    print(f\"📊 WeightedRandomSampler: counts={class_counts.astype(int).tolist()}\")\n    return sampler\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3. Dataset with Frame Dropout","metadata":{}},{"cell_type":"code","source":"class RSNADatasetV6(Dataset):\n    def __init__(self, df, seq_length=7, img_size=256, transform=None, \n                 is_training=False, frame_dropout=0.0):\n        self.df = df.reset_index(drop=True)\n        self.seq_length = seq_length\n        self.img_size = img_size\n        self.transform = transform\n        self.is_training = is_training\n        self.frame_dropout = frame_dropout\n        self.clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def load_dicom(self, path):\n        try:\n            dcm = pydicom.dcmread(path)\n            img = dcm.pixel_array.astype(np.float32)\n            \n            # DICOM windowing\n            if hasattr(dcm, 'WindowCenter') and hasattr(dcm, 'WindowWidth'):\n                wc = dcm.WindowCenter\n                ww = dcm.WindowWidth\n                if isinstance(wc, pydicom.multival.MultiValue):\n                    wc = float(wc[0])\n                else:\n                    wc = float(wc)\n                if isinstance(ww, pydicom.multival.MultiValue):\n                    ww = float(ww[0])\n                else:\n                    ww = float(ww)\n                img = np.clip((img - (wc - ww/2)) / max(ww, 1) * 255, 0, 255)\n            else:\n                if img.max() > img.min():\n                    img = (img - img.min()) / (img.max() - img.min()) * 255.0\n                else:\n                    img = np.zeros_like(img)\n            \n            img = img.astype(np.uint8)\n            img = self.clahe.apply(img)\n            return img\n        except:\n            return np.zeros((self.img_size, self.img_size), dtype=np.uint8)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        center_inst = int(row['instance_number'])\n        study_path = f\"{TRAIN_IMAGES}/{row['study_id']}/{row['series_id']}\"\n        cx, cy = int(row['x']), int(row['y'])\n        \n        # Crop jittering during training\n        if self.is_training:\n            jitter = self.img_size // 16  # ~6% jitter\n            cx += random.randint(-jitter, jitter)\n            cy += random.randint(-jitter, jitter)\n        \n        start = center_inst - (self.seq_length // 2)\n        indices = [start + i for i in range(self.seq_length)]\n        \n        # Frame dropout mask (keep center frame always)\n        if self.is_training and self.frame_dropout > 0:\n            frame_mask = [random.random() > self.frame_dropout for _ in range(self.seq_length)]\n            frame_mask[self.seq_length // 2] = True  # Always keep center\n        else:\n            frame_mask = [True] * self.seq_length\n        \n        images_list = []\n        for i, inst in enumerate(indices):\n            if not frame_mask[i]:\n                # Dropped frame — zero tensor\n                if self.transform:\n                    dummy = np.zeros((self.img_size, self.img_size, 3), dtype=np.uint8)\n                    res = self.transform(image=dummy)\n                    images_list.append(res['image'] * 0)  # Zero after normalization\n                else:\n                    images_list.append(torch.zeros(3, self.img_size, self.img_size))\n                continue\n            \n            path = os.path.join(study_path, f\"{inst}.dcm\")\n            if os.path.exists(path):\n                img = self.load_dicom(path)\n            else:\n                img = np.zeros((self.img_size, self.img_size), dtype=np.uint8)\n            \n            h, w = img.shape\n            crop_size = self.img_size // 2\n            x1 = max(0, cx - crop_size)\n            y1 = max(0, cy - crop_size)\n            x2 = min(w, cx + crop_size)\n            y2 = min(h, cy + crop_size)\n            crop = img[y1:y2, x1:x2]\n            \n            if crop.size == 0:\n                crop = np.zeros((self.img_size, self.img_size), dtype=np.uint8)\n            else:\n                crop = cv2.resize(crop, (self.img_size, self.img_size))\n            \n            crop = cv2.cvtColor(crop, cv2.COLOR_GRAY2RGB)\n            \n            if self.transform:\n                res = self.transform(image=crop)\n                images_list.append(res['image'])\n            else:\n                images_list.append(torch.tensor(crop).permute(2, 0, 1).float() / 255.0)\n            \n        sequence = torch.stack(images_list, dim=0)\n        label = torch.tensor(row['label'], dtype=torch.long)\n        level_idx = torch.tensor(row['level_idx'], dtype=torch.long)\n        \n        return sequence, label, level_idx\n\nprint(\"✅ RSNADatasetV6 with frame dropout ready\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4. Stronger Augmentation Pipeline","metadata":{}},{"cell_type":"code","source":"train_aug = A.Compose([\n    # Spatial — wider range than v5\n    A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.15, rotate_limit=15,\n                       border_mode=cv2.BORDER_CONSTANT, value=0, p=0.7),\n    \n    # Intensity — medical imaging appropriate\n    A.OneOf([\n        A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=1.0),\n        A.RandomGamma(gamma_limit=(70, 130), p=1.0),\n        A.CLAHE(clip_limit=4.0, tile_grid_size=(8, 8), p=1.0),\n    ], p=0.8),\n    \n    # Noise\n    A.OneOf([\n        A.GaussNoise(var_limit=(5.0, 40.0), p=1.0),\n        A.MultiplicativeNoise(multiplier=(0.85, 1.15), p=1.0),\n    ], p=0.4),\n    \n    # Geometric distortion\n    A.OneOf([\n        A.ElasticTransform(alpha=1, sigma=50, alpha_affine=25, p=1.0),\n        A.GridDistortion(num_steps=5, distort_limit=0.1, p=1.0),\n        A.OpticalDistortion(distort_limit=0.1, shift_limit=0.05, p=1.0),\n    ], p=0.3),\n    \n    # Dropout — forces model to use all spatial info\n    A.CoarseDropout(max_holes=6, max_height=32, max_width=32,\n                    min_holes=2, min_height=16, min_width=16,\n                    fill_value=0, p=0.4),\n    \n    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n    ToTensorV2()\n])\n\nval_aug = A.Compose([\n    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n    ToTensorV2()\n])\n\n# TTA — NO horizontal flip (wrong for sagittal spine)\ntta_augs = [\n    val_aug,  # Original\n    A.Compose([  # Slight brightness variation\n        A.RandomBrightnessContrast(brightness_limit=0.1, contrast_limit=0.1, p=1.0),\n        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n        ToTensorV2()\n    ]),\n    A.Compose([  # Slight scale\n        A.ShiftScaleRotate(shift_limit=0, scale_limit=0.05, rotate_limit=0, p=1.0),\n        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n        ToTensorV2()\n    ]),\n]\n\nprint(\"✅ Stronger augmentation pipeline\")\nprint(f\"   - CoarseDropout added (forces spatial diversity)\")\nprint(f\"   - Wider rotation: 15° (was 8°)\")\nprint(f\"   - TTA: {len(tta_augs)} augmentations (no horizontal flip)\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 5. Model Architecture\n\nSame core as v5 (AttentionPool, FiLM, BiGRU) but output layer adapted for CORAL ordinal loss.\n- CORAL outputs **K-1 = 2 logits** instead of K = 3\n- Each logit represents P(Y > k): P(Y > 0) and P(Y > 1)\n","metadata":{}},{"cell_type":"code","source":"class AttentionPool(nn.Module):\n    def __init__(self, dim):\n        super().__init__()\n        self.attn = nn.Sequential(\n            nn.Linear(dim, dim // 4),\n            nn.Tanh(),\n            nn.Linear(dim // 4, 1)\n        )\n    def forward(self, x):\n        weights = F.softmax(self.attn(x), dim=1)\n        pooled = (x * weights).sum(dim=1)\n        return pooled, weights.squeeze(-1)\n\n\nclass FiLMLayer(nn.Module):\n    def __init__(self, num_levels, feature_dim):\n        super().__init__()\n        self.gamma = nn.Embedding(num_levels, feature_dim)\n        self.beta = nn.Embedding(num_levels, feature_dim)\n        nn.init.ones_(self.gamma.weight)\n        nn.init.zeros_(self.beta.weight)\n    def forward(self, x, level_idx):\n        return self.gamma(level_idx) * x + self.beta(level_idx)\n\n\nclass SpineModelV6(nn.Module):\n    def __init__(self, num_classes=3, hidden_dim=256, gru_layers=2,\n                 dropout=0.4, num_levels=5, stochastic_depth=0.1):\n        super().__init__()\n        self.num_classes = num_classes\n        \n        effnet = models.efficientnet_v2_s(weights='IMAGENET1K_V1')\n        if stochastic_depth > 0:\n            blocks = list(effnet.features.children())\n            num_blocks = len(blocks)\n            for i, block in enumerate(blocks):\n                if hasattr(block, 'stochastic_depth'):\n                    block.stochastic_depth.p = stochastic_depth * (i / num_blocks)\n        \n        self.backbone = nn.Sequential(*list(effnet.children())[:-1])\n        self.feature_dim = 1280\n        \n        self.feature_proj = nn.Sequential(\n            nn.Dropout(dropout),\n            nn.Linear(self.feature_dim, hidden_dim),\n            nn.LayerNorm(hidden_dim),\n            nn.GELU()\n        )\n        \n        self.gru = nn.GRU(\n            input_size=hidden_dim, hidden_size=hidden_dim // 2,\n            num_layers=gru_layers, batch_first=True, bidirectional=True,\n            dropout=dropout if gru_layers > 1 else 0\n        )\n        \n        self.attn_pool = AttentionPool(hidden_dim)\n        self.film = FiLMLayer(num_levels, hidden_dim)\n        \n        # Dual output: CORAL logits + standard CE logits\n        # CORAL head: K-1 ordinal logits with shared features + rank-specific biases\n        self.ordinal_features = nn.Sequential(\n            nn.LayerNorm(hidden_dim),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, 128),\n            nn.GELU(),\n            nn.Dropout(dropout * 0.5),\n        )\n        self.ordinal_fc = nn.Linear(128, 1, bias=False)  # Shared weight\n        self.ordinal_bias = nn.Parameter(torch.tensor([-2.0, -2.95]))  # Prior: P(Y>0)≈12%, P(Y>1)≈5%  # Per-rank bias\n        \n        # CE head for stability\n        self.ce_head = nn.Linear(128, num_classes)\n        \n    def forward(self, x, level_idx=None):\n        b, s, c, h, w = x.size()\n        x = x.view(b * s, c, h, w)\n        \n        features = self.backbone(x)\n        features = features.view(b, s, -1)\n        features = self.feature_proj(features)\n        \n        gru_out, _ = self.gru(features)\n        context, attn_weights = self.attn_pool(gru_out)\n        \n        if level_idx is not None:\n            context = self.film(context, level_idx)\n        \n        # Shared feature extraction\n        shared = self.ordinal_features(context)\n        \n        # CORAL ordinal logits: shared_weight * features + per_rank_bias\n        ordinal_logits = self.ordinal_fc(shared) + self.ordinal_bias.unsqueeze(0)\n        # ordinal_logits shape: (B, K-1)\n        \n        # CE logits for stability\n        ce_logits = self.ce_head(shared)\n        \n        return {\n            'ordinal': ordinal_logits,\n            'ce': ce_logits,\n            'attention': attn_weights\n        }\n    \n    def predict_proba(self, ordinal_logits):\n        \"\"\"Convert CORAL ordinal logits to class probabilities.\"\"\"\n        cumprobs = torch.sigmoid(ordinal_logits)  # P(Y > k) for k=0,1\n        # P(Y=0) = 1 - P(Y>0)\n        # P(Y=1) = P(Y>0) - P(Y>1)\n        # P(Y=2) = P(Y>1)\n        probs = torch.zeros(cumprobs.size(0), self.num_classes, device=cumprobs.device)\n        probs[:, 0] = 1 - cumprobs[:, 0]\n        probs[:, 1] = cumprobs[:, 0] - cumprobs[:, 1]\n        probs[:, 2] = cumprobs[:, 1]\n        # Clamp to avoid negative probabilities from floating point\n        probs = probs.clamp(min=0)\n        # Renormalize\n        probs = probs / probs.sum(dim=1, keepdim=True).clamp(min=1e-8)\n        return probs\n\nprint(\"✅ SpineModelV6: CORAL ordinal output + CE head\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 6. CORAL Ordinal Loss + CE","metadata":{}},{"cell_type":"code","source":"class CoralLoss(nn.Module):\n    \"\"\"\n    CORAL (Consistent Rank Logits) ordinal loss.\n    \n    For K classes, uses K-1 binary classifiers: P(Y > k)\n    This naturally encodes the ordinal structure:\n    - Label 0 (Normal): both P(Y>0) and P(Y>1) should be low\n    - Label 1 (Moderate): P(Y>0) high, P(Y>1) low\n    - Label 2 (Severe): both P(Y>0) and P(Y>1) high\n    \n    Key advantage: predicting Normal→Severe requires TWO binary mistakes,\n    making distant ordinal errors much less likely.\n    \"\"\"\n    def __init__(self, num_classes=3):\n        super().__init__()\n        self.num_classes = num_classes\n    \n    def forward(self, ordinal_logits, labels):\n        # ordinal_logits: (B, K-1) — raw logits for P(Y > k)\n        # labels: (B,) — class indices 0, 1, 2\n        \n        # Create ordinal targets: for label=k, target is 1 for all j < k\n        # label=0: [0, 0]  label=1: [1, 0]  label=2: [1, 1]\n        batch_size = labels.size(0)\n        levels = torch.arange(self.num_classes - 1, device=labels.device)\n        targets = (labels.unsqueeze(1) > levels.unsqueeze(0)).float()\n        \n        # Binary cross-entropy for each ordinal threshold\n        loss = F.binary_cross_entropy_with_logits(ordinal_logits, targets, reduction='mean')\n        return loss\n\n\nclass CombinedOrdinalLoss(nn.Module):\n    \"\"\"CORAL + small CE for training stability.\"\"\"\n    def __init__(self, num_classes=3, ce_weight=0.3):\n        super().__init__()\n        self.coral = CoralLoss(num_classes)\n        self.ce_weight = ce_weight\n    \n    def forward(self, outputs, labels):\n        coral_loss = self.coral(outputs['ordinal'], labels)\n        ce_loss = F.cross_entropy(outputs['ce'], labels)\n        total = coral_loss + self.ce_weight * ce_loss\n        return total, {\n            'total': total.item(),\n            'coral': coral_loss.item(),\n            'ce': ce_loss.item()\n        }\n\n\ndef mixup_data(x, y, alpha=0.2):\n    if alpha > 0:\n        lam = np.random.beta(alpha, alpha)\n    else:\n        lam = 1.0\n    batch_size = x.size(0)\n    index = torch.randperm(batch_size, device=x.device)\n    mixed_x = lam * x + (1 - lam) * x[index]\n    return mixed_x, y, y[index], lam\n\nprint(\"✅ CORAL ordinal loss + CE ready\")\nprint(\"   - CORAL encodes Normal < Moderate < Severe ordering\")\nprint(\"   - CE provides gradient stability\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def compute_per_class_metrics(preds, labels, num_classes=3):\n    metrics = {}\n    for c in range(num_classes):\n        mask = (labels == c)\n        if mask.sum() > 0:\n            correct = ((preds == c) & mask).sum()\n            metrics[f'class_{c}_recall'] = correct / mask.sum()\n        else:\n            metrics[f'class_{c}_recall'] = 0.0\n    return metrics\n\nclass EarlyStopping:\n    def __init__(self, patience=5, min_delta=0.001, mode='max'):\n        self.patience = patience\n        self.min_delta = min_delta\n        self.mode = mode\n        self.counter = 0\n        self.best_score = None\n        \n    def __call__(self, val_score):\n        if self.best_score is None:\n            self.best_score = val_score\n            return False\n        improved = (val_score > self.best_score + self.min_delta) if self.mode == 'max' \\\n                   else (val_score < self.best_score - self.min_delta)\n        if improved:\n            self.best_score = val_score\n            self.counter = 0\n            return False\n        self.counter += 1\n        return self.counter >= self.patience\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 7. Training Function v6","metadata":{}},{"cell_type":"code","source":"def train_one_fold_v6(model, train_loader, val_loader, fold, config):\n    criterion = CombinedOrdinalLoss(num_classes=3, ce_weight=config['ce_weight'])\n    \n    optimizer = optim.AdamW([\n        {'params': model.backbone.parameters(), 'lr': config['backbone_lr']},\n        {'params': model.feature_proj.parameters(), 'lr': config['learning_rate']},\n        {'params': model.gru.parameters(), 'lr': config['learning_rate']},\n        {'params': model.attn_pool.parameters(), 'lr': config['learning_rate']},\n        {'params': model.film.parameters(), 'lr': config['learning_rate']},\n        {'params': model.ordinal_features.parameters(), 'lr': config['learning_rate']},\n        {'params': model.ordinal_fc.parameters(), 'lr': config['learning_rate']},\n        {'params': [model.ordinal_bias], 'lr': config['learning_rate']},\n        {'params': model.ce_head.parameters(), 'lr': config['learning_rate']},\n    ], weight_decay=config['weight_decay'])\n    \n    warmup_steps = config['warmup_epochs'] * len(train_loader)\n    total_steps = config['epochs'] * len(train_loader)\n    \n    def lr_lambda(step):\n        if step < warmup_steps:\n            return step / max(warmup_steps, 1)\n        progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)\n        return max(0.5 * (1 + np.cos(np.pi * progress)), 1e-6 / config['learning_rate'])\n    \n    scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)\n    scaler = GradScaler('cuda')\n    \n    swa_model = AveragedModel(model) if config['use_swa'] else None\n    swa_scheduler = SWALR(optimizer, swa_lr=config['swa_lr']) if config['use_swa'] else None\n    \n    early_stopping = EarlyStopping(patience=config['patience'], min_delta=0.003, mode='max')\n    \n    best_balanced_acc = 0.0\n    history = {\n        'train_loss': [], 'train_acc': [],\n        'val_loss': [], 'val_acc': [], 'balanced_acc': [],\n        'class_0_recall': [], 'class_1_recall': [], 'class_2_recall': [],\n        'coral_loss': [], 'ce_loss': []\n    }\n    \n    print(f\"\\n🚀 Training Fold {fold+1}/{config['num_folds']} (v6 — Ordinal)\")\n    print(f\"   Train: {len(train_loader.dataset)}, Val: {len(val_loader.dataset)}\")\n    print(f\"   CORAL + {config['ce_weight']}*CE, LR: {config['learning_rate']}\")\n    print(f\"   Backbone frozen for first {config['freeze_backbone_epochs']} epochs\")\n    \n    # Progressive backbone unfreezing\n    for param in model.backbone.parameters():\n        param.requires_grad = False\n    backbone_frozen = True\n    \n    for epoch in range(config['epochs']):\n        # Unfreeze backbone after N epochs\n        if backbone_frozen and epoch >= config['freeze_backbone_epochs']:\n            for param in model.backbone.parameters():\n                param.requires_grad = True\n            backbone_frozen = False\n            print(f\"   🔓 Backbone unfrozen at epoch {epoch+1}\")\n        \n        model.train()\n        train_loss = 0\n        train_correct = 0\n        train_total = 0\n        epoch_coral = 0\n        epoch_ce = 0\n        is_swa_phase = config['use_swa'] and epoch >= config['swa_start_epoch']\n        \n        status = \"[FROZEN]\" if backbone_frozen else (\"[SWA]\" if is_swa_phase else \"\")\n        loop = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{config['epochs']} {status}\")\n        \n        for images, labels, level_idx in loop:\n            images = images.to(config['device'])\n            labels = labels.to(config['device'])\n            level_idx = level_idx.to(config['device'])\n            \n            # Mixup (skip during SWA phase)\n            use_mixup = config['use_mixup'] and not is_swa_phase and random.random() < 0.5\n            if use_mixup:\n                images, labels_a, labels_b, lam = mixup_data(images, labels, config['mixup_alpha'])\n            \n            optimizer.zero_grad()\n            \n            with autocast('cuda'):\n                outputs = model(images, level_idx)\n                if use_mixup:\n                    loss_a, _ = criterion(outputs, labels_a)\n                    loss_b, _ = criterion(outputs, labels_b)\n                    loss = lam * loss_a + (1 - lam) * loss_b\n                    loss_dict = {'coral': 0, 'ce': 0}  # Skip tracking for mixup\n                else:\n                    loss, loss_dict = criterion(outputs, labels)\n                    epoch_coral += loss_dict['coral']\n                    epoch_ce += loss_dict['ce']\n            \n            scaler.scale(loss).backward()\n            scaler.unscale_(optimizer)\n            torch.nn.utils.clip_grad_norm_(model.parameters(), config['clip_grad_norm'])\n            scaler.step(optimizer)\n            scaler.update()\n            \n            if is_swa_phase:\n                swa_scheduler.step()\n            else:\n                scheduler.step()\n            \n            train_loss += loss.item()\n            \n            # Predictions from ordinal logits\n            with torch.no_grad():\n                probs = model.predict_proba(outputs['ordinal'])\n                predicted = probs.argmax(dim=1)\n            \n            if use_mixup:\n                train_correct += (lam * (predicted == labels_a).float() + \n                                  (1 - lam) * (predicted == labels_b).float()).sum().item()\n            else:\n                train_correct += (predicted == labels).sum().item()\n            train_total += labels.size(0)\n            \n            loop.set_postfix(\n                loss=f\"{train_loss/(loop.n+1):.4f}\",\n                acc=f\"{100*train_correct/train_total:.1f}%\",\n                lr=f\"{optimizer.param_groups[0]['lr']:.2e}\"\n            )\n        \n        n_batches = len(train_loader)\n        train_epoch_loss = train_loss / n_batches\n        train_acc = train_correct / train_total\n        \n        if swa_model is not None and is_swa_phase:\n            swa_model.update_parameters(model)\n        \n        # Validation\n        model.eval()\n        val_loss = 0\n        val_correct = 0\n        val_total = 0\n        all_preds = []\n        all_labels = []\n        \n        with torch.no_grad():\n            for images, labels, level_idx in val_loader:\n                images = images.to(config['device'])\n                labels = labels.to(config['device'])\n                level_idx = level_idx.to(config['device'])\n                \n                with autocast('cuda'):\n                    outputs = model(images, level_idx)\n                    loss, _ = criterion(outputs, labels)\n                \n                val_loss += loss.item()\n                probs = model.predict_proba(outputs['ordinal'])\n                predicted = probs.argmax(dim=1)\n                val_correct += (predicted == labels).sum().item()\n                val_total += labels.size(0)\n                \n                all_preds.extend(predicted.cpu().numpy())\n                all_labels.extend(labels.cpu().numpy())\n        \n        val_epoch_loss = val_loss / len(val_loader)\n        val_acc = val_correct / val_total\n        \n        all_preds = np.array(all_preds)\n        all_labels = np.array(all_labels)\n        per_class = compute_per_class_metrics(all_preds, all_labels)\n        \n        balanced_acc = (per_class['class_0_recall'] + \n                       per_class['class_1_recall'] + \n                       per_class['class_2_recall']) / 3\n        \n        history['train_loss'].append(train_epoch_loss)\n        history['train_acc'].append(train_acc)\n        history['val_loss'].append(val_epoch_loss)\n        history['val_acc'].append(val_acc)\n        history['balanced_acc'].append(balanced_acc)\n        history['class_0_recall'].append(per_class['class_0_recall'])\n        history['class_1_recall'].append(per_class['class_1_recall'])\n        history['class_2_recall'].append(per_class['class_2_recall'])\n        history['coral_loss'].append(epoch_coral / n_batches)\n        history['ce_loss'].append(epoch_ce / n_batches)\n        \n        print(f\"📊 Train Loss: {train_epoch_loss:.4f} | Train Acc: {100*train_acc:.1f}% | \"\n              f\"Val Loss: {val_epoch_loss:.4f} | Val Acc: {100*val_acc:.1f}%\")\n        print(f\"   Per-class: Normal={100*per_class['class_0_recall']:.1f}%, \"\n              f\"Moderate={100*per_class['class_1_recall']:.1f}%, \"\n              f\"Severe={100*per_class['class_2_recall']:.1f}%\")\n        print(f\"   🎯 Balanced Accuracy: {100*balanced_acc:.1f}%\"\n              f\"{' [SWA]' if is_swa_phase else ''}\"\n              f\"{' [FROZEN]' if backbone_frozen else ''}\")\n        \n        min_minority_recall = min(per_class['class_1_recall'], per_class['class_2_recall'])\n        \n        if balanced_acc > best_balanced_acc and min_minority_recall >= config.get('min_minority_recall', 0.1):\n            best_balanced_acc = balanced_acc\n            torch.save(model.state_dict(), f\"best_model_v6_fold{fold}.pth\")\n            print(f\"✅ Best Model Saved! (BA: {100*balanced_acc:.1f}%, \"\n                  f\"Min Minority: {100*min_minority_recall:.1f}%)\")\n        \n        if early_stopping(balanced_acc):\n            print(f\"⏹️ Early stopping at epoch {epoch+1}\")\n            break\n    \n    model.load_state_dict(torch.load(f\"best_model_v6_fold{fold}.pth\"))\n    return model, history, best_balanced_acc\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 8. Training","metadata":{}},{"cell_type":"code","source":"kfold = StratifiedGroupKFold(n_splits=CONFIG['num_folds'], shuffle=True, random_state=CONFIG['seed'])\nfold_results = []\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for fold, (train_idx, val_idx) in enumerate(kfold.split(df_final, df_final['label'], df_final['study_id'])):\n    if fold not in CONFIG['train_folds']:\n        continue\n    \n    print(f\"\\n{'='*60}\")\n    print(f\"FOLD {fold + 1}/{CONFIG['num_folds']} (v6 — Ordinal)\")\n    print(f\"{'='*60}\")\n    \n    train_df = df_final.iloc[train_idx].reset_index(drop=True)\n    val_df = df_final.iloc[val_idx].reset_index(drop=True)\n    \n    print(f\"\\n📊 Class Distribution:\")\n    for i in range(3):\n        count = (train_df['label'] == i).sum()\n        print(f\"   Class {i}: {count} ({100*count/len(train_df):.1f}%)\")\n    \n    sampler = create_weighted_sampler(train_df)\n    \n    train_dataset = RSNADatasetV6(\n        train_df, seq_length=CONFIG['seq_length'], img_size=CONFIG['img_size'],\n        transform=train_aug, is_training=True, \n        frame_dropout=CONFIG['frame_dropout']\n    )\n    val_dataset = RSNADatasetV6(\n        val_df, seq_length=CONFIG['seq_length'], img_size=CONFIG['img_size'],\n        transform=val_aug, is_training=False\n    )\n    \n    train_loader = DataLoader(\n        train_dataset, batch_size=CONFIG['batch_size'], sampler=sampler,\n        num_workers=2, pin_memory=True, drop_last=True\n    )\n    val_loader = DataLoader(\n        val_dataset, batch_size=CONFIG['batch_size'], shuffle=False,\n        num_workers=2, pin_memory=True\n    )\n    \n    model = SpineModelV6(\n        num_classes=3, hidden_dim=CONFIG['hidden_dim'],\n        dropout=CONFIG['dropout'], stochastic_depth=CONFIG['stochastic_depth_rate']\n    ).to(CONFIG['device'])\n    \n    param_count = sum(p.numel() for p in model.parameters())\n    print(f\"\\n🏗️  Model: SpineModelV6 ({param_count:,} params)\")\n    \n    model, history, best_balanced_acc = train_one_fold_v6(\n        model, train_loader, val_loader, fold, CONFIG\n    )\n    \n    fold_results.append({\n        'fold': fold,\n        'best_balanced_acc': best_balanced_acc,\n        'history': history\n    })\n    \n    print(f\"\\n✅ Fold {fold+1} Complete | Best BA: {100*best_balanced_acc:.1f}%\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"\\n\" + \"=\"*60)\nprint(\"TRAINING SUMMARY\")\nprint(\"=\"*60)\nfor r in fold_results:\n    print(f\"Fold {r['fold']+1}: Best BA = {100*r['best_balanced_acc']:.1f}%\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 9. Evaluation with TTA","metadata":{}},{"cell_type":"code","source":"def predict_with_tta_v6(model, dataset_df, config, tta_augs):\n    model.eval()\n    all_probs = []\n    all_labels = None\n    \n    for aug_idx, aug in enumerate(tta_augs):\n        ds = RSNADatasetV6(\n            dataset_df, seq_length=config['seq_length'], img_size=config['img_size'],\n            transform=aug, is_training=False\n        )\n        loader = DataLoader(ds, batch_size=config['batch_size'], shuffle=False, \n                          num_workers=2, pin_memory=True)\n        \n        aug_probs = []\n        if aug_idx == 0:\n            labels_list = []\n        \n        with torch.no_grad():\n            for images, labels, level_idx in loader:\n                images = images.to(config['device'])\n                level_idx = level_idx.to(config['device'])\n                \n                with autocast('cuda'):\n                    outputs = model(images, level_idx)\n                    probs = model.predict_proba(outputs['ordinal'])\n                \n                aug_probs.append(probs.cpu().numpy())\n                if aug_idx == 0:\n                    labels_list.extend(labels.numpy())\n        \n        all_probs.append(np.concatenate(aug_probs, axis=0))\n        if aug_idx == 0:\n            all_labels = np.array(labels_list)\n    \n    avg_probs = np.mean(all_probs, axis=0)\n    avg_preds = np.argmax(avg_probs, axis=1)\n    return avg_preds, all_labels, avg_probs\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Run evaluation\nmodel.eval()\ntta_preds, tta_labels, _ = predict_with_tta_v6(model, val_df, CONFIG, tta_augs)\nno_tta_preds, _, _ = predict_with_tta_v6(model, val_df, CONFIG, [val_aug])\n\npc_no = compute_per_class_metrics(no_tta_preds, tta_labels)\nba_no = np.mean([pc_no[f'class_{c}_recall'] for c in range(3)])\n\npc_tta = compute_per_class_metrics(tta_preds, tta_labels)\nba_tta = np.mean([pc_tta[f'class_{c}_recall'] for c in range(3)])\n\nprint(f\"\\n{'='*60}\")\nprint(f\"RESULTS COMPARISON\")\nprint(f\"{'='*60}\")\nprint(f\"\\nWithout TTA:  BA={100*ba_no:.1f}%  N={100*pc_no['class_0_recall']:.1f}%  \"\n      f\"M={100*pc_no['class_1_recall']:.1f}%  S={100*pc_no['class_2_recall']:.1f}%\")\nprint(f\"With TTA:     BA={100*ba_tta:.1f}%  N={100*pc_tta['class_0_recall']:.1f}%  \"\n      f\"M={100*pc_tta['class_1_recall']:.1f}%  S={100*pc_tta['class_2_recall']:.1f}%\")\nprint(f\"TTA delta:    {100*(ba_tta-ba_no):+.1f}%\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"\\n\" + \"=\"*50)\nprint(\"CLASSIFICATION REPORT\")\nprint(\"=\"*50)\nprint(classification_report(tta_labels, tta_preds,\n                           target_names=['Normal/Mild', 'Moderate', 'Severe']))\n\n# Confusion matrix\ncm = confusion_matrix(tta_labels, tta_preds)\ncm_norm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]\nplt.figure(figsize=(8, 6))\nsns.heatmap(cm_norm, annot=True, fmt='.2%', cmap='Blues',\n            xticklabels=['Normal/Mild', 'Moderate', 'Severe'],\n            yticklabels=['Normal/Mild', 'Moderate', 'Severe'])\nplt.ylabel('True Label')\nplt.xlabel('Predicted Label')\nplt.title(f'v6 Confusion Matrix (BA: {100*ba_tta:.1f}%)')\nplt.tight_layout()\nplt.show()\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Training history plots\nif fold_results:\n    h = fold_results[0]['history']\n    fig, axes = plt.subplots(1, 3, figsize=(18, 5))\n    epochs = range(1, len(h['train_loss']) + 1)\n    \n    axes[0].plot(epochs, h['train_loss'], 'b-', label='Train')\n    axes[0].plot(epochs, h['val_loss'], 'r-', label='Val')\n    axes[0].set_title('Loss'); axes[0].legend(); axes[0].grid(True, alpha=0.3)\n    \n    axes[1].plot(epochs, h['class_0_recall'], 'g-o', label='Normal', ms=3)\n    axes[1].plot(epochs, h['class_1_recall'], color='orange', marker='s', label='Moderate', ms=3)\n    axes[1].plot(epochs, h['class_2_recall'], 'r-^', label='Severe', ms=3)\n    axes[1].set_title('Per-Class Recall'); axes[1].legend(); axes[1].grid(True, alpha=0.3)\n    \n    axes[2].plot(epochs, h['balanced_acc'], 'purple', marker='d', lw=2, ms=3)\n    axes[2].set_title(f'Balanced Acc (Best: {100*max(h[\"balanced_acc\"]):.1f}%)')\n    axes[2].grid(True, alpha=0.3)\n    \n    plt.tight_layout()\n    plt.show()\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"\\n\" + \"=\"*60)\nprint(\"TRAINING COMPLETE — Version 6 (Ordinal-Aware)\")\nprint(\"=\"*60)\nprint(f\"  ✓ CORAL ordinal loss (encodes Normal < Moderate < Severe)\")\nprint(f\"  ✓ Lower LR: {CONFIG['learning_rate']} head / {CONFIG['backbone_lr']} backbone\")\nprint(f\"  ✓ 4-epoch warmup + 3-epoch backbone freeze\")\nprint(f\"  ✓ Frame dropout: {CONFIG['frame_dropout']}\")\nprint(f\"  ✓ Stronger augmentation (CoarseDropout, wider rotation)\")\nprint(f\"  ✓ Fixed TTA (no horizontal flip)\")\nprint(f\"  ✓ Gradient clipping + SWA\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}