{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":112899,"databundleVersionId":13449579,"sourceType":"competition"},{"sourceId":13320246,"sourceType":"datasetVersion","datasetId":8444312}],"dockerImageVersionId":31090,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms, models\nfrom torchvision.models import EfficientNet_B0_Weights\nfrom PIL import Image\nfrom tqdm import tqdm\nimport gc\nfrom sklearn.model_selection import StratifiedKFold\nfrom iterstrat.ml_stratifiers import MultilabelStratifiedKFold\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# Configuration\nclass Config:\n    # Paths\n    TRAIN_CSV = '/kaggle/input/grand-xray-slam-division-a/train1.csv'\n    TRAIN_IMG_DIR = '/kaggle/input/grand-xray-slam-division-a/train1'\n    TEST_CSV = '/kaggle/input/grand-xray-slam-division-a/test1.csv'\n    TEST_IMG_DIR = '/kaggle/input/grand-xray-slam-division-a/test1'\n    \n    # Model checkpoint paths\n    CHECKPOINT_DIR = '/kaggle/working/checkpoints'\n    BEST_MODEL_PATH = '/kaggle/working/best_model_fold{}.pth'\n    EMA_MODEL_PATH = '/kaggle/working/ema_model_fold{}.pth'\n    \n    # Training parameters\n    IMG_SIZE = 512  # Increased from 224 to 512\n    BATCH_SIZE = 16  # Reduced due to larger image size\n    EPOCHS = 8\n    LEARNING_RATE = 2e-4\n    WEIGHT_DECAY = 1e-4\n    NUM_WORKERS = 4\n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    \n    # K-Fold CV\n    N_FOLDS = 5\n    TRAIN_FOLDS = [0, 1, 2, 3, 4]  # Which folds to train\n    \n    # Early stopping\n    PATIENCE = 3\n    \n    # EMA\n    EMA_DECAY = 0.999\n    \n    # Mixed precision\n    USE_AMP = True\n    \n    # Label columns for 14 thoracic conditions\n    LABELS = ['Atelectasis', 'Cardiomegaly', 'Consolidation', 'Edema', \n              'Enlarged Cardiomediastinum', 'Fracture', 'Lung Lesion', \n              'Lung Opacity', 'No Finding', 'Pleural Effusion', \n              'Pleural Other', 'Pneumonia', 'Pneumothorax', 'Support Devices']\n\n# Create checkpoint directory\nos.makedirs(Config.CHECKPOINT_DIR, exist_ok=True)\n\n# Exponential Moving Average helper\nclass EMA:\n    def __init__(self, model, decay=0.999):\n        self.model = model\n        self.decay = decay\n        self.shadow = {}\n        self.backup = {}\n        self.register()\n\n    def register(self):\n        for name, param in self.model.named_parameters():\n            if param.requires_grad:\n                self.shadow[name] = param.data.clone()\n\n    def update(self):\n        for name, param in self.model.named_parameters():\n            if param.requires_grad:\n                new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow[name]\n                self.shadow[name] = new_average.clone()\n\n    def apply_shadow(self):\n        for name, param in self.model.named_parameters():\n            if param.requires_grad:\n                self.backup[name] = param.data\n                param.data = self.shadow[name]\n\n    def restore(self):\n        for name, param in self.model.named_parameters():\n            if param.requires_grad:\n                param.data = self.backup[name]\n        self.backup = {}\n\n# Custom Dataset with proper dunder methods\nclass ChestXrayDataset(Dataset):\n    def __init__(self, dataframe, img_dir, transform=None, is_test=False):\n        self.df = dataframe.reset_index(drop=True)\n        self.img_dir = img_dir\n        self.transform = transform\n        self.is_test = is_test\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        img_name = self.df.iloc[idx]['image_id'] if 'image_id' in self.df.columns else self.df.iloc[idx][0]\n        img_path = os.path.join(self.img_dir, str(img_name))\n        \n        # Handle different image extensions\n        if not os.path.exists(img_path):\n            for ext in ['.jpg', '.jpeg', '.png', '.dcm']:\n                if os.path.exists(img_path + ext):\n                    img_path = img_path + ext\n                    break\n        \n        # Load image and ensure RGB (handle grayscale)\n        image = Image.open(img_path)\n        if image.mode != 'RGB':\n            image = image.convert('RGB')\n        \n        if self.transform:\n            image = self.transform(image)\n        \n        if self.is_test:\n            return image, img_name\n        else:\n            labels = self.df.iloc[idx][Config.LABELS].values.astype('float32')\n            return image, torch.tensor(labels, dtype=torch.float32)\n\n# Data Augmentation and Transforms\ndef get_transforms(train=True):\n    if train:\n        return transforms.Compose([\n            transforms.Resize((Config.IMG_SIZE, Config.IMG_SIZE)),\n            transforms.RandomHorizontalFlip(p=0.5),\n            transforms.RandomRotation(7),\n            transforms.ColorJitter(brightness=0.2, contrast=0.2),\n            transforms.ToTensor(),\n            transforms.Normalize(mean=[0.485, 0.456, 0.406], \n                               std=[0.229, 0.224, 0.225])\n        ])\n    else:\n        return transforms.Compose([\n            transforms.Resize((Config.IMG_SIZE, Config.IMG_SIZE)),\n            transforms.ToTensor(),\n            transforms.Normalize(mean=[0.485, 0.456, 0.406], \n                               std=[0.229, 0.224, 0.225])\n        ])\n\n# CBAM Attention Module\nclass ChannelAttention(nn.Module):\n    def __init__(self, in_channels, reduction=16):\n        super(ChannelAttention, self).__init__()\n        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n        self.max_pool = nn.AdaptiveMaxPool2d(1)\n        self.fc = nn.Sequential(\n            nn.Linear(in_channels, in_channels // reduction, bias=False),\n            nn.ReLU(inplace=True),\n            nn.Linear(in_channels // reduction, in_channels, bias=False)\n        )\n        self.sigmoid = nn.Sigmoid()\n\n    def forward(self, x):\n        b, c, _, _ = x.size()\n        avg_out = self.fc(self.avg_pool(x).view(b, c))\n        max_out = self.fc(self.max_pool(x).view(b, c))\n        out = avg_out + max_out\n        return self.sigmoid(out).view(b, c, 1, 1)\n\nclass SpatialAttention(nn.Module):\n    def __init__(self, kernel_size=7):\n        super(SpatialAttention, self).__init__()\n        self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2, bias=False)\n        self.sigmoid = nn.Sigmoid()\n\n    def forward(self, x):\n        avg_out = torch.mean(x, dim=1, keepdim=True)\n        max_out, _ = torch.max(x, dim=1, keepdim=True)\n        x = torch.cat([avg_out, max_out], dim=1)\n        x = self.conv(x)\n        return self.sigmoid(x)\n\nclass CBAM(nn.Module):\n    def __init__(self, in_channels, reduction=16, kernel_size=7):\n        super(CBAM, self).__init__()\n        self.channel_attention = ChannelAttention(in_channels, reduction)\n        self.spatial_attention = SpatialAttention(kernel_size)\n\n    def forward(self, x):\n        x = x * self.channel_attention(x)\n        x = x * self.spatial_attention(x)\n        return x\n\n# Model Definition with CBAM and Label Attention\nclass EfficientNetB0Model(nn.Module):\n    def __init__(self, num_classes=14, pretrained=True, use_attention=True):\n        super(EfficientNetB0Model, self).__init__()\n        # Load pretrained EfficientNet-B0 with updated syntax\n        if pretrained:\n            self.backbone = models.efficientnet_b0(weights=EfficientNet_B0_Weights.IMAGENET1K_V1)\n        else:\n            self.backbone = models.efficientnet_b0(weights=None)\n        \n        # Get feature dimension\n        in_features = self.backbone.classifier[1].in_features\n        \n        # Remove original classifier\n        self.backbone.classifier = nn.Identity()\n        \n        # Add CBAM attention\n        self.use_attention = use_attention\n        if use_attention:\n            self.cbam = CBAM(in_features, reduction=16)\n        \n        # Global Average Pooling\n        self.gap = nn.AdaptiveAvgPool2d(1)\n        \n        # Label attention head\n        self.label_embeddings = nn.Parameter(torch.randn(num_classes, in_features))\n        self.attention = nn.MultiheadAttention(in_features, num_heads=4, batch_first=True)\n        \n        # Final classifier\n        self.classifier = nn.Sequential(\n            nn.Dropout(p=0.3),\n            nn.Linear(in_features, num_classes)\n        )\n    \n    def forward(self, x):\n        # Extract features\n        features = self.backbone.features(x)\n        \n        # Apply CBAM if enabled\n        if self.use_attention:\n            features = self.cbam(features)\n        \n        # Global average pooling\n        pooled = self.gap(features).flatten(1)\n        \n        # Label attention\n        batch_size = pooled.size(0)\n        label_emb = self.label_embeddings.unsqueeze(0).expand(batch_size, -1, -1)\n        pooled_expanded = pooled.unsqueeze(1)\n        attended, _ = self.attention(label_emb, pooled_expanded, pooled_expanded)\n        attended = attended.mean(dim=1)\n        \n        # Combine pooled and attended features\n        combined = pooled + attended\n        \n        # Final classification\n        return self.classifier(combined)\n\n# Calculate pos_weight for class imbalance\ndef calculate_pos_weight(train_df, labels):\n    pos_counts = train_df[labels].sum(axis=0)\n    neg_counts = len(train_df) - pos_counts\n    pos_weight = neg_counts / (pos_counts + 1e-5)  # Add epsilon to avoid division by zero\n    return torch.tensor(pos_weight.values, dtype=torch.float32)\n\n# Training function with mixed precision\ndef train_epoch(model, dataloader, criterion, optimizer, scheduler, scaler, ema, device):\n    model.train()\n    running_loss = 0.0\n    \n    with tqdm(dataloader, desc='Training', leave=False) as pbar:\n        for images, labels in pbar:\n            images, labels = images.to(device), labels.to(device)\n            \n            optimizer.zero_grad()\n            \n            # Mixed precision training\n            if Config.USE_AMP:\n                with torch.cuda.amp.autocast():\n                    outputs = model(images)\n                    loss = criterion(outputs, labels)\n                scaler.scale(loss).backward()\n                scaler.step(optimizer)\n                scaler.update()\n            else:\n                outputs = model(images)\n                loss = criterion(outputs, labels)\n                loss.backward()\n                optimizer.step()\n            \n            # Update EMA\n            if ema is not None:\n                ema.update()\n            \n            # Update scheduler\n            if scheduler is not None:\n                scheduler.step()\n            \n            running_loss += loss.item() * images.size(0)\n            pbar.set_postfix({'loss': f'{loss.item():.4f}'})\n            \n            # Clear cache\n            if torch.cuda.is_available():\n                torch.cuda.empty_cache()\n    \n    epoch_loss = running_loss / len(dataloader.dataset)\n    return epoch_loss\n\n# Validation function with per-label AUC\ndef validate_epoch(model, dataloader, criterion, device):\n    from sklearn.metrics import roc_auc_score\n    \n    model.eval()\n    running_loss = 0.0\n    all_preds = []\n    all_labels = []\n    \n    with torch.no_grad():\n        with tqdm(dataloader, desc='Validation', leave=False) as pbar:\n            for images, labels in pbar:\n                images, labels = images.to(device), labels.to(device)\n                \n                if Config.USE_AMP:\n                    with torch.cuda.amp.autocast():\n                        outputs = model(images)\n                        loss = criterion(outputs, labels)\n                else:\n                    outputs = model(images)\n                    loss = criterion(outputs, labels)\n                \n                # Store predictions and labels\n                probs = torch.sigmoid(outputs).cpu().numpy()\n                all_preds.append(probs)\n                all_labels.append(labels.cpu().numpy())\n                \n                running_loss += loss.item() * images.size(0)\n                pbar.set_postfix({'loss': f'{loss.item():.4f}'})\n    \n    epoch_loss = running_loss / len(dataloader.dataset)\n    \n    # Calculate per-label and macro AUC\n    all_preds = np.vstack(all_preds)\n    all_labels = np.vstack(all_labels)\n    \n    label_aucs = []\n    for i, label_name in enumerate(Config.LABELS):\n        try:\n            auc = roc_auc_score(all_labels[:, i], all_preds[:, i])\n            label_aucs.append(auc)\n        except:\n            label_aucs.append(0.5)  # If AUC can't be calculated\n    \n    macro_auc = np.mean(label_aucs)\n    \n    return epoch_loss, macro_auc, label_aucs\n\n# Test inference with TTA\ndef generate_predictions(model, dataloader, device, tta=True):\n    model.eval()\n    predictions = []\n    image_ids = []\n    \n    with torch.no_grad():\n        for images, img_names in tqdm(dataloader, desc='Generating predictions'):\n            images = images.to(device)\n            \n            if Config.USE_AMP:\n                with torch.cuda.amp.autocast():\n                    outputs = model(images)\n            else:\n                outputs = model(images)\n            \n            probs = torch.sigmoid(outputs)\n            \n            # TTA: horizontal flip\n            if tta:\n                images_flipped = torch.flip(images, dims=[3])\n                if Config.USE_AMP:\n                    with torch.cuda.amp.autocast():\n                        outputs_flipped = model(images_flipped)\n                else:\n                    outputs_flipped = model(images_flipped)\n                probs_flipped = torch.sigmoid(outputs_flipped)\n                probs = (probs + probs_flipped) / 2\n            \n            predictions.append(probs.cpu().numpy())\n            image_ids.extend(img_names)\n    \n    predictions = np.vstack(predictions)\n    return predictions, image_ids\n\n# Main training pipeline\ndef train_fold(fold, train_df, val_df):\n    print(f\"\\n{'='*60}\")\n    print(f\"Training Fold {fold}\")\n    print(f\"{'='*60}\")\n    print(f\"Train: {len(train_df)} | Val: {len(val_df)}\")\n    \n    # Create datasets\n    train_dataset = ChestXrayDataset(train_df, Config.TRAIN_IMG_DIR, \n                                     transform=get_transforms(train=True))\n    val_dataset = ChestXrayDataset(val_df, Config.TRAIN_IMG_DIR, \n                                   transform=get_transforms(train=False))\n    \n    # Create dataloaders\n    train_loader = DataLoader(train_dataset, batch_size=Config.BATCH_SIZE, \n                             shuffle=True, num_workers=Config.NUM_WORKERS,\n                             pin_memory=True)\n    val_loader = DataLoader(val_dataset, batch_size=Config.BATCH_SIZE, \n                           shuffle=False, num_workers=Config.NUM_WORKERS,\n                           pin_memory=True)\n    \n    # Initialize model\n    model = EfficientNetB0Model(num_classes=len(Config.LABELS), pretrained=True, use_attention=True)\n    model = model.to(Config.DEVICE)\n    \n    # Initialize EMA\n    ema = EMA(model, decay=Config.EMA_DECAY)\n    \n    # Calculate pos_weight for class imbalance\n    pos_weight = calculate_pos_weight(train_df, Config.LABELS).to(Config.DEVICE)\n    criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)\n    \n    # Optimizer: AdamW\n    optimizer = optim.AdamW(model.parameters(), lr=Config.LEARNING_RATE, \n                           weight_decay=Config.WEIGHT_DECAY)\n    \n    # Scheduler: OneCycle\n    total_steps = len(train_loader) * Config.EPOCHS\n    scheduler = optim.lr_scheduler.OneCycleLR(\n        optimizer, \n        max_lr=Config.LEARNING_RATE,\n        total_steps=total_steps,\n        pct_start=0.1,\n        anneal_strategy='cos'\n    )\n    \n    # Mixed precision scaler\n    scaler = torch.cuda.amp.GradScaler() if Config.USE_AMP else None\n    \n    # Training loop\n    best_macro_auc = 0.0\n    patience_counter = 0\n    \n    for epoch in range(Config.EPOCHS):\n        print(f\"\\nEpoch [{epoch+1}/{Config.EPOCHS}]\")\n        \n        # Train\n        train_loss = train_epoch(model, train_loader, criterion, optimizer, \n                                scheduler, scaler, ema, Config.DEVICE)\n        \n        # Validate with EMA weights\n        ema.apply_shadow()\n        val_loss, macro_auc, label_aucs = validate_epoch(model, val_loader, criterion, Config.DEVICE)\n        ema.restore()\n        \n        print(f\"Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | Macro AUC: {macro_auc:.4f}\")\n        \n        # Print per-label AUCs\n        print(\"Per-label AUCs:\")\n        for i, (label, auc) in enumerate(zip(Config.LABELS, label_aucs)):\n            print(f\"  {label}: {auc:.4f}\", end=\" | \" if (i+1) % 3 != 0 else \"\\n\")\n        print()\n        \n        # Save best model based on macro AUC\n        if macro_auc > best_macro_auc:\n            best_macro_auc = macro_auc\n            patience_counter = 0\n            \n            # Save regular model\n            torch.save(model.state_dict(), Config.BEST_MODEL_PATH.format(fold))\n            \n            # Save EMA model\n            ema.apply_shadow()\n            torch.save(model.state_dict(), Config.EMA_MODEL_PATH.format(fold))\n            ema.restore()\n            \n            print(f\"✓ Best model saved! Macro AUC: {macro_auc:.4f}\")\n        else:\n            patience_counter += 1\n            print(f\"No improvement. Patience: {patience_counter}/{Config.PATIENCE}\")\n        \n        # Early stopping\n        if patience_counter >= Config.PATIENCE:\n            print(f\"\\nEarly stopping triggered after {epoch+1} epochs\")\n            break\n        \n        print(\"-\" * 60)\n        \n        # Memory cleanup\n        gc.collect()\n        if torch.cuda.is_available():\n            torch.cuda.empty_cache()\n    \n    return best_macro_auc\n\ndef main():\n    print(\"=\"*60)\n    print(\"EfficientNet-B0 with CBAM & Label Attention\")\n    print(\"Grand X-ray Slam Division A\")\n    print(\"=\"*60)\n    print(f\"\\nDevice: {Config.DEVICE}\")\n    print(f\"Image Size: {Config.IMG_SIZE}x{Config.IMG_SIZE}\")\n    print(f\"Batch Size: {Config.BATCH_SIZE}\")\n    print(f\"Learning Rate: {Config.LEARNING_RATE}\")\n    print(f\"Weight Decay: {Config.WEIGHT_DECAY}\")\n    print(f\"Epochs: {Config.EPOCHS}\")\n    print(f\"Mixed Precision: {Config.USE_AMP}\")\n    print(f\"EMA Decay: {Config.EMA_DECAY}\\n\")\n    \n    # Load training data\n    print(\"Loading training data...\")\n    train_df = pd.read_csv(Config.TRAIN_CSV)\n    print(f\"Training samples: {len(train_df)}\")\n    print(f\"Labels: {Config.LABELS}\\n\")\n    \n    # Create stratified K-folds\n    print(f\"Creating {Config.N_FOLDS}-fold MultilabelStratifiedKFold splits...\")\n    mskf = MultilabelStratifiedKFold(n_splits=Config.N_FOLDS, shuffle=True, random_state=42)\n    \n    fold_results = []\n    \n    # Train each fold\n    for fold, (train_idx, val_idx) in enumerate(mskf.split(train_df, train_df[Config.LABELS])):\n        if fold not in Config.TRAIN_FOLDS:\n            continue\n        \n        train_data = train_df.iloc[train_idx]\n        val_data = train_df.iloc[val_idx]\n        \n        best_auc = train_fold(fold, train_data, val_data)\n        fold_results.append(best_auc)\n    \n    # Print fold summary\n    print(\"\\n\" + \"=\"*60)\n    print(\"Cross-Validation Results\")\n    print(\"=\"*60)\n    for fold, auc in enumerate(fold_results):\n        print(f\"Fold {fold}: Macro AUC = {auc:.4f}\")\n    print(f\"\\nMean Macro AUC: {np.mean(fold_results):.4f} ± {np.std(fold_results):.4f}\")\n    print(\"=\"*60)\n    \n    # Generate test predictions (ensemble across folds)\n    print(\"\\nGenerating test predictions...\")\n    test_df = pd.read_csv(Config.TEST_CSV)\n    test_dataset = ChestXrayDataset(test_df, Config.TEST_IMG_DIR, \n                                    transform=get_transforms(train=False), is_test=True)\n    test_loader = DataLoader(test_dataset, batch_size=Config.BATCH_SIZE, \n                            shuffle=False, num_workers=Config.NUM_WORKERS,\n                            pin_memory=True)\n    \n    all_fold_predictions = []\n    \n    for fold in Config.TRAIN_FOLDS:\n        print(f\"Loading fold {fold} EMA model...\")\n        model = EfficientNetB0Model(num_classes=len(Config.LABELS), pretrained=False, use_attention=True)\n        model.load_state_dict(torch.load(Config.EMA_MODEL_PATH.format(fold)))\n        model = model.to(Config.DEVICE)\n        \n        predictions, image_ids = generate_predictions(model, test_loader, Config.DEVICE, tta=True)\n        all_fold_predictions.append(predictions)\n        \n        del model\n        gc.collect()\n        torch.cuda.empty_cache()\n    \n    # Average predictions across folds\n    final_predictions = np.mean(all_fold_predictions, axis=0)\n    \n    # Create submission\n    submission_df = pd.DataFrame(final_predictions, columns=Config.LABELS)\n    submission_df.insert(0, 'image_id', image_ids)\n    submission_df.to_csv('/kaggle/working/submission.csv', index=False)\n    \n    print(\"\\n✓ Submission file created: /kaggle/working/submission.csv\")\n    print(\"=\"*60)\n\nif __name__ == '__main__':\n    main()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}