{"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,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":272104372,"sourceType":"kernelVersion"}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"##### Reference:\n##### [1] Aic, CSIRO-The best baseline to start with, https://www.kaggle.com/code/gunianxue/csiro-the-best-baseline-to-start-with","metadata":{}},{"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 import Normalize, HorizontalFlip, VerticalFlip, Rotate, Resize, Compose\nfrom albumentations.pytorch import ToTensorV2\nimport timm\nimport cv2\nfrom tqdm import tqdm\nimport gc\n\n# ===============================================================\n# Configuration\n# ===============================================================\nclass CFG:\n    BASE_PATH = '/kaggle/input/csiro-biomass'\n    TEST_CSV = os.path.join(BASE_PATH, 'test.csv')\n    TEST_IMAGE_DIR = os.path.join(BASE_PATH, 'test')\n    MODEL_DIR = '/kaggle/input/csiro/'\n    SUBMISSION_FILE = 'submission.csv'\n    \n    MODEL_NAME = 'convnext_tiny'\n    IMG_SIZE = 768\n    \n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    BATCH_SIZE = 2  # Increased for faster inference\n    NUM_WORKERS = 2\n    N_FOLDS = 5\n    \n    TARGET_COLS = ['Dry_Total_g', 'GDM_g', 'Dry_Green_g']\n    ALL_TARGET_COLS = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n\nprint(f\"Device: {CFG.DEVICE} | Model: {CFG.MODEL_NAME} | Size: {CFG.IMG_SIZE}\")\n\n# ===============================================================\n# Enhanced TTA with 5 augmentations\n# ===============================================================\ndef get_tta_transforms():\n    base = [\n        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2()\n    ]\n    \n    transforms_list = [\n        # Original\n        A.Compose([A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE), *base]),\n        # Horizontal flip\n        A.Compose([A.HorizontalFlip(p=1.0), A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE), *base]),\n        # Vertical flip\n        A.Compose([A.VerticalFlip(p=1.0), A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE), *base]),\n        # Both flips\n        A.Compose([A.HorizontalFlip(p=1.0), A.VerticalFlip(p=1.0), \n                   A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE), *base]),\n        # Slight rotation\n        A.Compose([A.Rotate(limit=5, p=1.0, border_mode=cv2.BORDER_REFLECT),\n                   A.Resize(CFG.IMG_SIZE, CFG.IMG_SIZE), *base])\n    ]\n    return transforms_list\n\n# ===============================================================\n# Dataset\n# ===============================================================\nclass TestBiomassDataset(Dataset):\n    def __init__(self, df, transform_pipeline, image_dir):\n        self.df = df\n        self.transforms = transform_pipeline\n        self.image_dir = image_dir\n        self.image_paths = df['image_path'].values\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        img_path = self.image_paths[idx]\n        full_path = os.path.join(self.image_dir, os.path.basename(img_path))\n        \n        image = cv2.imread(full_path)\n        if image is None:\n            image = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        # Split into left and right\n        mid = image.shape[1] // 2\n        img_left = image[:, :mid]\n        img_right = image[:, mid:]\n        \n        # Apply transforms\n        img_left_tensor = self.transforms(image=img_left)['image']\n        img_right_tensor = self.transforms(image=img_right)['image']\n        \n        return img_left_tensor, img_right_tensor\n\n# ===============================================================\n# Model Architecture\n# ===============================================================\nclass BiomassModel(nn.Module):\n    def __init__(self, model_name, pretrained, n_targets=3):\n        super(BiomassModel, self).__init__()\n        \n        self.backbone = timm.create_model(\n            model_name, pretrained=pretrained, num_classes=0, global_pool='avg'\n        )\n        \n        n_features = self.backbone.num_features\n        n_combined = n_features * 2\n        \n        # Three prediction heads\n        self.head_total = nn.Sequential(\n            nn.Linear(n_combined, n_combined // 2), nn.ReLU(), nn.Dropout(0.3),\n            nn.Linear(n_combined // 2, 1)\n        )\n        self.head_gdm = nn.Sequential(\n            nn.Linear(n_combined, n_combined // 2), nn.ReLU(), nn.Dropout(0.3),\n            nn.Linear(n_combined // 2, 1)\n        )\n        self.head_green = nn.Sequential(\n            nn.Linear(n_combined, n_combined // 2), nn.ReLU(), nn.Dropout(0.3),\n            nn.Linear(n_combined // 2, 1)\n        )\n\n    def forward(self, img_left, img_right):\n        feat_l = self.backbone(img_left)\n        feat_r = self.backbone(img_right)\n        combined = torch.cat([feat_l, feat_r], dim=1)\n        \n        return self.head_total(combined), self.head_gdm(combined), self.head_green(combined)\n\n# ===============================================================\n# Inference Functions\n# ===============================================================\ndef predict_one_view(models_list, test_loader, device):\n    \"\"\"Predict with 5-fold ensemble for one TTA view\"\"\"\n    preds = {'total': [], 'gdm': [], 'green': []}\n    \n    with torch.no_grad():\n        for img_left, img_right in tqdm(test_loader, desc=\"  View\", leave=False):\n            img_left, img_right = img_left.to(device), img_right.to(device)\n            \n            batch_preds = {'total': [], 'gdm': [], 'green': []}\n            for model in models_list:\n                p_total, p_gdm, p_green = model(img_left, img_right)\n                batch_preds['total'].append(p_total.cpu())\n                batch_preds['gdm'].append(p_gdm.cpu())\n                batch_preds['green'].append(p_green.cpu())\n            \n            # Average across folds\n            preds['total'].append(torch.mean(torch.stack(batch_preds['total']), dim=0).numpy())\n            preds['gdm'].append(torch.mean(torch.stack(batch_preds['gdm']), dim=0).numpy())\n            preds['green'].append(torch.mean(torch.stack(batch_preds['green']), dim=0).numpy())\n    \n    return {\n        'total': np.concatenate(preds['total']).flatten(),\n        'gdm': np.concatenate(preds['gdm']).flatten(),\n        'green': np.concatenate(preds['green']).flatten()\n    }\n\ndef run_inference_with_tta():\n    \"\"\"Main inference with TTA and ensemble\"\"\"\n    print(f\"\\n{'='*50}\\n🚀 Starting Inference with TTA\\n{'='*50}\")\n    \n    # Load test data\n    test_df_long = pd.read_csv(CFG.TEST_CSV)\n    test_df_unique = test_df_long.drop_duplicates(subset=['image_path']).reset_index(drop=True)\n    print(f\"Found {len(test_df_unique)} unique test images\")\n    \n    # Load 5 models\n    print(\"\\nLoading 5 trained models...\")\n    models_list = []\n    for fold in range(CFG.N_FOLDS):\n        model_path = os.path.join(CFG.MODEL_DIR, f'best_model_fold{fold}.pth')\n        model = BiomassModel(CFG.MODEL_NAME, pretrained=False)\n        \n        try:\n            model.load_state_dict(torch.load(model_path, map_location=CFG.DEVICE))\n        except RuntimeError:\n            # Handle DataParallel weights\n            state_dict = torch.load(model_path, map_location=CFG.DEVICE)\n            new_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}\n            model.load_state_dict(new_state_dict)\n        \n        model.eval()\n        model.to(CFG.DEVICE)\n        models_list.append(model)\n    print(f\"✓ Loaded {len(models_list)} models\")\n    \n    # TTA loop\n    tta_transforms = get_tta_transforms()\n    print(f\"\\nRunning predictions with {len(tta_transforms)} TTA views...\")\n    \n    all_view_preds = []\n    for i, tta_transform in enumerate(tta_transforms):\n        print(f\"--- TTA View {i+1}/{len(tta_transforms)} ---\")\n        \n        test_dataset = TestBiomassDataset(test_df_unique, tta_transform, CFG.TEST_IMAGE_DIR)\n        test_loader = DataLoader(\n            test_dataset, batch_size=CFG.BATCH_SIZE, shuffle=False,\n            num_workers=CFG.NUM_WORKERS, pin_memory=True\n        )\n        \n        view_preds = predict_one_view(models_list, test_loader, CFG.DEVICE)\n        all_view_preds.append(view_preds)\n        print(f\"✓ Completed View {i+1}\")\n    \n    # Ensemble TTA results (weighted average: give more weight to original)\n    weights = [0.3, 0.2, 0.2, 0.15, 0.15]  # Original gets highest weight\n    final_preds = {\n        'total': np.average([d['total'] for d in all_view_preds], axis=0, weights=weights),\n        'gdm': np.average([d['gdm'] for d in all_view_preds], axis=0, weights=weights),\n        'green': np.average([d['green'] for d in all_view_preds], axis=0, weights=weights)\n    }\n    \n    print(\"✓ Predictions complete\")\n    \n    # Cleanup\n    del models_list, test_loader, test_dataset\n    gc.collect()\n    torch.cuda.empty_cache()\n    \n    return final_preds, test_df_long, test_df_unique\n\n# ===============================================================\n# Create Submission\n# ===============================================================\ndef create_submission(preds_np, test_df_long, test_df_unique):\n    \"\"\"Post-process predictions and create submission file\"\"\"\n    if preds_np is None:\n        return\n    \n    print(\"\\nCreating submission file...\")\n    \n    # Get 3 ensemble predictions\n    pred_total = preds_np['total']\n    pred_gdm = preds_np['gdm']\n    pred_green = preds_np['green']\n    \n    # Calculate remaining 2 targets with constraints\n    pred_clover = np.maximum(0, pred_gdm - pred_green)\n    pred_dead = np.maximum(0, pred_total - pred_gdm)\n    \n    # Apply consistency constraints\n    pred_gdm = np.minimum(pred_gdm, pred_total)  # GDM can't exceed Total\n    pred_green = np.minimum(pred_green, pred_gdm)  # Green can't exceed GDM\n    pred_clover = np.maximum(0, pred_gdm - pred_green)  # Recalculate clover\n    pred_dead = np.maximum(0, pred_total - pred_gdm)  # Recalculate dead\n    \n    # Create wide DataFrame\n    preds_wide = pd.DataFrame({\n        'image_path': test_df_unique['image_path'],\n        'Dry_Green_g': pred_green,\n        'Dry_Dead_g': pred_dead,\n        'Dry_Clover_g': pred_clover,\n        'GDM_g': pred_gdm,\n        'Dry_Total_g': pred_total\n    })\n    \n    # Convert to long format\n    preds_long = preds_wide.melt(\n        id_vars=['image_path'],\n        value_vars=CFG.ALL_TARGET_COLS,\n        var_name='target_name',\n        value_name='target'\n    )\n    \n    # Merge with original test data to get sample_id\n    submission = pd.merge(\n        test_df_long[['sample_id', 'image_path', 'target_name']],\n        preds_long,\n        on=['image_path', 'target_name'],\n        how='left'\n    )[['sample_id', 'target']]\n    \n    submission.to_csv(CFG.SUBMISSION_FILE, index=False)\n    print(f\"\\n🎉 Submission saved: {CFG.SUBMISSION_FILE}\")\n    print(f\"Preview:\\n{submission.head(10)}\")\n\n# ===============================================================\n# Main Execution\n# ===============================================================\nif __name__ == \"__main__\":\n    preds, df_long, df_unique = run_inference_with_tta()\n    create_submission(preds, df_long, df_unique)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-10-31T07:08:10.335699Z","iopub.execute_input":"2025-10-31T07:08:10.335925Z","iopub.status.idle":"2025-10-31T07:09:14.144483Z","shell.execute_reply.started":"2025-10-31T07:08:10.335908Z","shell.execute_reply":"2025-10-31T07:09:14.143619Z"}},"outputs":[],"execution_count":null}]}