{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.12"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":126118,"databundleVersionId":14933142,"sourceType":"competition"}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -qU rasterio","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-05T04:53:04.113250Z","iopub.execute_input":"2026-02-05T04:53:04.113758Z","iopub.status.idle":"2026-02-05T04:53:09.980060Z","shell.execute_reply.started":"2026-02-05T04:53:04.113709Z","shell.execute_reply":"2026-02-05T04:53:09.979190Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Import Libraris","metadata":{}},{"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\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport rasterio\nfrom rasterio.enums import Resampling\nimport timm\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom tqdm.notebook import tqdm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-05T04:53:09.982509Z","iopub.execute_input":"2026-02-05T04:53:09.982939Z","iopub.status.idle":"2026-02-05T04:53:22.481618Z","shell.execute_reply.started":"2026-02-05T04:53:09.982899Z","shell.execute_reply":"2026-02-05T04:53:22.481021Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Config\n","metadata":{}},{"cell_type":"code","source":"DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nIMG_SIZE = 256\nBATCH_SIZE = 32\nSSL_EPOCHS = 30\nSUP_EPOCHS = 20\nSSL_LR = 5e-4\nSUP_LR = 1e-4\n\nBAND_STATS = {\n    \"mean\": [\n        788.23,\n        923.93,\n        1171.00,\n        1209.15,\n        1631.89,\n        2419.85,\n        2745.96,\n        2853.41,\n        2919.35,\n        3002.29,\n        2322.61,\n        1662.19,\n    ],\n    \"std\": [\n        1139.46,\n        1188.03,\n        1163.91,\n        1203.45,\n        1195.66,\n        1185.31,\n        1276.67,\n        1270.89,\n        1254.47,\n        1419.39,\n        1082.24,\n        994.31,\n    ],\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,"execution":{"iopub.status.busy":"2026-02-05T04:53:22.482484Z","iopub.execute_input":"2026-02-05T04:53:22.483071Z","iopub.status.idle":"2026-02-05T04:53:22.535244Z","shell.execute_reply.started":"2026-02-05T04:53:22.483045Z","shell.execute_reply":"2026-02-05T04:53:22.534497Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Dataset Class\n","metadata":{}},{"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\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        folder_path, label = self.samples[idx]\n        channels = []\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(\n                        1,\n                        out_shape=(IMG_SIZE, IMG_SIZE),\n                        resampling=Resampling.bilinear,\n                    )\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\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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-05T04:53:22.536146Z","iopub.execute_input":"2026-02-05T04:53:22.536515Z","iopub.status.idle":"2026-02-05T04:53:22.553886Z","shell.execute_reply.started":"2026-02-05T04:53:22.536487Z","shell.execute_reply":"2026-02-05T04:53:22.553234Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Models (Swin & SimMIM)\n","metadata":{}},{"cell_type":"code","source":"class SentinelSwin(nn.Module):\n    def __init__(self, num_classes=4, pretrained=True):\n        super().__init__()\n        self.model = timm.create_model(\n            \"swinv2_tiny_window16_256\", pretrained=pretrained, num_classes=num_classes\n        )\n\n        original_patch_embed = self.model.patch_embed.proj\n        new_patch_embed = nn.Conv2d(\n            in_channels=12,\n            out_channels=original_patch_embed.out_channels,\n            kernel_size=original_patch_embed.kernel_size,\n            stride=original_patch_embed.stride,\n        )\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\n        self.model.patch_embed.proj = new_patch_embed\n\n    def forward(self, x):\n        return self.model(x)\n\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(\n            mask.unsqueeze(1).float(), size=(img_h, img_w), mode=\"nearest\"\n        )\n        x_masked = x * (1 - mask_expanded)\n\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\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,"execution":{"iopub.status.busy":"2026-02-05T04:53:22.554716Z","iopub.execute_input":"2026-02-05T04:53:22.554983Z","iopub.status.idle":"2026-02-05T04:53:22.564186Z","shell.execute_reply.started":"2026-02-05T04:53:22.554956Z","shell.execute_reply":"2026-02-05T04:53:22.563608Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Helper Functions\n","metadata":{}},{"cell_type":"code","source":"def generate_mask(B, ratio=0.6):\n    H = IMG_SIZE // 32\n    W = IMG_SIZE // 32\n    mask = np.zeros((B, H * W))\n    mask[:, : int(H * W * ratio)] = 1\n    for i in range(B):\n        np.random.shuffle(mask[i])\n    return torch.from_numpy(mask.reshape(B, H, W)).to(DEVICE)\n\n\ndef predict_with_tta(model, image_tensor):\n    model.eval()\n    batch_stack = []\n    batch_stack.append(image_tensor)\n    batch_stack.append(torch.flip(image_tensor, [2]))\n    batch_stack.append(torch.flip(image_tensor, [3]))\n    batch_stack.append(torch.rot90(image_tensor, 1, [2, 3]))\n    batch_stack.append(torch.rot90(image_tensor, 2, [2, 3]))\n    batch_stack.append(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 torch.argmax(avg_probs).item()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-05T04:53:22.565206Z","iopub.execute_input":"2026-02-05T04:53:22.565944Z","iopub.status.idle":"2026-02-05T04:53:22.579301Z","shell.execute_reply.started":"2026-02-05T04:53:22.565910Z","shell.execute_reply":"2026-02-05T04:53:22.578643Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## SSL Training\n","metadata":{}},{"cell_type":"code","source":"transforms = A.Compose(\n    [\n        A.HorizontalFlip(p=0.5),\n        A.VerticalFlip(p=0.5),\n        A.Rotate(limit=90, p=0.5),\n        ToTensorV2(),\n    ]\n)\n\nroot_dirs = [\"/kaggle/input/beyond-visible-spectrum-ai-for-agriculture-2026p2\"]\nssl_ds = SentinelDataset(root_dirs, mode=\"ssl\", transform=transforms)\nssl_loader = DataLoader(\n    ssl_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=2, drop_last=True\n)\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\nssl_model.train()\nfor epoch in range(SSL_EPOCHS):\n    total_loss = 0\n    pbar = tqdm(ssl_loader, desc=f\"SSL Ep {epoch+1}/{SSL_EPOCHS}\")\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(\n            mask.unsqueeze(1).float(), size=(IMG_SIZE, IMG_SIZE), mode=\"nearest\"\n        )\n        loss = (torch.abs(rec - imgs) * mask_exp).sum() / (mask_exp.sum() + 1e-6)\n\n        loss.backward()\n        ssl_opt.step()\n        total_loss += loss.item()\n        pbar.set_postfix(loss=f\"{loss.item():.4f}\")\n\ntorch.save(base_model.model.state_dict(), \"sentinel_ssl_brain.pth\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-05T04:53:22.581099Z","iopub.execute_input":"2026-02-05T04:53:22.581386Z","iopub.status.idle":"2026-02-05T05:02:15.222337Z","shell.execute_reply.started":"2026-02-05T04:53:22.581365Z","shell.execute_reply":"2026-02-05T05:02:15.221386Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Supervised Training\n","metadata":{}},{"cell_type":"code","source":"sup_ds = SentinelDataset(root_dirs, mode=\"supervised\", transform=transforms)\nsup_loader = DataLoader(sup_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)\n\nfinal_model = SentinelSwin(num_classes=NUM_CLASSES, pretrained=False).to(DEVICE)\n\ntry:\n    state_dict = torch.load(\"sentinel_ssl_brain.pth\")\n    final_model.model.load_state_dict(state_dict, strict=False)\n    print(\"SSL Weights Loaded.\")\nexcept Exception as e:\n    print(f\"SSL Load Failed: {e}\")\n\nclass_weights = torch.tensor([1.0, 4.0, 0.6, 12.0]).to(DEVICE)\ncriterion = nn.CrossEntropyLoss(weight=class_weights)\noptimizer = optim.AdamW(final_model.parameters(), lr=SUP_LR)\nscheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=SUP_EPOCHS)\n\nfinal_model.train()\nfor epoch in range(SUP_EPOCHS):\n    total_loss = 0\n    correct = 0\n    total = 0\n    pbar = tqdm(sup_loader, desc=f\"Sup Ep {epoch+1}/{SUP_EPOCHS}\")\n\n    for imgs, labels in pbar:\n        imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)\n\n        optimizer.zero_grad()\n        outputs = final_model(imgs)\n        loss = criterion(outputs, labels)\n        loss.backward()\n        optimizer.step()\n\n        total_loss += loss.item()\n        _, pred = torch.max(outputs, 1)\n        total += labels.size(0)\n        correct += (pred == labels).sum().item()\n        pbar.set_postfix(acc=f\"{100*correct/total:.1f}%\", loss=f\"{loss.item():.4f}\")\n\n    scheduler.step()\n\ntorch.save(final_model.state_dict(), \"model.pth\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-05T05:02:15.223840Z","iopub.execute_input":"2026-02-05T05:02:15.224178Z","iopub.status.idle":"2026-02-05T05:03:15.717130Z","shell.execute_reply.started":"2026-02-05T05:02:15.224146Z","shell.execute_reply":"2026-02-05T05:03:15.716401Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Submission\n","metadata":{}},{"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\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(\n                    1, out_shape=(IMG_SIZE, IMG_SIZE), resampling=Resampling.bilinear\n                )\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\n\nmodel = SentinelSwin(num_classes=4)\nmodel.load_state_dict(torch.load(\"model.pth\", map_location=DEVICE))\nmodel.to(DEVICE)\nmodel.eval()\n\nsample_ids = [\n    d for d in os.listdir(EVAL_DIR) if os.path.isdir(os.path.join(EVAL_DIR, d))\n]\nresults = []\n\nfor sample_id in tqdm(sample_ids):\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    pred_idx = predict_with_tta(model, img_tensor)\n    results.append({\"Id\": sample_id, \"Category\": ID_TO_CLASS[pred_idx]})\n\ndf = pd.DataFrame(results)\ndf = df[[\"Id\", \"Category\"]]\ndf.to_csv(OUTPUT_FILE, index=False)\nprint(\"Submission Saved.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-05T05:03:15.718467Z","iopub.execute_input":"2026-02-05T05:03:15.718821Z","iopub.status.idle":"2026-02-05T05:03:19.694321Z","shell.execute_reply.started":"2026-02-05T05:03:15.718778Z","shell.execute_reply":"2026-02-05T05:03:19.693571Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pd.read_csv(\"/kaggle/working/submission.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-05T05:03:54.138099Z","iopub.execute_input":"2026-02-05T05:03:54.138432Z","iopub.status.idle":"2026-02-05T05:03:54.162016Z","shell.execute_reply.started":"2026-02-05T05:03:54.138403Z","shell.execute_reply":"2026-02-05T05:03:54.161276Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}