{"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":"gpu","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"},{"sourceId":272736244,"sourceType":"kernelVersion"}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# A EfficientNet-B2 Based Model Coding in Pytorch Lightning","metadata":{}},{"cell_type":"markdown","source":"## 🧠 Model Overview\n\nThis model was developed for the CSIRO competition using PyTorch Lightning.\nIt leverages a dual-branch architecture designed to process paired image halves and learn their combined representations efficiently.","metadata":{}},{"cell_type":"markdown","source":"## ⚙️ Data Preprocessing\n\nEach input image is split in half along the center, and both halves are fed into the model simultaneously.\nThis setup allows the network to capture cross-regional relationships within each sample.","metadata":{}},{"cell_type":"markdown","source":"## 🏗️ Model Architecture\n\nThe model structure follows this pipeline:\n\n2 × EfficientNet-B2 (shared weights) -> Linear -> ReLU -> Linear\n\nOnly the last block of the EfficientNet backbone is trainable, keeping the remaining layers frozen for stability and efficiency.","metadata":{}},{"cell_type":"markdown","source":"## ⚠️ Important Notice\n\nVersion 2 and Version 1 of the model have different architectures and are not compatible.\nPlease make sure to use the matching training notebook for the correct version.","metadata":{}},{"cell_type":"markdown","source":"## Code","metadata":{}},{"cell_type":"code","source":"# I would like to thank these generous sharings which make this notebook possible:\n# https://www.kaggle.com/code/easterndundrey/csiro-gold-solution\n# https://www.kaggle.com/code/easterndundrey/bad-assumptions?scriptVersionId=272451747\n# https://www.kaggle.com/code/mks2192/csiro-notebook-training\n# https://www.kaggle.com/code/mks2192/csiro-image2biomass\n\nimport timm\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom PIL import Image\nimport torchvision.transforms as T\nfrom torchvision.transforms import functional as TF\nimport pytorch_lightning as pl\nfrom pytorch_lightning.callbacks import ModelCheckpoint\nimport random\nimport os\nimport numpy as np\nimport pandas as pd\nfrom sklearn.model_selection import KFold","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-02T02:34:07.019225Z","iopub.execute_input":"2025-11-02T02:34:07.019913Z","iopub.status.idle":"2025-11-02T02:34:23.482947Z","shell.execute_reply.started":"2025-11-02T02:34:07.019888Z","shell.execute_reply":"2025-11-02T02:34:23.482257Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======== Load Datafile ========\ntest_df = pd.read_csv('/kaggle/input/csiro-biomass/test.csv')\ntest_df[\"id\"] = test_df[\"sample_id\"].str.split(\"__\").str[0]\nimage_paths_df = test_df[[\"id\",\"image_path\"]].drop_duplicates(\"image_path\").reset_index(drop=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-02T02:34:23.484380Z","iopub.execute_input":"2025-11-02T02:34:23.484837Z","iopub.status.idle":"2025-11-02T02:34:23.511054Z","shell.execute_reply.started":"2025-11-02T02:34:23.484817Z","shell.execute_reply":"2025-11-02T02:34:23.510381Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======== Seed ========\ndef seed_everything(seed: int = 114514):\n    pl.seed_everything(seed, workers=True)\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\nseed_everything(114514)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-02T02:34:23.512171Z","iopub.execute_input":"2025-11-02T02:34:23.512379Z","iopub.status.idle":"2025-11-02T02:34:23.522679Z","shell.execute_reply.started":"2025-11-02T02:34:23.512364Z","shell.execute_reply":"2025-11-02T02:34:23.522024Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======== Dataset ========\nclass InferenceDataset(Dataset):\n    def __init__(self, df, full_image_aug=None, per_half_transform=None):\n        self.df = df\n        self.full_image_aug = full_image_aug\n        self.per_half_transform = per_half_transform\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        image = Image.open('/kaggle/input/csiro-biomass/' + row[\"image_path\"]).convert(\"RGB\")\n    \n        base = self.full_image_aug(image) if self.full_image_aug is not None else image\n    \n        left = base.crop((0, 0, 1000, 1000))\n        right = base.crop((1000, 0, 2000, 1000))\n    \n        if self.per_half_transform is not None:\n            left = self.per_half_transform(left)\n            right = self.per_half_transform(right)\n        else:\n            left = F.to_tensor(left)\n            right = F.to_tensor(right)\n    \n        imgs = torch.stack([left, right], dim=0)   # (2, C, H, W)\n    \n        return imgs","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-02T02:34:23.523395Z","iopub.execute_input":"2025-11-02T02:34:23.523607Z","iopub.status.idle":"2025-11-02T02:34:23.530125Z","shell.execute_reply.started":"2025-11-02T02:34:23.523592Z","shell.execute_reply":"2025-11-02T02:34:23.529484Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======== Loss Funtion ========\nclass WeightedR2Loss(nn.Module):\n    def __init__(self, weights=None, eps=1e-8):\n        super().__init__()\n        if weights is None:\n            weights = torch.tensor([0.1, 0.1, 0.1, 0.2, 0.5])\n        self.register_buffer('weights', weights)\n        self.eps = eps\n\n    def forward(self, y_pred, y_true):\n        \"\"\"\n        y_pred: (B, 3)\n        y_true: (B, 5)\n        \"\"\"\n        DG = y_pred[:,0]\n        GDM = y_pred[:,1]\n        DT = y_pred[:,2]\n        DD = DT - GDM\n        DC = GDM - DG\n        y_hat = torch.stack([DG, DD, DC, GDM, DT], dim=1)\n        \n        y_mean = torch.mean(y_true, dim=0, keepdim=True)\n        ss_res = torch.sum((y_true - y_hat) ** 2, dim=0)\n        ss_tot = torch.sum((y_true - y_mean) ** 2, dim=0)\n        r2 = 1 - ss_res / (ss_tot + self.eps)\n\n        weighted_r2 = torch.sum(self.weights * r2)\n        loss = 1 - weighted_r2\n\n        return loss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-02T02:34:23.531504Z","iopub.execute_input":"2025-11-02T02:34:23.531921Z","iopub.status.idle":"2025-11-02T02:34:23.545004Z","shell.execute_reply.started":"2025-11-02T02:34:23.531897Z","shell.execute_reply":"2025-11-02T02:34:23.544487Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======== Model ========\nclass MultiRegressionModel(pl.LightningModule):\n    def __init__(\n        self,\n        model_name=\"efficientnet_b2\",\n        pretrained=False,\n        lr=5e-3,\n        output_dim=3,\n        hidden_dim=1536,\n        open_last_n_blocks=1,\n        train_bn=False,\n    ):\n        super().__init__()\n        self.save_hyperparameters()\n        \n        self.backbone = timm.create_model(\n            model_name, \n            pretrained=pretrained, \n            num_classes=0, \n            global_pool=\"avg\"\n        )\n        feat_dim = getattr(self.backbone, \"num_features\", None)\n        if feat_dim is None:\n            x_dummy = torch.zeros(1, 3, 1000, 1000)\n            with torch.no_grad():\n                feat_dim = self.backbone(x_dummy).shape[-1]\n        self.head = nn.Sequential(\n            nn.Linear(feat_dim * 2, hidden_dim),\n            nn.ReLU(inplace=True),\n            nn.Linear(hidden_dim, output_dim),\n        )\n\n        self.criterion = WeightedR2Loss()\n        self.val_outputs = []\n        self._freeze_backbone(open_last_n_blocks=open_last_n_blocks, train_bn=train_bn)\n\n    # Helpers\n    def _set_requires_grad(self, module, flag: bool):\n        for p in module.parameters():\n            p.requires_grad = flag\n            \n    def _freeze_backbone(self, open_last_n_blocks=1, train_bn=False):\n        self._set_requires_grad(self.backbone, False)\n\n        if hasattr(self.backbone, \"blocks\"):\n            blocks = self.backbone.blocks\n            n = len(blocks)\n            assert open_last_n_blocks >= 1 and open_last_n_blocks <= n, \\\n                f\"Expect open_last_n_blocks between [1, {n}], get {open_last_n_blocks}\"\n            for blk in blocks[n - open_last_n_blocks:]:\n                self._set_requires_grad(blk, True)\n\n        if not train_bn:\n            for m in self.backbone.modules():\n                if isinstance(m, (nn.BatchNorm2d, nn.SyncBatchNorm)):\n                    if all(not p.requires_grad for p in m.parameters()):\n                        m.eval()\n\n        self._set_requires_grad(self.head, True)\n\n\n    def forward(self, x_imgs):\n        if x_imgs.ndim == 5:\n            B, T, C, H, W = x_imgs.shape\n            assert T == 2, f\"Expect 2 pics, get T={T}\"\n            x = x_imgs.view(B * T, C, H, W)\n            feats = self.backbone(x)\n            F = feats.shape[-1]\n            feats = feats.view(B, T, F)\n            fused = torch.cat([feats[:, 0, :], feats[:, 1, :]], dim=1)\n        elif x_imgs.ndim == 4:\n            T, C, H, W = x_imgs.shape\n            assert T == 2, f\"Expect 2 pics, get T={T}\"\n            feats = self.backbone(x_imgs)\n            fused = torch.cat([feats[0], feats[1]], dim=0).unsqueeze(0)\n        else:\n            raise ValueError(\"SHAPE should be train/val-(B, 2, C, H, W) or single-sample val-(2, C, H, W)\")\n\n        y_hat = self.head(fused)\n        return y_hat\n\n    def training_step(self, batch, batch_idx):\n        x_img, y = batch\n        y_hat = self(x_img)\n        loss = self.criterion(y_hat, y)\n        self.log(\"train_loss\", loss, on_step=False, on_epoch=True)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        x_img, y = batch\n        y_hat = self(x_img)\n        loss = self.criterion(y_hat, y)\n        self.val_outputs.append((y_hat.detach().cpu(), y.detach().cpu()))\n        self.log(\"val_loss\", loss, on_step=False, on_epoch=True)\n        return loss\n\n    def on_validation_epoch_end(self):\n        y_hats = torch.cat([x[0] for x in self.val_outputs], dim=0)\n        y_trues = torch.cat([x[1] for x in self.val_outputs], dim=0)\n        self.val_outputs.clear()\n\n        DG = y_hats[:, 0]\n        GDM = y_hats[:, 1]\n        DT = y_hats[:, 2]\n        DD = DT - GDM\n        DC = GDM - DG\n        y_hat_full = torch.stack([DG, DD, DC, GDM, DT], dim=1)\n\n        y_mean = torch.mean(y_trues, dim=0, keepdim=True)\n        ss_res = torch.sum((y_trues - y_hat_full) ** 2, dim=0)\n        ss_tot = torch.sum((y_trues - y_mean) ** 2, dim=0)\n        r2 = 1 - ss_res / (ss_tot + 1e-8)\n\n        weights = self.criterion.weights.cpu()\n        weighted_r2 = torch.sum(weights * r2)\n\n        for i, name in enumerate([\"DG\", \"DD\", \"DC\", \"GDM\", \"DT\"]):\n            self.log(f\"val_r2_{name}\", r2[i], prog_bar=True, on_epoch=True)\n        self.log(\"val_weighted_r2\", weighted_r2, prog_bar=True, on_epoch=True)\n\n    def configure_optimizers(self):\n        optimizer = torch.optim.AdamW(self.parameters(), lr=self.hparams.lr)\n        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n            optimizer,\n            mode=\"min\",\n            factor=0.5,\n            patience=2,\n            threshold=0.001,\n            min_lr=1e-7,\n            verbose=False,\n        )\n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": {\n                \"scheduler\": scheduler,\n                \"monitor\": \"val_loss\",\n                \"interval\": \"epoch\",\n                \"frequency\": 1,\n            },\n        }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-02T02:34:23.545686Z","iopub.execute_input":"2025-11-02T02:34:23.545870Z","iopub.status.idle":"2025-11-02T02:34:23.570734Z","shell.execute_reply.started":"2025-11-02T02:34:23.545856Z","shell.execute_reply":"2025-11-02T02:34:23.569945Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======== Helpers ========\ndef tta_inference(model, images):\n    preds = model(images)\n    preds_lr = model(torch.flip(images, dims=[-1]))\n    preds_ud = model(torch.flip(images, dims=[-2]))\n    preds_lrud = model(torch.flip(images, dims=[-1, -2]))\n    preds_mean = 0.4 * preds + 0.3 * preds_lr + 0.2 * preds_ud + 0.1 * preds_lrud\n    return preds_mean","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-02T02:34:23.571386Z","iopub.execute_input":"2025-11-02T02:34:23.571632Z","iopub.status.idle":"2025-11-02T02:34:23.585045Z","shell.execute_reply.started":"2025-11-02T02:34:23.571612Z","shell.execute_reply":"2025-11-02T02:34:23.584258Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======== DataLoader ========\n# Transform\nimg_size = 1000\ninfer_transform = T.Compose([\n    T.Resize((img_size, img_size)),\n    T.ToTensor(),\n    T.Normalize([0.485, 0.456, 0.406],\n                [0.229, 0.224, 0.225])\n])\n\n# DataLoader\ndataset = InferenceDataset(image_paths_df, per_half_transform=infer_transform)\ndataloader = DataLoader(dataset, batch_size=8, shuffle=False, num_workers=4)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-02T02:34:23.585813Z","iopub.execute_input":"2025-11-02T02:34:23.586242Z","iopub.status.idle":"2025-11-02T02:34:23.597231Z","shell.execute_reply.started":"2025-11-02T02:34:23.586213Z","shell.execute_reply":"2025-11-02T02:34:23.596602Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======== Infer ========\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nresults_dict = {}\n\nfor fold in range(3):\n    best_ckpt = f\"/kaggle/input/csiro-naive-cnn-training/lightning_logs/version_{fold}/checkpoints/best_model_fold{fold}.ckpt\"\n    model = MultiRegressionModel.load_from_checkpoint(\n        best_ckpt,\n        model_name=\"efficientnet_b2\",\n        pretrained=False,\n        lr=1e-4,\n        output_dim=3,\n        hidden_dim=1536,\n        open_last_n_blocks=1,\n        train_bn=False,\n        map_location=device,\n    )\n    model.eval().to(device)\n\n    results = []\n    with torch.no_grad():\n        use_amp = (device.type == \"cuda\")\n        amp_dtype = torch.float16\n        for images in dataloader:\n            images = images.to(device, non_blocking=True)\n            with torch.autocast(\"cuda\", dtype=amp_dtype, enabled=use_amp):\n                preds = tta_inference(model, images)\n            results.append(preds.cpu().numpy())\n\n    results_dict[fold] = np.concatenate(results, axis=0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-02T02:34:23.597831Z","iopub.execute_input":"2025-11-02T02:34:23.598020Z","iopub.status.idle":"2025-11-02T02:34:29.660740Z","shell.execute_reply.started":"2025-11-02T02:34:23.598007Z","shell.execute_reply":"2025-11-02T02:34:29.659912Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======== Construct Submission ========\nresult_df = pd.DataFrame(np.mean([results_dict[fold] for fold in range(3)], axis=0), columns=[\"Dry_Green_g\", \"GDM_g\", \"Dry_Total_g\"])\nresult_df[\"Dry_Dead_g\"] = (result_df[\"Dry_Total_g\"] - result_df[\"GDM_g\"]).clip(lower=0)\nresult_df[\"Dry_Clover_g\"] = (result_df[\"GDM_g\"] - result_df[\"Dry_Green_g\"]).clip(lower=0)\nresult_df['sample_id'] = image_paths_df['id']\nresult_df = pd.melt(result_df, id_vars='sample_id', value_vars=[\"Dry_Green_g\", \"Dry_Dead_g\", \"Dry_Clover_g\", \"GDM_g\", \"Dry_Total_g\"], value_name='target')\nresult_df['sample_id'] = result_df['sample_id'] + '__' + result_df['variable']\nresult_df['target'] = result_df['target']\nresult_df[['sample_id', 'target']].to_csv('submission.csv', index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-02T02:34:29.662008Z","iopub.execute_input":"2025-11-02T02:34:29.662701Z","iopub.status.idle":"2025-11-02T02:34:29.687385Z","shell.execute_reply.started":"2025-11-02T02:34:29.662665Z","shell.execute_reply":"2025-11-02T02:34:29.686757Z"}},"outputs":[],"execution_count":null}]}