{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":126118,"databundleVersionId":14933142}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -qU rasterio timm albumentations scikit-learn\n\nimport os\nimport math\nimport copy\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader, Subset\nfrom torch.optim.lr_scheduler import LambdaLR\nimport rasterio\nfrom rasterio.enums import Resampling\nimport timm\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom tqdm.notebook import tqdm\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.metrics import f1_score\n\n# Configuration\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nIMG_SIZE = 256\nBATCH_SIZE = 32\nSSL_EPOCHS = 30  # Increase to 100+ for best results if compute allows\nSUP_EPOCHS = 20\nSSL_LR = 5e-4\nSUP_LR = 1e-4\nN_SPLITS = 5     # For Stratified K-Fold\n\nBAND_STATS = {\n    \"mean\": [788.23, 923.93, 1171.00, 1209.15, 1631.89, 2419.85, 2745.96, 2853.41, 2919.35, 3002.29, 2322.61, 1662.19],\n    \"std\": [1139.46, 1188.03, 1163.91, 1203.45, 1195.66, 1185.31, 1276.67, 1270.89, 1254.47, 1419.39, 1082.24, 994.31],\n}\nBANDS = [\"B1\", \"B2\", \"B3\", \"B4\", \"B5\", \"B6\", \"B7\", \"B8\", \"B8A\", \"B9\", \"B11\", \"B12\"]\nCLASSES = {\"Aphid\": 0, \"Blast\": 1, \"RPH\": 2, \"Rust\": 3}\nNUM_CLASSES = 4","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class FocalLoss(nn.Module):\n    \"\"\"Multi-class Focal Loss to handle severe class imbalances.\"\"\"\n    def __init__(self, weight=None, gamma=2.0, reduction='mean'):\n        super(FocalLoss, self).__init__()\n        self.weight = weight\n        self.gamma = gamma\n        self.reduction = reduction\n\n    def forward(self, inputs, targets):\n        ce_loss = F.cross_entropy(inputs, targets, weight=self.weight, reduction='none')\n        pt = torch.exp(-ce_loss)\n        focal_loss = ((1 - pt) ** self.gamma) * ce_loss\n        if self.reduction == 'mean': return focal_loss.mean()\n        elif self.reduction == 'sum': return focal_loss.sum()\n        else: return focal_loss\n\ndef get_cosine_schedule_with_warmup(optimizer, warmup_epochs, total_epochs):\n    \"\"\"Linear warmup followed by cosine decay.\"\"\"\n    def lr_lambda(current_epoch):\n        if current_epoch < warmup_epochs:\n            return float(current_epoch) / float(max(1, warmup_epochs))\n        progress = float(current_epoch - warmup_epochs) / float(max(1, total_epochs - warmup_epochs))\n        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))\n    return LambdaLR(optimizer, lr_lambda)\n\ndef generate_mask(B, ratio=0.6):\n    \"\"\"Contiguous Block Masking for SimMIM.\"\"\"\n    H, W = IMG_SIZE // 32, IMG_SIZE // 32\n    num_patches = H * W\n    num_mask = int(num_patches * ratio)\n    masks = np.zeros((B, H, W), dtype=np.float32)\n    \n    for i in range(B):\n        mask = np.zeros((H, W), dtype=int)\n        mask_count = 0\n        while mask_count < num_mask:\n            block_w, block_h = np.random.randint(2, 5), np.random.randint(2, 5)\n            top, left = np.random.randint(0, H - block_h + 1), np.random.randint(0, W - block_w + 1)\n            mask[top:top+block_h, left:left+block_w] = 1\n            mask_count = mask.sum()\n            \n        if mask_count > num_mask:\n            masked_indices = np.argwhere(mask == 1)\n            np.random.shuffle(masked_indices)\n            excess = mask_count - num_mask\n            for j in range(excess):\n                y, x = masked_indices[j]\n                mask[y, x] = 0\n        masks[i] = mask\n    return torch.from_numpy(masks).to(DEVICE)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def mixup_data(x, y, alpha=0.2):\n    \"\"\"Returns mixed inputs, pairs of targets, and lambda\"\"\"\n    if alpha > 0:\n        lam = np.random.beta(alpha, alpha)\n    else:\n        lam = 1\n\n    batch_size = x.size()[0]\n    index = torch.randperm(batch_size).to(DEVICE)\n\n    mixed_x = lam * x + (1 - lam) * x[index, :]\n    y_a, y_b = y, y[index]\n    return mixed_x, y_a, y_b, lam\n\ndef mixup_criterion(criterion, pred, y_a, y_b, lam):\n    \"\"\"Calculates the blended loss.\"\"\"\n    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class SentinelDataset(Dataset):\n    def __init__(self, root_dirs, mode=\"ssl\", transform=None):\n        self.mode = mode\n        self.transform = transform\n        self.samples = []\n\n        for root_dir in root_dirs:\n            for root, dirs, files in os.walk(root_dir):\n                if \"B2.tif\" in files:\n                    label = -1\n                    path_parts = root.split(os.sep)\n                    for part in path_parts:\n                        if part in CLASSES:\n                            label = CLASSES[part]\n                            break\n                    if mode == \"ssl\":\n                        self.samples.append((root, 0))\n                    elif mode == \"supervised\" and label != -1:\n                        self.samples.append((root, label))\n\n    def __len__(self):\n        return len(self.samples)\n\n    def __getitem__(self, idx):\n        # Bulletproof loading loop to avoid pure-zero matrices\n        while True:\n            folder_path, label = self.samples[idx]\n            channels = []\n            success = True\n\n            for i, band in enumerate(BANDS):\n                path = os.path.join(folder_path, f\"{band}.tif\")\n                try:\n                    with rasterio.open(path) as src:\n                        img = src.read(1, out_shape=(IMG_SIZE, IMG_SIZE), resampling=Resampling.bilinear)\n                        img = img.astype(np.float32)\n                        img = (img - BAND_STATS[\"mean\"][i]) / (BAND_STATS[\"std\"][i] + 1e-6)\n                        channels.append(img)\n                except Exception:\n                    success = False\n                    break\n            \n            if success and len(channels) == len(BANDS):\n                break\n            else:\n                idx = np.random.randint(0, len(self.samples))\n\n        image = np.dstack(channels)\n\n        if self.transform:\n            augmented = self.transform(image=image)\n            image = augmented[\"image\"]\n        else:\n            image = torch.from_numpy(image.transpose(2, 0, 1))\n\n        return image, label\n\nclass SentinelSwin(nn.Module):\n    def __init__(self, num_classes=4, pretrained=True):\n        super().__init__()\n        self.model = timm.create_model(\"swinv2_tiny_window16_256\", pretrained=pretrained, num_classes=num_classes)\n        original_patch_embed = self.model.patch_embed.proj\n        new_patch_embed = nn.Conv2d(\n            in_channels=12, out_channels=original_patch_embed.out_channels,\n            kernel_size=original_patch_embed.kernel_size, stride=original_patch_embed.stride,\n        )\n        with torch.no_grad():\n            w = original_patch_embed.weight.clone()\n            new_patch_embed.weight[:, :3] = w\n            new_patch_embed.weight[:, 3:6] = w\n            new_patch_embed.weight[:, 6:9] = w\n            new_patch_embed.weight[:, 9:] = w\n            new_patch_embed.bias = original_patch_embed.bias\n        self.model.patch_embed.proj = new_patch_embed\n\n    def forward(self, x):\n        return self.model(x)\n\nclass SimMIM(nn.Module):\n    def __init__(self, encoder):\n        super().__init__()\n        self.encoder = encoder\n        self.encoder_dim = encoder.model.num_features\n        self.upscale_factor = 32\n        output_channels = 12 * (self.upscale_factor**2)\n        self.decoder = nn.Conv2d(self.encoder_dim, output_channels, kernel_size=1)\n\n    def forward(self, x, mask):\n        img_h, img_w = x.shape[2], x.shape[3]\n        mask_expanded = torch.nn.functional.interpolate(mask.unsqueeze(1).float(), size=(img_h, img_w), mode=\"nearest\")\n        x_masked = x * (1 - mask_expanded)\n        z = self.encoder.model.forward_features(x_masked)\n        if z.shape[-1] == self.encoder_dim:\n            z = z.permute(0, 3, 1, 2)\n        x_rec = self.decoder(z)\n        x_rec = torch.nn.functional.pixel_shuffle(x_rec, self.upscale_factor)\n        return x_rec","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"transforms_ssl = A.Compose([\n    A.HorizontalFlip(p=0.5),\n    A.VerticalFlip(p=0.5),\n    A.Rotate(limit=90, p=0.5),\n    ToTensorV2(),\n])\n\nroot_dirs = [\"/kaggle/input/beyond-visible-spectrum-ai-for-agriculture-2026p2\"]\n\nssl_ds = SentinelDataset(root_dirs, mode=\"ssl\", transform=transforms_ssl)\nssl_loader = DataLoader(ssl_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=2, drop_last=True)\n\nbase_model = SentinelSwin(num_classes=0, pretrained=True)\nssl_model = SimMIM(base_model).to(DEVICE)\nssl_opt = optim.AdamW(ssl_model.parameters(), lr=SSL_LR)\n\n# Apply Cosine Warmup Scheduler\nWARMUP_EPOCHS = 5 \nssl_scheduler = get_cosine_schedule_with_warmup(ssl_opt, WARMUP_EPOCHS, SSL_EPOCHS)\n\nprint(\"Starting SSL Pre-training...\")\nssl_model.train()\nfor epoch in range(SSL_EPOCHS):\n    total_loss = 0\n    current_lr = ssl_opt.param_groups[0]['lr']\n    pbar = tqdm(ssl_loader, desc=f\"SSL Ep {epoch+1}/{SSL_EPOCHS} [LR: {current_lr:.6f}]\")\n    \n    for imgs, _ in pbar:\n        imgs = imgs.to(DEVICE)\n        mask = generate_mask(imgs.shape[0])\n\n        ssl_opt.zero_grad()\n        rec = ssl_model(imgs, mask)\n\n        mask_exp = torch.nn.functional.interpolate(mask.unsqueeze(1).float(), size=(IMG_SIZE, IMG_SIZE), mode=\"nearest\")\n        loss = (torch.abs(rec - imgs) * mask_exp).sum() / (mask_exp.sum() + 1e-6)\n\n        loss.backward()\n        ssl_opt.step()\n        \n        total_loss += loss.item()\n        pbar.set_postfix(loss=f\"{loss.item():.4f}\")\n\n    ssl_scheduler.step()\n\ntorch.save(base_model.model.state_dict(), \"sentinel_ssl_brain.pth\")\nprint(\"SSL Pre-training Complete and Saved.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ==========================================\n# 7. SUPERVISED FINE-TUNING (SAFE TRANSFER MODE)\n# ==========================================\nprint(\"\\n--- Starting Supervised Fine-Tuning (Continuous LLRD + MixUp) ---\")\n\ndef mixup_data(x, y, alpha=0.2):\n    if alpha > 0: lam = np.random.beta(alpha, alpha)\n    else: lam = 1\n    batch_size = x.size()[0]\n    index = torch.randperm(batch_size).to(DEVICE)\n    mixed_x = lam * x + (1 - lam) * x[index, :]\n    y_a, y_b = y, y[index]\n    return mixed_x, y_a, y_b, lam\n\ndef mixup_criterion(criterion, pred, y_a, y_b, lam):\n    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)\n\nsup_ds = SentinelDataset(root_dirs, mode=\"supervised\", transform=transforms_ssl)\nall_labels = [label for _, label in sup_ds.samples]\nall_indices = np.arange(len(sup_ds))\nskf = StratifiedKFold(n_splits=N_SPLITS, shuffle=True, random_state=42)\nfold_f1_scores = []\n\n# Restored the balanced weights \nclass_weights = torch.tensor([1.0, 6.0, 0.4, 12.0]).to(DEVICE)\ncriterion = nn.CrossEntropyLoss(weight=class_weights).to(DEVICE)\n\nfor fold, (train_idx, val_idx) in enumerate(skf.split(all_indices, all_labels)):\n    print(f\"\\n{'='*15} FOLD {fold + 1}/{N_SPLITS} {'='*15}\")\n    \n    train_subset = Subset(sup_ds, train_idx)\n    val_subset = Subset(sup_ds, val_idx)\n    train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)\n    val_loader = DataLoader(val_subset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)\n    \n    fold_model = SentinelSwin(num_classes=NUM_CLASSES, pretrained=False).to(DEVICE)\n    try:\n        fold_model.model.load_state_dict(torch.load(\"sentinel_ssl_brain.pth\"), strict=False)\n    except Exception as e:\n        print(f\"SSL Load Failed: {e}\")\n        \n    # STAGE 1: Continuous LLRD from Epoch 1\n    backbone_params, head_params = [], []\n    for name, param in fold_model.named_parameters():\n        if \"head\" in name: head_params.append(param)\n        else: backbone_params.append(param)\n        \n    # Backbone gets 50% LR: Strong enough to learn, gentle enough to not forget SSL features\n    optimizer = optim.AdamW([\n        {'params': backbone_params, 'lr': SUP_LR * 0.5}, \n        {'params': head_params, 'lr': SUP_LR}            \n    ], weight_decay=1e-4)\n    \n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=SUP_EPOCHS)\n    \n    best_fold_f1 = 0.0\n    best_fold_weights = None\n    \n    for epoch in range(SUP_EPOCHS):\n        fold_model.train()\n        total_loss = 0\n        \n        for imgs, labels in train_loader:\n            imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)\n            \n            use_mixup = np.random.rand() < 0.5\n            if use_mixup:\n                imgs, targets_a, targets_b, lam = mixup_data(imgs, labels, alpha=0.2)\n\n            optimizer.zero_grad()\n            outputs = fold_model(imgs)\n            \n            if use_mixup:\n                loss = mixup_criterion(criterion, outputs, targets_a, targets_b, lam)\n            else:\n                loss = criterion(outputs, labels)\n            \n            loss.backward()\n            torch.nn.utils.clip_grad_norm_(fold_model.parameters(), max_norm=1.0)\n            optimizer.step()\n            total_loss += loss.item()\n            \n        scheduler.step()\n        \n        fold_model.eval()\n        val_preds, val_targets = [], []\n        with torch.no_grad():\n            for imgs, labels in val_loader:\n                imgs = imgs.to(DEVICE)\n                outputs = fold_model(imgs)\n                preds = torch.argmax(outputs, 1).cpu().numpy()\n                val_preds.extend(preds)\n                val_targets.extend(labels.numpy())\n                \n        epoch_f1 = f1_score(val_targets, val_preds, average='macro')\n        print(f\"Epoch {epoch+1}/{SUP_EPOCHS} - Loss: {total_loss/len(train_loader):.4f} - Val Macro-F1: {epoch_f1:.4f}\")\n        \n        if epoch_f1 > best_fold_f1:\n            best_fold_f1 = epoch_f1\n            best_fold_weights = copy.deepcopy(fold_model.state_dict())\n\n    print(f\"Best F1 for Fold {fold + 1}: {best_fold_f1:.4f}\")\n    fold_f1_scores.append(best_fold_f1)\n    torch.save(best_fold_weights, f\"model_fold_{fold+1}.pth\")\n\nprint(f\"\\n=== OVERALL CV MACRO-F1: {np.mean(fold_f1_scores):.4f} ===\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"EVAL_DIR = \"/kaggle/input/beyond-visible-spectrum-ai-for-agriculture-2026p2/ICPR02/kaggle/evaluation\"\nOUTPUT_FILE = \"submission.csv\"\nID_TO_CLASS = {0: \"Aphid\", 1: \"Blast\", 2: \"RPH\", 3: \"Rust\"}\n\ndef load_sample(folder_path):\n    channels = []\n    for i, band in enumerate(BANDS):\n        path = os.path.join(folder_path, f\"{band}.tif\")\n        try:\n            with rasterio.open(path) as src:\n                img = src.read(1, out_shape=(IMG_SIZE, IMG_SIZE), resampling=Resampling.bilinear)\n                img = img.astype(np.float32)\n                img = (img - BAND_STATS[\"mean\"][i]) / (BAND_STATS[\"std\"][i] + 1e-6)\n                channels.append(img)\n        except:\n            channels.append(np.zeros((IMG_SIZE, IMG_SIZE), dtype=np.float32))\n    return np.dstack(channels)\n\ndef get_tta_probs(model, image_tensor):\n    \"\"\"Returns average probabilities across 6 augmented views.\"\"\"\n    model.eval()\n    batch_stack = [\n        image_tensor,\n        torch.flip(image_tensor, [2]),\n        torch.flip(image_tensor, [3]),\n        torch.rot90(image_tensor, 1, [2, 3]),\n        torch.rot90(image_tensor, 2, [2, 3]),\n        torch.rot90(image_tensor, 3, [2, 3])\n    ]\n    batch = torch.cat(batch_stack, dim=0)\n    with torch.no_grad():\n        logits = model(batch)\n        probs = F.softmax(logits, dim=1)\n        avg_probs = torch.mean(probs, dim=0)\n    return avg_probs\n\n# Load All 5 Models into an Ensemble List\nensemble_models = []\nfor fold in range(1, N_SPLITS + 1):\n    model = SentinelSwin(num_classes=NUM_CLASSES, pretrained=False)\n    model.load_state_dict(torch.load(f\"model_fold_{fold}.pth\", map_location=DEVICE))\n    model.to(DEVICE)\n    model.eval()\n    ensemble_models.append(model)\n    \nprint(f\"Successfully loaded {len(ensemble_models)} models for ensembling.\")\n\nsample_ids = [d for d in os.listdir(EVAL_DIR) if os.path.isdir(os.path.join(EVAL_DIR, d))]\nresults = []\n\nfor sample_id in tqdm(sample_ids, desc=\"Ensemble Inference\"):\n    folder_path = os.path.join(EVAL_DIR, sample_id)\n    img_np = load_sample(folder_path)\n    img_tensor = torch.from_numpy(img_np.transpose(2, 0, 1)).unsqueeze(0).to(DEVICE)\n    \n    # Accumulate probabilities from all 5 models\n    ensemble_probs = torch.zeros(NUM_CLASSES).to(DEVICE)\n    for model in ensemble_models:\n        ensemble_probs += get_tta_probs(model, img_tensor)\n        \n    # The final prediction is the class with the highest combined probability\n    ensemble_probs /= len(ensemble_models)\n    pred_idx = torch.argmax(ensemble_probs).item()\n    \n    results.append({\"Id\": sample_id, \"Category\": ID_TO_CLASS[pred_idx]})\n\ndf = pd.DataFrame(results)[[\"Id\", \"Category\"]]\ndf.to_csv(OUTPUT_FILE, index=False)\nprint(\"Ensemble Submission Saved successfully!\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}