{"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":112509,"databundleVersionId":14254895,"sourceType":"competition"}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport pandas as pd\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport timm  # Excellent library for pretrained model backbones\nimport cv2  # OpenCV for image reading and processing\nfrom tqdm import tqdm  # Progress bar\nimport matplotlib.pyplot as plt\n\n\n# --- CENTRAL CONFIGURATION CLASS ---\n# Manage all hyperparameters and paths here\nclass CFG:\n    # --- Paths ---\n    BASE_PATH = '/kaggle/input/csiro-biomass'\n    TRAIN_CSV = os.path.join(BASE_PATH, 'train.csv')\n    IMAGE_DIR = os.path.join(BASE_PATH, 'train')\n    SEED = 42\n    \n    # --- Model Settings ---\n    MODEL_NAME = 'convnext_tiny'  # You can change to 'resnet50' or others\n    PRETRAINED = True\n    IMG_SIZE = 512  # Input image size (downscaled from 1000x1000)\n    \n    # --- Training Settings ---\n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    BATCH_SIZE = 2   # Small batch because two 768x768 images consume lots of VRAM\n    EPOCHS = 40\n    LEARNING_RATE = 3e-4      # LR for Stage 1 (head training)\n    ACCUMULATION_STEPS = 8\n    \n    # --- Fine-Tuning Settings ---\n    FREEZE_EPOCHS = 10         # Number of epochs to train only the heads\n    FINETUNE_LR = 2.5e-5        # Smaller LR for Stage 2 (full model fine-tuning)\n    \n    NUM_WORKERS = 2  # DataLoader threads\n    \n    # --- Targets & Loss Weights ---\n    TARGET_COLS = ['Dry_Total_g', 'GDM_g', 'Dry_Green_g']\n    LOSS_WEIGHTS = {\n        'total_loss': 0.5,\n        'gdm_loss': 0.2,\n        'green_loss': 0.1\n    }\n    # All 5 competition targets (for evaluation)\n    ALL_TARGET_COLS = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n    \n    # R² weights matching competition metric\n    R2_WEIGHTS = [0.1, 0.1, 0.1, 0.2, 0.5]\n# --- Print to confirm configuration ---\nprint(f\"Using device: {CFG.DEVICE}\")\nprint(f\"Model backbone: {CFG.MODEL_NAME}\")\nprint(f\"Training image size: {CFG.IMG_SIZE}x{CFG.IMG_SIZE}\")\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:40:57.115180Z","iopub.execute_input":"2025-11-03T05:40:57.115461Z","iopub.status.idle":"2025-11-03T05:41:09.252734Z","shell.execute_reply.started":"2025-11-03T05:40:57.115439Z","shell.execute_reply":"2025-11-03T05:41:09.252087Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Updates\n1. Convnext_tiny\n2. FP16 and Gradient Accumulation\n3. OOFS for ensembling","metadata":{}},{"cell_type":"code","source":"import random\ndef set_seed(seed: int = 42):\n    \"\"\"\n    Ensures full reproducibility across:\n    - Python's random module\n    - NumPy operations\n    - PyTorch (CPU & GPU)\n    - CUDA convolution algorithms\n\n    Args:\n        seed (int): The random seed to set (default = 42)\n    \"\"\"\n    # Python\n    random.seed(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n\n    # NumPy\n    np.random.seed(seed)\n\n    # PyTorch\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n\n    # Ensures deterministic behavior\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\n    print(f\"✅ Random seed set to {seed} for full reproducibility.\")\n\nset_seed(CFG.SEED)\noof_dfs = []\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:41:09.255832Z","iopub.execute_input":"2025-11-03T05:41:09.256005Z","iopub.status.idle":"2025-11-03T05:41:09.266802Z","shell.execute_reply.started":"2025-11-03T05:41:09.255991Z","shell.execute_reply":"2025-11-03T05:41:09.266276Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- Load data and perform pivot ---\n\nprint(f\"Loading {CFG.TRAIN_CSV}...\")\ntry:\n    # 1. Read the original CSV file (in long format)\n    df_long = pd.read_csv(CFG.TRAIN_CSV)\n    print(f\"The original 'long' table has {len(df_long)} rows.\")\n    \n    # 2. Perform pivot\n    # Use 'image_path' as the unique index\n    # Convert 'target_name' values into new columns\n    # Use values from the 'target' column\n    df_wide = df_long.pivot(\n        index='image_path',\n        columns='target_name',\n        values='target'\n    )\n    \n    # 3. Clean up the DataFrame\n    # After pivoting, 'image_path' becomes the index.\n    # Use reset_index() to turn it back into a normal column.\n    df_wide = df_wide.reset_index()\n    df_wide.columns.name = None  # Remove the 'target_name' label from column axis\n    \n    print(f\"Pivot complete! The new 'wide' table has {len(df_wide)} rows (one per image).\")\n    \n    # 4. Display the first 5 rows of the new wide table\n    print(\"\\n--- First 5 rows of df_wide ---\")\n    print(df_wide.head())\n    \n    # 5. (Optional) Check if our target columns are present\n    print(\"\\n--- Columns present in df_wide ---\")\n    print(df_wide.columns.tolist())\n    \nexcept FileNotFoundError:\n    print(f\"ERROR: File not found at {CFG.TRAIN_CSV}\")\n    print(\"Please verify CFG.TRAIN_CSV path.\")\n    # Create an empty placeholder df_wide so later steps don’t fail\n    df_wide = pd.DataFrame(columns=['image_path'] + CFG.TARGET_COLS)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:41:09.268056Z","iopub.execute_input":"2025-11-03T05:41:09.268382Z","iopub.status.idle":"2025-11-03T05:41:09.319130Z","shell.execute_reply.started":"2025-11-03T05:41:09.268356Z","shell.execute_reply":"2025-11-03T05:41:09.318542Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from albumentations import (\n    Compose, \n    Resize, \n    Normalize,\n    HorizontalFlip, \n    VerticalFlip,\n    RandomRotate90,  # Only rotates by 90, 180, or 270 degrees\n    ColorJitter\n)\n\n# --- Define Data Augmentation Pipelines ---\n\ndef get_train_transforms():\n    \"\"\"\n    Data augmentation pipeline for training.\n    Applied independently to img_left and img_right.\n    \"\"\"\n    return Compose([\n        # 1. Geometric augmentations\n        HorizontalFlip(p=0.5),\n        VerticalFlip(p=0.5),\n        RandomRotate90(p=0.5),  # Randomly rotate by 90, 180, or 270 degrees\n\n        # 2. Photometric (color-based) augmentations\n        ColorJitter(\n            brightness=0.2, \n            contrast=0.2, \n            saturation=0.2, \n            hue=0.1, \n            p=0.75\n        ),\n\n        # 3. Normalization (standard ImageNet mean/std)\n        Normalize(\n            mean=[0.485, 0.456, 0.406],\n            std=[0.229, 0.224, 0.225]\n        ),\n\n        # 4. Resize and convert to tensor\n        # (Resize after augmentations to retain visual quality)\n        Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n        ToTensorV2()\n    ])\n\ndef get_valid_transforms():\n    \"\"\"\n    Validation pipeline — only normalization and resizing (no augmentations).\n    \"\"\"\n    return Compose([\n        Normalize(\n            mean=[0.485, 0.456, 0.406],\n            std=[0.229, 0.224, 0.225]\n        ),\n        Resize(CFG.IMG_SIZE, CFG.IMG_SIZE),\n        ToTensorV2()\n    ])\n\n# --- Print to confirm ---\nprint(\"Augmentation functions defined successfully.\")\nprint(f\"Training images will be augmented and resized to: {CFG.IMG_SIZE}x{CFG.IMG_SIZE}\")\nprint(f\"Validation images will only be resized to: {CFG.IMG_SIZE}x{CFG.IMG_SIZE}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:41:09.319706Z","iopub.execute_input":"2025-11-03T05:41:09.319867Z","iopub.status.idle":"2025-11-03T05:41:09.326301Z","shell.execute_reply.started":"2025-11-03T05:41:09.319854Z","shell.execute_reply":"2025-11-03T05:41:09.325438Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# (This is the NEW version of the BiomassDataset class)\n\nclass BiomassDataset(Dataset):\n    \"\"\"\n    Custom dataset for the 'Two-Stream' strategy.\n\n    Returns:\n        (img_left, img_right, train_targets (3), all_targets (5))\n    \"\"\"\n    def __init__(self, df, transforms_fn, image_dir, train_target_cols, all_target_cols):\n        self.df = df\n        self.transforms_fn = transforms_fn\n        self.image_dir = image_dir\n        \n        # Pre-store for faster access\n        self.image_paths = df['image_path'].values\n        # 3 targets used for training loss computation\n        self.train_targets = df[train_target_cols].values\n        # 5 targets used for R² evaluation\n        self.all_targets = df[all_target_cols].values\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        # 1. Retrieve paths and target values\n        img_path_suffix = self.image_paths[idx]\n        train_target_vals = self.train_targets[idx]\n        all_target_vals = self.all_targets[idx]\n        \n        # 2. Read the original image (2000x1000)\n        filename = os.path.basename(img_path_suffix)\n        full_path = os.path.join(self.image_dir, filename)\n        image = cv2.imread(full_path)\n        if image is None:\n            raise FileNotFoundError(f\"Cannot read image: {full_path}\")\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        # 3. Crop into two halves (Left and Right)\n        height, width, _ = image.shape\n        mid_point = width // 2\n        img_left = image[:, :mid_point]\n        img_right = image[:, mid_point:]\n        \n        # 4. Apply augmentations independently\n        transforms = self.transforms_fn()\n        img_left_tensor = transforms(image=img_left)['image']\n        \n        # Use a NEW augmentation pipeline for the right image\n        transforms_2 = self.transforms_fn()\n        img_right_tensor = transforms_2(image=img_right)['image']\n        \n        # 5. Convert target values to tensors\n        train_target_tensor = torch.tensor(train_target_vals, dtype=torch.float32)\n        all_targets_tensor = torch.tensor(all_target_vals, dtype=torch.float32)\n        \n        # 6. Return tuple of tensors\n        return img_left_tensor, img_right_tensor, train_target_tensor, all_targets_tensor\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:41:09.913563Z","iopub.execute_input":"2025-11-03T05:41:09.913989Z","iopub.status.idle":"2025-11-03T05:41:09.920573Z","shell.execute_reply.started":"2025-11-03T05:41:09.913972Z","shell.execute_reply":"2025-11-03T05:41:09.919746Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- Test Dataset and DataLoader ---\n\nprint(\"\\nTesting Dataset & DataLoader...\")\n\ntry:\n    # 1. Create a small dataset (using training augmentations)\n    # Take the first 10 rows from df_wide for testing\n    test_df_subset = df_wide.head(10)\n    \n    train_dataset = BiomassDataset(\n        df=test_df_subset,\n        transforms_fn=get_train_transforms,  # Pass the function (not called yet)\n        image_dir=CFG.IMAGE_DIR,\n        train_target_cols=CFG.TARGET_COLS,\n        all_target_cols=CFG.ALL_TARGET_COLS\n    )\n    \n    # 2. Create a DataLoader\n    train_loader = DataLoader(\n        train_dataset,\n        batch_size=CFG.BATCH_SIZE,\n        shuffle=True,\n        num_workers=CFG.NUM_WORKERS\n    )\n    \n    # 3. Load one batch\n    print(f\"Loading 1 batch (batch_size={CFG.BATCH_SIZE})...\")\n    img_left_batch, img_right_batch, train_targets_batch, all_targets_batch = next(iter(train_loader))\n    \n    # 4. Print output shapes\n    print(\"\\n--- Output Batch Shapes ---\")\n    print(f\"  Left images:  {img_left_batch.shape}\")\n    print(f\"  Right images: {img_right_batch.shape}\")\n    print(f\"  Train targets: {train_targets_batch.shape}\")\n    print(f\"  All targets:   {all_targets_batch.shape}\")\n    \n    print(\"\\n--- Expected Shapes ---\")\n    print(f\"  Images:   [Batch, Channels, Height, Width] -> [{CFG.BATCH_SIZE}, 3, {CFG.IMG_SIZE}, {CFG.IMG_SIZE}]\")\n    print(f\"  Targets:  [Batch, Num_Train_Targets] -> [{CFG.BATCH_SIZE}, {len(CFG.TARGET_COLS)}]\")\n    \n    # (Optional) Display one sample image\n    img = img_left_batch[0].permute(1, 2, 0).cpu().numpy()\n    img = img * 0.229 + 0.485  # Undo normalization (approx.)\n    plt.imshow(np.clip(img, 0, 1))\n    plt.title(\"Sample Image from Batch (Left)\")\n    plt.axis('off')\n    plt.show()\n\nexcept Exception as e:\n    print(f\"\\nERROR while testing DataLoader: {e}\")\n    print(\"Please check CFG paths and image directory structure.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:41:13.843591Z","iopub.execute_input":"2025-11-03T05:41:13.844124Z","iopub.status.idle":"2025-11-03T05:41:14.951377Z","shell.execute_reply.started":"2025-11-03T05:41:13.844092Z","shell.execute_reply":"2025-11-03T05:41:14.950554Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class BiomassModel(nn.Module):\n    \"\"\"\n    Two-Stream, Three-Head Specialized Model.\n    \n    Architecture:\n    1. A shared backbone (e.g., ConvNeXt) processes both left and right image halves.\n    2. The same backbone is applied independently to img_left and img_right.\n    3. The two resulting feature vectors are concatenated.\n    4. The combined feature vector is fed into three separate MLP heads:\n       - Head 1 → Dry_Total_g\n       - Head 2 → GDM_g\n       - Head 3 → Dry_Green_g\n    \"\"\"\n    def __init__(self, model_name, pretrained=True, n_targets=3):\n        super(BiomassModel, self).__init__()\n        \n        # 1. Load the shared backbone\n        # num_classes=0 removes the original classification layer\n        # global_pool='avg' adds global average pooling to produce a 1D feature vector\n        self.backbone = timm.create_model(\n            model_name,\n            pretrained=pretrained,\n            num_classes=0,\n            global_pool='avg'\n        )\n        \n        # 2. Get the feature dimension from the backbone\n        # Example: convnext_tiny -> 768, resnet50 -> 2048\n        self.n_features = self.backbone.num_features\n        \n        # 3. Compute the combined feature dimension after fusion\n        # (features_left + features_right)\n        self.n_combined_features = self.n_features * 2\n        \n        # 4. Define three independent MLP heads (one for each target)\n        \n        # --- Head for Dry_Total_g ---\n        self.head_total = nn.Sequential(\n            nn.Linear(self.n_combined_features, self.n_combined_features // 2),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(self.n_combined_features // 2, 1)  # Single output value\n        )\n        \n        # --- Head for GDM_g ---\n        self.head_gdm = nn.Sequential(\n            nn.Linear(self.n_combined_features, self.n_combined_features // 2),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(self.n_combined_features // 2, 1)\n        )\n        \n        # --- Head for Dry_Green_g ---\n        self.head_green = nn.Sequential(\n            nn.Linear(self.n_combined_features, self.n_combined_features // 2),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(self.n_combined_features // 2, 1)\n        )\n\n    def forward(self, img_left, img_right):\n        # 1. Forward pass through stream 1 (Left image)\n        features_left = self.backbone(img_left)   # Shape: [batch, n_features]\n        \n        # 2. Forward pass through stream 2 (Right image)\n        features_right = self.backbone(img_right) # Shape: [batch, n_features]\n        \n        # 3. Feature fusion (concatenate left and right)\n        combined = torch.cat([features_left, features_right], dim=1)  # [batch, n_combined_features]\n        \n        # 4. Pass through three specialized MLP heads\n        out_total = self.head_total(combined)\n        out_gdm = self.head_gdm(combined)\n        out_green = self.head_green(combined)\n        \n        # 5. Return three outputs (as separate tensors for weighted loss computation)\n        return out_total, out_gdm, out_green\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:41:14.952619Z","iopub.execute_input":"2025-11-03T05:41:14.952839Z","iopub.status.idle":"2025-11-03T05:41:14.961117Z","shell.execute_reply.started":"2025-11-03T05:41:14.952819Z","shell.execute_reply":"2025-11-03T05:41:14.960469Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- Test Model Architecture ---\nprint(\"\\nTesting model architecture...\")\n\ntry:\n    # 1. Initialize the model\n    model = BiomassModel(\n        model_name=CFG.MODEL_NAME,\n        pretrained=False  # No need to load pretrained weights for architecture check\n    ).to(CFG.DEVICE)\n    \n    # 2. Create dummy input data\n    # (Batch size, Channels, Height, Width)\n    dummy_left = torch.randn(CFG.BATCH_SIZE, 3, CFG.IMG_SIZE, CFG.IMG_SIZE).to(CFG.DEVICE)\n    dummy_right = torch.randn(CFG.BATCH_SIZE, 3, CFG.IMG_SIZE, CFG.IMG_SIZE).to(CFG.DEVICE)\n    \n    print(f\"Feeding two dummy batches of shape {dummy_left.shape} into the model...\")\n    \n    # 3. Run a forward pass\n    out_total, out_gdm, out_green = model(dummy_left, dummy_right)\n    \n    # 4. Print output shapes\n    print(\"\\n--- Model Output Shapes ---\")\n    print(f\"  Output (Total): {out_total.shape}\")\n    print(f\"  Output (GDM):   {out_gdm.shape}\")\n    print(f\"  Output (Green): {out_green.shape}\")\n    \n    print(\"\\n--- Expected Shapes for Each Output ---\")\n    print(f\"  Expected: [Batch, 1] -> [{CFG.BATCH_SIZE}, 1]\")\n\nexcept Exception as e:\n    print(f\"\\nERROR while testing model: {e}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:41:14.961661Z","iopub.execute_input":"2025-11-03T05:41:14.961883Z","iopub.status.idle":"2025-11-03T05:41:16.414443Z","shell.execute_reply.started":"2025-11-03T05:41:14.961865Z","shell.execute_reply":"2025-11-03T05:41:16.413643Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class WeightedBiomassLoss(nn.Module):\n    \"\"\"\n    Weighted composite loss function for biomass prediction.\n\n    Calculates three separate SmoothL1 losses (for Total, GDM, and Green)\n    and combines them using the predefined weights from CFG.\n    \"\"\"\n    def __init__(self, loss_weights_dict):\n        super(WeightedBiomassLoss, self).__init__()\n        \n        # Use a single SmoothL1 (Huber) loss function for stability\n        self.criterion = nn.SmoothL1Loss()\n        \n        # Store weight dictionary (e.g., {'total_loss': 0.5, 'gdm_loss': 0.2, 'green_loss': 0.1})\n        self.weights = loss_weights_dict\n\n    def forward(self, predictions, targets):\n        \"\"\"\n        Args:\n            predictions: tuple of (out_total, out_gdm, out_green) from the model\n            targets: tensor of shape [batch_size, 3] containing ground truth values\n                     Columns correspond to:\n                       0 - Dry_Total_g\n                       1 - GDM_g\n                       2 - Dry_Green_g\n        Returns:\n            Weighted total loss (scalar)\n        \"\"\"\n        # 1. Unpack predictions\n        pred_total, pred_gdm, pred_green = predictions\n        \n        # 2. Extract true target components\n        true_total = targets[:, 0].unsqueeze(-1)  # [batch, 1]\n        true_gdm   = targets[:, 1].unsqueeze(-1)  # [batch, 1]\n        true_green = targets[:, 2].unsqueeze(-1)  # [batch, 1]\n        \n        # 3. Compute individual losses\n        loss_total = self.criterion(pred_total, true_total)\n        loss_gdm   = self.criterion(pred_gdm, true_gdm)\n        loss_green = self.criterion(pred_green, true_green)\n        \n        # 4. Compute weighted sum of all three losses\n        total_loss = (\n            self.weights['total_loss'] * loss_total +\n            self.weights['gdm_loss'] * loss_gdm +\n            self.weights['green_loss'] * loss_green\n        )\n        \n        return total_loss\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:41:17.352910Z","iopub.execute_input":"2025-11-03T05:41:17.353463Z","iopub.status.idle":"2025-11-03T05:41:17.358956Z","shell.execute_reply.started":"2025-11-03T05:41:17.353440Z","shell.execute_reply":"2025-11-03T05:41:17.358336Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# --- Initialize and Test the Loss Function ---\n\nprint(\"\\nInitializing loss function...\")\n\n# 1. Initialize\ncriterion = WeightedBiomassLoss(loss_weights_dict=CFG.LOSS_WEIGHTS)\ncriterion.to(CFG.DEVICE)\n\nprint(f\"Loss function created with weights: {CFG.LOSS_WEIGHTS}\")\n\n# 2. Test the loss function\nprint(\"\\nTesting loss function...\")\ntry:\n    # Create dummy prediction data (simulating model outputs)\n    dummy_preds = (\n        torch.randn(CFG.BATCH_SIZE, 1).to(CFG.DEVICE),  # out_total\n        torch.randn(CFG.BATCH_SIZE, 1).to(CFG.DEVICE),  # out_gdm\n        torch.randn(CFG.BATCH_SIZE, 1).to(CFG.DEVICE)   # out_green\n    )\n    \n    # Create dummy target data (simulating dataloader outputs)\n    dummy_targets = torch.randn(CFG.BATCH_SIZE, len(CFG.TARGET_COLS)).to(CFG.DEVICE)\n    \n    # Compute loss\n    loss_value = criterion(dummy_preds, dummy_targets)\n    \n    print(f\"  Dummy data: {CFG.BATCH_SIZE} samples\")\n    print(f\"  Computed loss value: {loss_value.item():.4f}\")\n    print(\"✓ Loss function works correctly.\")\n    \nexcept Exception as e:\n    print(f\"\\nERROR while testing loss function: {e}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:41:17.578506Z","iopub.execute_input":"2025-11-03T05:41:17.579188Z","iopub.status.idle":"2025-11-03T05:41:17.611803Z","shell.execute_reply.started":"2025-11-03T05:41:17.579162Z","shell.execute_reply":"2025-11-03T05:41:17.611228Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.model_selection import StratifiedKFold\nimport numpy as np\n\n# --- Add to the CFG class ---\nCFG.N_FOLDS = 5          # Use 5-Fold Cross-Validation\nCFG.RANDOM_STATE = 42     # For reproducibility\n\nprint(f\"\\nPreparing {CFG.N_FOLDS}-Fold Cross-Validation...\")\n\n# 1. Create a new column 'fold' initialized to -1\ndf_wide['fold'] = -1\n\n# 2. Create bins (groups) for the most important target: 'Dry_Total_g'\n# We use pd.cut to divide continuous values into discrete bins.\n# This helps StratifiedKFold maintain balanced distributions across folds.\n\nnum_bins = int(np.floor(1 + np.log2(len(df_wide))))  # Sturges' rule\nif len(df_wide) > 100:\n    num_bins = 10  # Use 10 bins if dataset is large enough\n\nprint(f\"Using {num_bins} bins to stratify based on 'Dry_Total_g'.\")\n\ndf_wide['total_bin'] = pd.cut(\n    df_wide['Dry_Total_g'], \n    bins=num_bins, \n    labels=False  # Only numerical labels are needed\n)\n\n# 3. Initialize StratifiedKFold\n# The split is based on the bin labels to preserve distribution of 'Dry_Total_g'\nskf = StratifiedKFold(\n    n_splits=CFG.N_FOLDS,\n    shuffle=True,\n    random_state=CFG.RANDOM_STATE\n)\n\n# 4. Assign fold numbers to each sample\n# skf.split returns (train_indices, valid_indices)\n# We only need valid_indices to assign fold IDs\nfor fold_num, (train_idx, valid_idx) in enumerate(skf.split(df_wide, df_wide['total_bin'])):\n    df_wide.loc[valid_idx, 'fold'] = fold_num\n\n# 5. Inspect the fold distribution\nprint(\"\\n--- Sample Count per Fold ---\")\nprint(df_wide['fold'].value_counts().sort_index())\n\nprint(\"\\n--- df_wide after adding 'fold' column (first 5 rows) ---\")\nprint(df_wide.head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:41:17.771981Z","iopub.execute_input":"2025-11-03T05:41:17.772551Z","iopub.status.idle":"2025-11-03T05:41:18.178495Z","shell.execute_reply.started":"2025-11-03T05:41:17.772531Z","shell.execute_reply":"2025-11-03T05:41:18.177756Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.metrics import r2_score\n\ndef calculate_competition_score(all_preds_3, all_targets_5):\n    \"\"\"\n    Compute the weighted R² score used in the competition.\n    \n    Args:\n        all_preds_3 (dict): Dictionary containing 3 NumPy arrays predicted by the model.\n                            Keys: 'total', 'gdm', 'green'\n        all_targets_5 (np.ndarray): Array of shape [N, 5] containing ground-truth values\n                                    for all biomass components.\n    \n    Returns:\n        float: Final weighted R² score.\n    \"\"\"\n    \n    # 1. Reconstruct all 5 predictions from the 3 model outputs\n    pred_total = all_preds_3['total']\n    pred_gdm = all_preds_3['gdm']\n    pred_green = all_preds_3['green']\n    \n    # Derived predictions based on biological relationships\n    pred_clover = np.maximum(0, pred_gdm - pred_green)\n    pred_dead = np.maximum(0, pred_total - pred_gdm)\n    \n    # Combine all 5 predictions in the exact order of CFG.ALL_TARGET_COLS\n    # ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n    y_preds = np.stack([\n        pred_green,\n        pred_dead,\n        pred_clover,\n        pred_gdm,\n        pred_total\n    ], axis=1)  # Shape: [N, 5]\n    \n    y_true = all_targets_5  # Shape: [N, 5]\n\n    # 2. Compute R² for each target\n    # sklearn's r2_score supports multi-output regression\n    r2_scores = r2_score(y_true, y_preds, multioutput='raw_values')\n    \n    # 3. Compute weighted average R² using competition weights\n    weighted_r2_total = 0.0\n    for i, weight in enumerate(CFG.R2_WEIGHTS):\n        weighted_r2_total += r2_scores[i] * weight\n        \n    return weighted_r2_total\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:41:18.202466Z","iopub.execute_input":"2025-11-03T05:41:18.203138Z","iopub.status.idle":"2025-11-03T05:41:18.208326Z","shell.execute_reply.started":"2025-11-03T05:41:18.203118Z","shell.execute_reply":"2025-11-03T05:41:18.207611Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch.optim as optim\nfrom tqdm import tqdm\nimport time\nimport gc\n\nfrom torch.cuda.amp import autocast, GradScaler\n\ndef train_one_epoch(model, loader, criterion, optimizer, device, scaler=None):\n    \"\"\"\n    Run one training epoch with FP16 mixed precision and gradient accumulation.\n    \"\"\"\n    model.train()\n    epoch_loss = 0.0\n\n    pbar = tqdm(loader, desc=\"Training\", leave=False)\n\n    optimizer.zero_grad()\n    for step, (img_left, img_right, train_targets, _all_targets_ignored) in enumerate(pbar):\n        img_left = img_left.to(device)\n        img_right = img_right.to(device)\n        targets = train_targets.to(device)\n\n        # --- FP16 Mixed Precision ---\n        with autocast(enabled=True):\n            predictions = model(img_left, img_right)\n            loss = criterion(predictions, targets)\n            loss = loss / CFG.ACCUMULATION_STEPS  # normalize loss for accumulation\n\n        # --- Scale loss for FP16 stability ---\n        scaler.scale(loss).backward()\n\n        # --- Gradient Accumulation ---\n        if (step + 1) % CFG.ACCUMULATION_STEPS == 0:\n            scaler.step(optimizer)\n            scaler.update()\n            optimizer.zero_grad()\n\n        epoch_loss += loss.item() * CFG.ACCUMULATION_STEPS  # revert scaling for reporting\n        pbar.set_postfix(loss=f'{loss.item() * CFG.ACCUMULATION_STEPS:.4f}')\n\n    return epoch_loss / len(loader)\n\n\ndef validate_one_epoch(model, loader, criterion, device):\n    \"\"\"\n    Run one validation epoch and compute weighted R² score.\n    \"\"\"\n    model.eval()\n    epoch_loss = 0.0\n    \n    # Lists to collect predictions and ground truths\n    all_preds_3 = {'total': [], 'gdm': [], 'green': []}\n    all_targets_list = []\n\n    with torch.no_grad():\n        pbar = tqdm(loader, desc=\"Validating\", leave=False)\n        \n        # The DataLoader now returns 4 values\n        for (img_left, img_right, train_targets, all_targets) in pbar:\n            img_left = img_left.to(device)\n            img_right = img_right.to(device)\n            train_targets = train_targets.to(device)  # 3 targets for loss\n            \n            # 1. Forward pass\n            pred_total, pred_gdm, pred_green = model(img_left, img_right)\n            \n            # 2. Compute loss (only based on 3 primary targets)\n            predictions_tuple = (pred_total, pred_gdm, pred_green)\n            loss = criterion(predictions_tuple, train_targets)\n            epoch_loss += loss.item()\n            \n            # 3. Collect predictions for R² calculation\n            all_preds_3['total'].append(pred_total.cpu().numpy())\n            all_preds_3['gdm'].append(pred_gdm.cpu().numpy())\n            all_preds_3['green'].append(pred_green.cpu().numpy())\n            all_targets_list.append(all_targets.cpu().numpy())\n\n    # --- End of epoch ---\n    \n    # 4. Concatenate all batches\n    preds_dict_np = {\n        'total': np.concatenate(all_preds_3['total']).flatten(),\n        'gdm':   np.concatenate(all_preds_3['gdm']).flatten(),\n        'green': np.concatenate(all_preds_3['green']).flatten()\n    }\n    targets_np_5 = np.concatenate(all_targets_list)  # Shape: [N, 5]\n    \n    # 5. Compute weighted R² competition score\n    competition_score = calculate_competition_score(preds_dict_np, targets_np_5)\n    \n    avg_epoch_loss = epoch_loss / len(loader)\n    \n    return avg_epoch_loss, competition_score\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:41:23.203745Z","iopub.execute_input":"2025-11-03T05:41:23.204327Z","iopub.status.idle":"2025-11-03T05:41:23.214448Z","shell.execute_reply.started":"2025-11-03T05:41:23.204302Z","shell.execute_reply":"2025-11-03T05:41:23.213681Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import time\nimport torch.optim as optim\nimport gc\n\n# def run_training(fold_to_run):\n#     \"\"\"\n#     Main training function implementing a 2-stage (Freeze/Unfreeze) strategy.\n#     \"\"\"\n#     print(f\"\\n{'='*50}\")\n#     print(f\"🚀 STARTING TRAINING FOR FOLD {fold_to_run} (2-Stage Strategy) 🚀\")\n#     print(f\"{'='*50}\")\n    \n#     start_time = time.time()\n    \n#     # 1. Split data for this fold\n#     print(f\"Splitting data for Fold {fold_to_run}...\")\n#     train_df = df_wide[df_wide['fold'] != fold_to_run].reset_index(drop=True)\n#     valid_df = df_wide[df_wide['fold'] == fold_to_run].reset_index(drop=True)\n    \n#     # 2. Create Datasets (BiomassDataset returns 4 outputs)\n#     train_dataset = BiomassDataset(\n#         df=train_df, transforms_fn=get_train_transforms, image_dir=CFG.IMAGE_DIR,\n#         train_target_cols=CFG.TARGET_COLS, all_target_cols=CFG.ALL_TARGET_COLS\n#     )\n#     valid_dataset = BiomassDataset(\n#         df=valid_df, transforms_fn=get_valid_transforms, image_dir=CFG.IMAGE_DIR,\n#         train_target_cols=CFG.TARGET_COLS, all_target_cols=CFG.ALL_TARGET_COLS\n#     )\n    \n#     # 3. Create DataLoaders\n#     train_loader = DataLoader(\n#         train_dataset, batch_size=CFG.BATCH_SIZE, shuffle=True,\n#         num_workers=CFG.NUM_WORKERS, pin_memory=True\n#     )\n#     valid_loader = DataLoader(\n#         valid_dataset, batch_size=CFG.BATCH_SIZE * 2, shuffle=False,\n#         num_workers=CFG.NUM_WORKERS, pin_memory=True\n#     )\n    \n#     # 4. Initialize Model & Loss (with Multi-GPU support)\n#     print(f\"Loading backbone '{CFG.MODEL_NAME}'...\")\n#     model_base = BiomassModel(CFG.MODEL_NAME, CFG.PRETRAINED)\n    \n#     if torch.cuda.device_count() > 1:\n#         print(f\"Using {torch.cuda.device_count()} GPUs with nn.DataParallel.\")\n#         model = nn.DataParallel(model_base)\n#     else:\n#         model = model_base\n#     model.to(CFG.DEVICE)\n    \n#     criterion = WeightedBiomassLoss(CFG.LOSS_WEIGHTS).to(CFG.DEVICE)\n    \n#     # =================================================================\n#     # ✨ STAGE 1: Train \"Heads\" (Freeze Backbone)\n#     # =================================================================\n#     print(f\"\\n--- STAGE 1: Freezing Backbone (Training Heads Only) ---\")\n#     print(f\"Epochs: 1 to {CFG.FREEZE_EPOCHS} | LR: {CFG.LEARNING_RATE}\")\n    \n#     # Freeze all backbone parameters\n#     for param in model.module.backbone.parameters() if isinstance(model, nn.DataParallel) else model.backbone.parameters():\n#         param.requires_grad = False\n        \n#     # Optimizer for only the heads\n#     optimizer = optim.Adam(\n#         filter(lambda p: p.requires_grad, model.parameters()), \n#         lr=CFG.LEARNING_RATE\n#     )\n#     scheduler = optim.lr_scheduler.ReduceLROnPlateau(\n#         optimizer, mode='min', factor=0.1, patience=2\n#     )\n    \n#     best_score = -float('inf')  # Track best R² score across both stages\n\n#     for epoch in range(1, CFG.FREEZE_EPOCHS + 1):\n#         print(f\"\\n--- Epoch {epoch}/{CFG.EPOCHS} (Stage 1) ---\")\n        \n#         train_loss = train_one_epoch(model, train_loader, criterion, optimizer, CFG.DEVICE)\n#         valid_loss, competition_score = validate_one_epoch(model, valid_loader, criterion, CFG.DEVICE)\n        \n#         scheduler.step(valid_loss)\n        \n#         print(f\"Epoch {epoch} - Train Loss: {train_loss:.4f} | Valid Loss: {valid_loss:.4f} | R² Score: {competition_score:.4f}\")\n        \n#         if competition_score > best_score:\n#             best_score = competition_score\n#             print(f\"✨ R² improved! Saving model as 'best_model_fold{fold_to_run}.pth'...\")\n#             torch.save(\n#                 model.module.state_dict() if isinstance(model, nn.DataParallel) else model.state_dict(),\n#                 f'best_model_fold{fold_to_run}.pth'\n#             )\n\n#     # =================================================================\n#     # ✨ STAGE 2: Fine-tuning (Unfreeze Backbone)\n#     # =================================================================\n#     print(f\"\\n--- STAGE 2: Unfreezing Backbone (Fine-tuning) ---\")\n#     print(f\"Epochs: {CFG.FREEZE_EPOCHS + 1} to {CFG.EPOCHS} | LR: {CFG.FINETUNE_LR}\")\n\n#     # Unfreeze all model parameters\n#     for param in model.module.backbone.parameters() if isinstance(model, nn.DataParallel) else model.backbone.parameters():\n#         param.requires_grad = True\n        \n#     # New optimizer for all model parameters (lower LR)\n#     optimizer = optim.Adam(model.parameters(), lr=CFG.FINETUNE_LR)\n#     scheduler = optim.lr_scheduler.ReduceLROnPlateau(\n#         optimizer, mode='min', factor=0.2, patience=3\n#     )\n    \n#     # Continue training through remaining epochs\n#     for epoch in range(CFG.FREEZE_EPOCHS + 1, CFG.EPOCHS + 1):\n#         print(f\"\\n--- Epoch {epoch}/{CFG.EPOCHS} (Stage 2) ---\")\n        \n#         train_loss = train_one_epoch(model, train_loader, criterion, optimizer, CFG.DEVICE)\n#         valid_loss, competition_score = validate_one_epoch(model, valid_loader, criterion, CFG.DEVICE)\n        \n#         scheduler.step(valid_loss)\n        \n#         print(f\"Epoch {epoch} - Train Loss: {train_loss:.4f} | Valid Loss: {valid_loss:.4f} | R² Score: {competition_score:.4f}\")\n        \n#         if competition_score > best_score:\n#             best_score = competition_score\n#             print(f\"✨ R² improved! Saving model as 'best_model_fold{fold_to_run}.pth'...\")\n#             torch.save(\n#                 model.module.state_dict() if isinstance(model, nn.DataParallel) else model.state_dict(),\n#                 f'best_model_fold{fold_to_run}.pth'\n#             )\n            \n#     # --- End of Training ---\n#     end_time = time.time()\n#     print(f\"\\n🎉 Finished Fold {fold_to_run} in {(end_time - start_time)/60:.2f} minutes.\")\n#     print(f\"Best R² Score: {best_score:.4f}\")\n    \n#     # Cleanup\n#     del model, train_loader, valid_loader, train_dataset, valid_dataset\n#     gc.collect()\n#     torch.cuda.empty_cache()\n\n\ndef run_training(fold_to_run):\n    \"\"\"\n    Main training function implementing 2-stage (Freeze/Unfreeze) strategy\n    with FP16 mixed precision and gradient accumulation.\n    \"\"\"\n    print(f\"\\n{'='*50}\")\n    print(f\"🚀 STARTING TRAINING FOR FOLD {fold_to_run} (FP16 + Grad Accumulation) 🚀\")\n    print(f\"{'='*50}\")\n    \n    start_time = time.time()\n\n    # 1. Split data\n    print(f\"Splitting data for Fold {fold_to_run}...\")\n    train_df = df_wide[df_wide['fold'] != fold_to_run].reset_index(drop=True)\n    valid_df = df_wide[df_wide['fold'] == fold_to_run].reset_index(drop=True)\n\n    # 2. Create datasets and dataloaders\n    train_dataset = BiomassDataset(\n        df=train_df, transforms_fn=get_train_transforms, image_dir=CFG.IMAGE_DIR,\n        train_target_cols=CFG.TARGET_COLS, all_target_cols=CFG.ALL_TARGET_COLS\n    )\n    valid_dataset = BiomassDataset(\n        df=valid_df, transforms_fn=get_valid_transforms, image_dir=CFG.IMAGE_DIR,\n        train_target_cols=CFG.TARGET_COLS, all_target_cols=CFG.ALL_TARGET_COLS\n    )\n\n    train_loader = DataLoader(\n        train_dataset, batch_size=CFG.BATCH_SIZE, shuffle=True,\n        num_workers=CFG.NUM_WORKERS, pin_memory=True\n    )\n    valid_loader = DataLoader(\n        valid_dataset, batch_size=CFG.BATCH_SIZE * 2, shuffle=False,\n        num_workers=CFG.NUM_WORKERS, pin_memory=True\n    )\n\n    # 3. Initialize model, criterion, and optimizer\n    print(f\"Loading backbone '{CFG.MODEL_NAME}'...\")\n    model_base = BiomassModel(CFG.MODEL_NAME, CFG.PRETRAINED)\n\n    if torch.cuda.device_count() > 1:\n        print(f\"Using {torch.cuda.device_count()} GPUs with nn.DataParallel.\")\n        model = nn.DataParallel(model_base)\n    else:\n        model = model_base\n    model.to(CFG.DEVICE)\n\n    criterion = WeightedBiomassLoss(CFG.LOSS_WEIGHTS).to(CFG.DEVICE)\n\n    # Initialize FP16 GradScaler\n    scaler = GradScaler()\n\n    # =================================================================\n    # ✨ STAGE 1: Train \"Heads\" (Freeze Backbone)\n    # =================================================================\n    print(f\"\\n--- STAGE 1: Freezing Backbone (Training Heads Only) ---\")\n    print(f\"Epochs: 1 to {CFG.FREEZE_EPOCHS} | LR: {CFG.LEARNING_RATE}\")\n\n    for param in model.module.backbone.parameters() if isinstance(model, nn.DataParallel) else model.backbone.parameters():\n        param.requires_grad = False\n\n    optimizer = optim.Adam(\n        filter(lambda p: p.requires_grad, model.parameters()),\n        lr=CFG.LEARNING_RATE\n    )\n    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=2)\n\n    best_score = -float('inf')\n\n    for epoch in range(1, CFG.FREEZE_EPOCHS + 1):\n        print(f\"\\n--- Epoch {epoch}/{CFG.EPOCHS} (Stage 1) ---\")\n\n        train_loss = train_one_epoch(model, train_loader, criterion, optimizer, CFG.DEVICE, scaler)\n        valid_loss, competition_score = validate_one_epoch(model, valid_loader, criterion, CFG.DEVICE)\n\n        scheduler.step(valid_loss)\n\n        print(f\"Epoch {epoch} - Train Loss: {train_loss:.4f} | Valid Loss: {valid_loss:.4f} | R² Score: {competition_score:.4f}\")\n\n        if competition_score > best_score:\n            best_score = competition_score\n            print(f\"✨ R² improved! Saving model as 'best_model_fold{fold_to_run}.pth'...\")\n            torch.save(\n                model.module.state_dict() if isinstance(model, nn.DataParallel) else model.state_dict(),\n                f'best_model_fold{fold_to_run}.pth'\n            )\n\n    # =================================================================\n    # ✨ STAGE 2: Fine-tuning (Unfreeze Backbone)\n    # =================================================================\n    print(f\"\\n--- STAGE 2: Unfreezing Backbone (Fine-tuning) ---\")\n    print(f\"Epochs: {CFG.FREEZE_EPOCHS + 1} to {CFG.EPOCHS} | LR: {CFG.FINETUNE_LR}\")\n\n    for param in model.module.backbone.parameters() if isinstance(model, nn.DataParallel) else model.backbone.parameters():\n        param.requires_grad = True\n\n    optimizer = optim.Adam(model.parameters(), lr=CFG.FINETUNE_LR)\n    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, patience=3)\n\n    for epoch in range(CFG.FREEZE_EPOCHS + 1, CFG.EPOCHS + 1):\n        print(f\"\\n--- Epoch {epoch}/{CFG.EPOCHS} (Stage 2) ---\")\n\n        train_loss = train_one_epoch(model, train_loader, criterion, optimizer, CFG.DEVICE, scaler)\n        valid_loss, competition_score = validate_one_epoch(model, valid_loader, criterion, CFG.DEVICE)\n\n        scheduler.step(valid_loss)\n\n        print(f\"Epoch {epoch} - Train Loss: {train_loss:.4f} | Valid Loss: {valid_loss:.4f} | R² Score: {competition_score:.4f}\")\n\n        if competition_score > best_score:\n            best_score = competition_score\n            print(f\"✨ R² improved! Saving model as 'best_model_fold{fold_to_run}.pth'...\")\n            torch.save(\n                model.module.state_dict() if isinstance(model, nn.DataParallel) else model.state_dict(),\n                f'best_model_fold{fold_to_run}.pth'\n            )\n\n    # --- End of Training ---\n    end_time = time.time()\n    print(f\"\\n🎉 Finished Fold {fold_to_run} in {(end_time - start_time)/60:.2f} minutes.\")\n    print(f\"Best R² Score: {best_score:.4f}\")\n        # ================================================================\n    # 🧾 GENERATE OOF PREDICTIONS FOR THIS FOLD\n    # ================================================================\n    print(f\"\\nGenerating OOF predictions for Fold {fold_to_run}...\")\n\n    # Load the best saved model (highest R²)\n    model_base = BiomassModel(CFG.MODEL_NAME, pretrained=False)\n    model_path = f'best_model_fold{fold_to_run}.pth'\n    state_dict = torch.load(model_path, map_location=CFG.DEVICE)\n    model_base.load_state_dict(state_dict)\n    model_base.to(CFG.DEVICE)\n    model_base.eval()\n\n    # Collect predictions\n    preds_dict = {'total': [], 'gdm': [], 'green': []}\n    all_targets_np = []\n    image_paths = valid_df[\"image_path\"].values\n\n    with torch.no_grad(), autocast(enabled=True):\n        for (img_left, img_right, train_targets, all_targets) in tqdm(valid_loader, desc=f\"OOF Inference Fold {fold_to_run}\"):\n            img_left = img_left.to(CFG.DEVICE)\n            img_right = img_right.to(CFG.DEVICE)\n\n            out_total, out_gdm, out_green = model_base(img_left, img_right)\n\n            preds_dict['total'].append(out_total.cpu().numpy())\n            preds_dict['gdm'].append(out_gdm.cpu().numpy())\n            preds_dict['green'].append(out_green.cpu().numpy())\n            all_targets_np.append(all_targets.cpu().numpy())\n\n    # Concatenate all predictions\n    preds_dict_np = {\n        'total': np.concatenate(preds_dict['total']).flatten(),\n        'gdm':   np.concatenate(preds_dict['gdm']).flatten(),\n        'green': np.concatenate(preds_dict['green']).flatten()\n    }\n    all_targets_np = np.concatenate(all_targets_np)\n\n    # Derive the remaining two predictions\n    pred_clover = np.maximum(0, preds_dict_np['gdm'] - preds_dict_np['green'])\n    pred_dead   = np.maximum(0, preds_dict_np['total'] - preds_dict_np['gdm'])\n\n    # Create fold DataFrame\n    fold_df = pd.DataFrame({\n        \"image_path\": image_paths,\n        \"fold\": fold_to_run,\n        \"Dry_Green_g\": all_targets_np[:, 0],\n        \"Dry_Dead_g\": all_targets_np[:, 1],\n        \"Dry_Clover_g\": all_targets_np[:, 2],\n        \"GDM_g\": all_targets_np[:, 3],\n        \"Dry_Total_g\": all_targets_np[:, 4],\n        \"Dry_Green_g_pred\": preds_dict_np['green'],\n        \"Dry_Dead_g_pred\": pred_dead,\n        \"Dry_Clover_g_pred\": pred_clover,\n        \"GDM_g_pred\": preds_dict_np['gdm'],\n        \"Dry_Total_g_pred\": preds_dict_np['total']\n    })\n    oof_dfs.append(fold_df)\n\n    # Compute fold R²\n    y_true = fold_df[[\"Dry_Green_g\", \"Dry_Dead_g\", \"Dry_Clover_g\", \"GDM_g\", \"Dry_Total_g\"]].values\n    y_pred = fold_df[[\"Dry_Green_g_pred\", \"Dry_Dead_g_pred\", \"Dry_Clover_g_pred\", \"GDM_g_pred\", \"Dry_Total_g_pred\"]].values\n\n    fold_r2 = r2_score(y_true, y_pred, multioutput='raw_values')\n    weighted_r2 = np.dot(fold_r2, CFG.R2_WEIGHTS)\n    print(f\"Fold {fold_to_run} Weighted R²: {weighted_r2:.4f}\")\n    for name, r2 in zip(CFG.ALL_TARGET_COLS, fold_r2):\n        print(f\"   {name}: {r2:.4f}\")\n\n    # Cleanup\n    del model, train_loader, valid_loader, train_dataset, valid_dataset\n    gc.collect()\n    torch.cuda.empty_cache()\n\n\n\n# --- RUN CROSS-VALIDATION TRAINING ---\ntry:\n    for i in range(CFG.N_FOLDS):\n        run_training(fold_to_run=i)\n    # ================================================================\n# 🧾 COMBINE AND SAVE OOF PREDICTIONS\n# ================================================================\n    print(\"\\n========== Aggregating OOF Predictions Across Folds ==========\")\n    \n    oof_df = pd.concat(oof_dfs, ignore_index=True)\n    \n    # Compute overall OOF weighted R²\n    y_true = oof_df[[\"Dry_Green_g\", \"Dry_Dead_g\", \"Dry_Clover_g\", \"GDM_g\", \"Dry_Total_g\"]].values\n    y_pred = oof_df[[\"Dry_Green_g_pred\", \"Dry_Dead_g_pred\", \"Dry_Clover_g_pred\", \"GDM_g_pred\", \"Dry_Total_g_pred\"]].values\n    \n    r2s = r2_score(y_true, y_pred, multioutput='raw_values')\n    oof_weighted_r2 = np.dot(r2s, CFG.R2_WEIGHTS)\n    \n    print(\"\\n========== Overall OOF Scores ==========\")\n    print(f\"✅ OOF Weighted R²: {oof_weighted_r2:.4f}\")\n    for name, r2 in zip(CFG.ALL_TARGET_COLS, r2s):\n        print(f\"   {name}: {r2:.4f}\")\n    \n    # Save OOF results\n    oof_df.to_csv(\"OOF_predictions.csv\", index=False)\n    print(\"✅ Saved OOF predictions to 'OOF_predictions.csv'\")\n\nexcept Exception as e:\n    gc.collect()\n    torch.cuda.empty_cache()\n    raise e\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T05:41:23.632669Z","iopub.execute_input":"2025-11-03T05:41:23.633103Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Note: \n1. Kaggle environment is low to train base and large model with 768 image size based on above training method\n2. To train that ideally use jarvis ai lab","metadata":{}},{"cell_type":"markdown","source":"## SCOPE OF Improvement\n","metadata":{}},{"cell_type":"markdown","source":"#### 1. Data-Level Improvements (Most Impactful for Small Datasets)\n a. Advanced Augmentations\n\nYou’re already using flips, rotations, and color jitter — great start.\nYou can make augmentations smarter and biologically plausible:\n\nUse A.RandomResizedCrop(scale=(0.85, 1.0), ratio=(1.0, 1.2)) — introduces scale and zoom variation.\n\nAdd A.GaussianBlur, A.MotionBlur, or A.RandomShadow to simulate camera noise or lighting differences.\n\nUse A.CoarseDropout or A.Cutout to simulate occlusions (e.g., shadow patches or partial grass coverage).\n\nApply mild A.RandomBrightnessContrast instead of ColorJitter for more controlled light simulation.\n\n💡 Tip: Use Albumentations’ A.OneOf([...], p=0.5) to randomize augmentation combinations — this prevents over-augmentation.\n\n\n#### 2 2. Model Architecture Improvements\n✅ a. Backbone Upgrades\n\nTry stronger backbones that generalize better for vegetation imagery:\n\nconvnext_base or convnextv2_base\n\nefficientnetv2_m\n\ncoatnet_0 or beit_base_patch16_224 (transformer-based)\n\n💡 Tip: Use pretrained weights from ImageNet-21K or SSL (e.g., DINOv2) if available.\n\n\n### 3 .  Multi-fold + ensemble averaging (Already Implemented)\n\n### 4 . TTA + physical consistency correction (Already Implemented)\n\n\n### 5 .✅ d. Advanced Checkpoint Averaging\n\nInstead of saving a single best epoch, average several top checkpoints (e.g., last 3 or top 5 by R²).\nThis smooths out model fluctuations.\n\n# Simple example:\nstate_dicts = [torch.load(f) for f in best_ckpt_paths]\navg_state_dict = {k: torch.mean(torch.stack([sd[k] for sd in state_dicts]), dim=0) for k in state_dicts[0]}\nmodel.load_state_dict(avg_state_dict)\n\n\n\n\n#### 6 ⚖️ 3. Training Strategy Enhancements\n✅ a. Cosine Annealing or OneCycleLR\n\nInstead of ReduceLROnPlateau, try:\n\ntorch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n\nOr even better: OneCycleLR with warm restarts.\n\nThese produce smoother LR transitions and often yield +0.01–0.03 improvement in R².\n\n✅ b. Mixed Precision Training (AMP)\n\nEnable FP16 training using PyTorch’s torch.cuda.amp:\n\nscaler = torch.cuda.amp.GradScaler()\nwith torch.cuda.amp.autocast():\n    predictions = model(img_left, img_right)\n    loss = criterion(predictions, targets)\n\n\nThis improves training speed and numerical stability, especially on GPUs with Tensor Cores.\n\n✅ c. Gradient Accumulation\n\nIf GPU memory is limited due to 768×768 resolution, use:\n\naccumulation_steps = 2\nloss = loss / accumulation_steps\nloss.backward()\nif (batch_idx + 1) % accumulation_steps == 0:\n    optimizer.step()\n    optimizer.zero_grad()\n\n\nThis effectively doubles batch size stability.\n","metadata":{}},{"cell_type":"code","source":"\n# c. Multi-Task / Shared Decoder\n\n# Instead of 3 independent MLP heads, use:\n\n# Shared first layer, then branch out:\n\n\n# shared = nn.Sequential(\n#     nn.Linear(2 * n_features, 512),\n#     nn.ReLU(),\n#     nn.Dropout(0.3)\n# )\n# head_total = nn.Linear(512, 1)\n# head_gdm = nn.Linear(512, 1)\n# head_green = nn.Linear(512, 1)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}