{"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":"none","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\"\"\"\nTwo-stage model: Image predicts ratios, then ratios+metadata predict targets\n\"\"\"\nimport torch\nimport torch.nn as nn\nimport timm\n\n\nclass TwoStageBiomassModel(nn.Module):\n    \"\"\"\n    Two-stage architecture:\n    1. Image → Ratio predictions (composition)\n    2. Ratios + Metadata → Final target predictions (absolute biomass)\n    \n    This separates:\n    - Visual composition (what the camera sees: green vs clover vs dead)\n    - Absolute biomass (affected by height, NDVI, etc.)\n    \"\"\"\n    \n    def __init__(self, \n                 backbone='vit_base_patch14_dinov2.lvd142m',\n                #  fft_backbone='convnext_tiny',\n                #  n_im_feats=2,\n                 n_targets=5,\n                 n_species=20,\n                 n_months=12,\n                 n_weather=0,\n                 pretrained=True,\n                 n_states=4,\n                 dropout=0.3,\n                 intermediate_size=100,\n                 ):\n        \"\"\"\n        Args:\n            backbone: CNN backbone for image processing\n            n_ratios: Number of ratio predictions from image (2: green_clover, total_gdm)\n            n_targets: Number of final target predictions (5 biomass values)\n            n_species: Number of species (size of one-hot vector)\n            pretrained: Use ImageNet pretrained weights\n            dropout: Dropout rate\n        \"\"\"\n        super(TwoStageBiomassModel, self).__init__()\n        \n        # ==================== STAGE 1: Image → Ratios ====================\n        self.image_backbone = timm.create_model(\n            backbone, \n            pretrained=pretrained,\n            num_classes=0,\n            global_pool='token',\n        )        \n        \n        # Get backbone feature dimension\n        with torch.no_grad():\n            dummy_input = torch.randn(1, 3, 518, 518)\n            backbone_features = self.image_backbone(dummy_input).shape[1]\n\n        self.aux_ndvi_height_head = nn.Sequential(\n            nn.Linear(backbone_features, 2),\n            nn.ReLU(),\n        )\n\n        self.aux_species_head = nn.Sequential(\n            nn.Linear(backbone_features, n_species),\n            nn.Sigmoid()\n        )\n\n        final_input_size = backbone_features + 2 + n_species\n\n        self.final_head = nn.Sequential(\n            nn.Linear(final_input_size, intermediate_size),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(intermediate_size, n_targets),\n            nn.ReLU(),\n        )\n        \n    def forward(self, inputs, return_ratios=False):\n        \"\"\"\n        Args:\n            image: (batch, 3, H, W)\n            metadata: dict with 'ndvi', 'height', 'cam_scale', 'month', 'species'\n            return_ratios: If True, also return predicted ratios\n        \n        Returns:\n            predictions: (batch, n_targets) - final biomass predictions\n            ratios (optional): (batch, n_ratios) - intermediate ratio predictions\n        \"\"\"\n        # ===== Stage 1: Image → Ratios =====\n        image = inputs['img']\n        image_features = self.image_backbone(image)\n\n        pred_species = self.aux_species_head(image_features)\n        pred_ndvi_height = self.aux_ndvi_height_head(image_features)\n        \n        # Directly combine ratios with raw metadata\n        combined = torch.cat([\n            image_features, \n            pred_ndvi_height,\n            pred_species,\n        ], dim=1)\n        final_predictions = self.final_head(combined)\n        \n        return {\n            'biomass': final_predictions,\n            'aux_ndvi_height': pred_ndvi_height,\n            'aux_species': pred_species,\n        }\n\n\n\nbackbone='vit_base_patch14_dinov2.lvd142m'\nmodel = TwoStageBiomassModel(backbone=backbone, pretrained=True)\n\n# 2. Extract only the weights for the backbone\n# This is usually what you want to reuse across different experiments\nbackbone_weights = model.image_backbone.state_dict()\n\n# 3. Save the weights to the current directory\n# This file will appear in the \"Output\" section after you 'Save & Run All'\ntorch.save(backbone_weights, f'{backbone}.pth')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-19T12:01:15.406907Z","iopub.execute_input":"2025-12-19T12:01:15.407298Z","iopub.status.idle":"2025-12-19T12:01:20.264311Z","shell.execute_reply.started":"2025-12-19T12:01:15.407273Z","shell.execute_reply":"2025-12-19T12:01:20.263437Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}