{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","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,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":4534,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":3326,"modelId":986},{"sourceId":64765,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":54020,"modelId":74163},{"sourceId":690934,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":523821,"modelId":537838}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport sys\n\n# --- 1. ENVIRONMENT & DEPENDENCY FIXES ---\nos.environ[\"PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION\"] = \"python\"\nos.environ['TRANSFORMERS_OFFLINE'] = '1'\nos.environ['HF_HUB_OFFLINE'] = '1'\nos.environ['HF_DATASETS_OFFLINE'] = '1'\nos.environ['WANDB_DISABLED'] = 'true'\n\n# # Auto-uninstall TF blindly (prevents CUDA/Protobuf conflicts)\n# # Do NOT import tensorflow here to check (it poisons the protobuf env)\n# print(\"Triggering blind uninstall of tensorflow...\")\n# os.system(\"pip uninstall -y tensorflow tensorflow-io tensorflow-estimator tensorboard keras\")\n\nimport cv2\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport pandas as pd\nimport numpy as np\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom tqdm import tqdm\nimport warnings\nimport timm\nfrom transformers import AutoModel\nfrom torchvision import models\n\nwarnings.filterwarnings(\"ignore\")\n\n# --- 2. CONFIGURATION (ABSOLUTE PATHS) ---\nclass Config:\n    # PROJECT\n    PROJECT_NAME = \"CSIRO_Biomass_Fusion\"\n    \n    # PATHS (Dynamic)\n    if os.path.exists('/kaggle/input'):\n        # STORAGE IS LINUX/KAGGLE\n        ROOT_DIR = '/kaggle/input/csiro-biomass'\n        OUTPUT_DIR = '/kaggle/working'\n        \n        FUSION_CHECKPOINT = \"/kaggle/input/csrio-test-2/pytorch/1/1/model_fusion.pth\"\n        if not os.path.exists(FUSION_CHECKPOINT):\n             # Try input datasets\n             FUSION_CHECKPOINT = \"/kaggle/input/csiro-weights/model_fusion.pth\"\n             \n        DINOV2_HF_PATH = \"/kaggle/input/dinov2/pytorch/base/1\"\n        EFFNET_PATH = \"/kaggle/input/csiro-weights/model_efficientnet.pth\"\n        \n        TEST_CSV = os.path.join(ROOT_DIR, \"test.csv\")\n    \n    else:\n        # STORAGE IS WINDOWS (LOCAL)\n        ROOT_DIR = r\"d:\\Kaggle\\CSIRO\\csiro-biomass\"\n        OUTPUT_DIR = r\"d:\\Kaggle\\CSIRO\"\n        \n        FUSION_CHECKPOINT = os.path.join(OUTPUT_DIR, \"model_fusion.pth\")\n        DINOV2_HF_PATH = \"\" \n        EFFNET_PATH = \"\" \n        \n        TEST_CSV = os.path.join(ROOT_DIR, \"test.csv\")\n    SUBMISSION_FILE = \"submission.csv\"\n    \n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    print(f\"Running on Device: {DEVICE}\")\n    \n    IMAGE_SIZE = 224\n    BATCH_SIZE = 16\n    \n    TARGETS = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n    NUM_SPECIES = 15\n    MAX_HEIGHT = 200.0\n\nconfig = Config()\n\n# --- 3. MODEL DEFINITION (FUSION) ---\nclass MultiTaskHead(nn.Module):\n    def __init__(self, in_features, num_targets, num_species):\n        super().__init__()\n        self.head_biomass = nn.Sequential(nn.Linear(in_features, 512), nn.SiLU(), nn.Dropout(0.2), nn.Linear(512, num_targets))\n        self.head_height = nn.Sequential(nn.Linear(in_features, 64), nn.ReLU(), nn.Linear(64, 1))\n        self.head_ndvi = nn.Sequential(nn.Linear(in_features, 64), nn.ReLU(), nn.Linear(64, 1))\n        self.head_species = nn.Sequential(nn.Linear(in_features, 128), nn.ReLU(), nn.Linear(128, num_species))\n\n    def forward(self, x):\n        return {\n            'biomass': self.head_biomass(x),\n            'height': self.head_height(x),\n            'ndvi': self.head_ndvi(x),\n            'species': self.head_species(x)\n        }\n\nclass FusionBiomassModel(nn.Module):\n    def __init__(self, num_targets=5, num_species=15):\n        super().__init__()\n        print(f\"Initializing Fusion Model...\")\n        \n        # 1. DinoV2 Backbone\n        self.dino_backbone = None\n        if os.path.exists(config.DINOV2_HF_PATH):\n             try:\n                print(f\"Loading DinoV2 from {config.DINOV2_HF_PATH}\")\n                self.dino_backbone = AutoModel.from_pretrained(config.DINOV2_HF_PATH)\n                self.dino_backbone.gradient_checkpointing_enable() \n                self.dino_dim = 768\n                print(\"Success: DinoV2 loaded.\")\n             except Exception as e: \n                print(f\"DinoV2 Load Failed: {e}\")\n        \n        if self.dino_backbone is None:\n             try:\n                 print(\"Trying Hub load...\")\n                 self.dino_backbone = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14')\n                 self.dino_dim = 384\n             except: pass\n\n        if self.dino_backbone is None:\n             print(\"WARNING: Using Dummy DinoV2 (Identity/Pooling)\")\n             self.dino_backbone = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Flatten())\n             self.dino_dim = 3\n\n        # 2. EfficientNet Backbone\n        self.eff_backbone = models.efficientnet_v2_s(weights=None)\n        self.eff_dim = self.eff_backbone.classifier[1].in_features\n        self.eff_backbone.classifier = nn.Identity()\n        \n        # 3. Head\n        self.fusion_dim = self.dino_dim + self.eff_dim\n        print(f\"Fusion Dimension: {self.dino_dim} (Dino) + {self.eff_dim} (Eff) = {self.fusion_dim}\")\n        self.heads = MultiTaskHead(self.fusion_dim, num_targets, num_species)\n        \n    def forward(self, x):\n        # A. Dino\n        if hasattr(self.dino_backbone, 'embeddings'):\n            dino_out = self.dino_backbone(x).last_hidden_state[:, 0, :]\n        else:\n            dino_out = self.dino_backbone(x)\n            if isinstance(dino_out, dict): dino_out = dino_out['x_norm_clstoken']\n            \n        # B. EffNet\n        eff_out = self.eff_backbone(x)\n        \n        # Check dim\n        if dino_out.dim() != eff_out.dim():\n            if dino_out.dim() == 4: dino_out = torch.flatten(dino_out, 1)\n\n        # C. Concat\n        combined = torch.cat([dino_out, eff_out], dim=1)\n        return self.heads(combined)\n\n# --- 4. DATASET (Exact Logic from kaggle_inference.py) ---\nclass CSIRODataset(Dataset):\n    def __init__(self, csv_file, root_dir):\n        self.root_dir = root_dir\n        self.df = pd.read_csv(csv_file)\n        if 'image_path' in self.df.columns:\n            self.df_grouped = self.df.groupby('image_path').first().reset_index()\n        else:\n            self.df_grouped = self.df\n        \n        self.transform = A.Compose([\n            A.Resize(config.IMAGE_SIZE, config.IMAGE_SIZE),\n            A.Normalize(),\n            ToTensorV2()\n        ])\n\n    def __len__(self):\n        return len(self.df_grouped)\n\n    def __getitem__(self, idx):\n        row = self.df_grouped.iloc[idx]\n        image_path = os.path.join(self.root_dir, row['image_path'])\n        \n        image = cv2.imread(image_path)\n        if image is None:\n            image = np.zeros((config.IMAGE_SIZE, config.IMAGE_SIZE, 3), dtype=np.uint8)\n        else:\n            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n            \n        image = self.transform(image=image)['image']\n        return image, row['image_path']\n\n# --- 5. MAIN INFERENCE LOOP ---\ndef run_inference():\n    print(\"--- Starting Fusion Inference ---\")\n    \n    # A. Load Model\n    model = FusionBiomassModel(num_targets=5, num_species=config.NUM_SPECIES).to(config.DEVICE)\n    \n    weights_loaded = False\n    if os.path.exists(config.FUSION_CHECKPOINT):\n        print(f\"Loading weights from {config.FUSION_CHECKPOINT}\")\n        try:\n            state_dict = torch.load(config.FUSION_CHECKPOINT, map_location=config.DEVICE)\n            model.load_state_dict(state_dict)\n            model.eval()\n            weights_loaded = True\n        except Exception as e:\n            print(f\"Error loading weights: {e}\")\n            try:\n                # Fallback: maybe structure mismatch? Try non-strict\n                model.load_state_dict(state_dict, strict=False)\n                model.eval()\n                print(\"Loaded with strict=False\")\n                weights_loaded = True\n            except:\n                print(\"Failed strict=False load too.\")\n    else:\n        print(f\"Weights not found: {config.FUSION_CHECKPOINT}\")\n        \n    if not weights_loaded:\n        print(\"CRITICAL: Running with random init weights (Zero-Submission Mode).\")\n    \n    # B. Generate Submission\n    # Fallback init\n    test_df = pd.read_csv(config.TEST_CSV) if os.path.exists(config.TEST_CSV) else pd.DataFrame()\n    results = []\n    \n    if os.path.exists(config.TEST_CSV):\n        ds = CSIRODataset(config.TEST_CSV, config.ROOT_DIR)\n        loader = DataLoader(ds, batch_size=config.BATCH_SIZE, shuffle=False, num_workers=0)\n        \n        # We will collect raw predictions first to apply global scaling/consistency\n        raw_rows = []\n        \n        if weights_loaded:\n            with torch.no_grad():\n                for imgs, paths in tqdm(loader, desc=\"Predicting (TTA)\"):\n                    imgs = imgs.to(config.DEVICE)\n                    \n                    # TTA Variants (HFlip, VFlip)\n                    imgs_hflip = torch.flip(imgs, [3]) \n                    imgs_vflip = torch.flip(imgs, [2]) \n                    \n                    def predict(batch_imgs):\n                        out = model(batch_imgs)\n                        if isinstance(out, dict):\n                            return torch.expm1(out['biomass'])\n                        return torch.zeros(batch_imgs.shape[0], 5).to(config.DEVICE)\n                    \n                    p1 = predict(imgs)\n                    p2 = predict(imgs_hflip)\n                    p3 = predict(imgs_vflip)\n                    \n                    batch_pred_bio = (p1 + p2 + p3) / 3.0\n                    batch_pred_bio = batch_pred_bio.cpu().numpy()\n                    \n                    for path, pred in zip(paths, batch_pred_bio):\n                        # pred is array of 5 targets: 'Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g'\n                        # We map these to the post-processing naming convention\n                        row_data = {\n                            'image_path': path,\n                            'pred_Green': pred[0],\n                            'pred_Dead': pred[1],\n                            'pred_Clover': pred[2],\n                            # pred[3] is GDM, we skip for consistency calc usually, or use check\n                            'pred_Total': pred[4]\n                        }\n                        raw_rows.append(row_data)\n        \n        # --- APPLY AGGRESSIVE POST-PROCESSING ---\n        if raw_rows:\n            print(\"Applying Aggressive Post-Processing...\")\n            scale_df = pd.DataFrame(raw_rows)\n            \n            # 1. Clip negative values\n            target_cols = ['pred_Total', 'pred_Green', 'pred_Dead', 'pred_Clover']\n            for c in target_cols:\n                scale_df[c] = scale_df[c].clip(lower=0.0)\n            \n            # 2. Noise Removal: Zero out very low biomass predictions\n             # Use 0.5 as threshold (from 03_make_submission.py)\n            mask_low_total = scale_df['pred_Total'] < 0.5\n            scale_df.loc[mask_low_total, target_cols] = 0.0\n            \n            # 3. Component Ratio Thresholding\n            #    If a component contributes < 1% to total, treat as noise\n            for c in ['pred_Green', 'pred_Dead', 'pred_Clover']:\n                ratio = scale_df[c] / (scale_df['pred_Total'] + 1e-6)\n                scale_df.loc[ratio < 0.01, c] = 0.0\n\n            # 4. Consistency Scaling: Re-distribute to match Total exactly\n            sum_comp = scale_df['pred_Green'] + scale_df['pred_Dead'] + scale_df['pred_Clover']\n            factor = scale_df['pred_Total'] / (sum_comp + 1e-6)\n            \n            scale_df['Final_Green'] = scale_df['pred_Green'] * factor\n            scale_df['Final_Dead'] = scale_df['pred_Dead'] * factor\n            scale_df['Final_Clover'] = scale_df['pred_Clover'] * factor\n            scale_df['Final_GDM'] = scale_df['Final_Green'] + scale_df['Final_Clover']\n            scale_df['Final_Total'] = scale_df['pred_Total']\n            \n            # Convert to Map for O(1) Lookup\n            # map: image_path -> {target_name: value}\n            preds_map = {}\n            for _, row in scale_df.iterrows():\n                preds_map[row['image_path']] = {\n                    'Dry_Green_g': row['Final_Green'],\n                    'Dry_Dead_g': row['Final_Dead'],\n                    'Dry_Clover_g': row['Final_Clover'],\n                    'GDM_g': row['Final_GDM'], # Derived\n                    'Dry_Total_g': row['Final_Total']\n                }\n        else:\n            preds_map = {}\n\n        # Map back to sample_id\n        for _, row in test_df.iterrows():\n            val = preds_map.get(row['image_path'], {}).get(row['target_name'], 0.0)\n            results.append({'sample_id': row['sample_id'], 'target': max(0.0, val)})\n            \n    else:\n        print(\"Test CSV not found.\")\n\n    # C. Save (Final)\n    sub = pd.DataFrame(results)\n    if sub.empty:\n        # Ultimate fallback\n        sub = pd.DataFrame(columns=['sample_id', 'target'])\n    else:\n        # --- ENFORCE SORT ORDER ---\n        # Request: Green -> Dead -> Clover -> GDM -> Total\n        \n        # 1. Extract Target Name from sample_id\n        # Assuming format: ID__TargetName OR ID_TargetName (if kaggle changes)\n        # But we relied on test.csv having 'target_name' earlier. \n        # Ideally we join with test_df or parse sample_id.\n        # Let's parse sample_id for robustness if column missing in 'sub'.\n        \n        def get_target_rank(sid):\n            if 'Dry_Green_g' in sid: return 0\n            if 'Dry_Dead_g' in sid: return 1\n            if 'Dry_Clover_g' in sid: return 2\n            if 'GDM_g' in sid: return 3\n            if 'Dry_Total_g' in sid: return 4\n            return 99 # Unknown\n\n        # We also need to sort by Image ID to keep groups together\n        # Extract base ID\n        def get_base_id(sid):\n             return sid.split('__')[0]\n\n        sub['rank'] = sub['sample_id'].apply(get_target_rank)\n        sub['base_id'] = sub['sample_id'].apply(get_base_id)\n        \n        # Sort by Base ID then Rank\n        sub = sub.sort_values(by=['base_id', 'rank'])\n        \n        # Clean up\n        sub = sub.drop(columns=['rank', 'base_id'])\n        \n    sub.to_csv(config.SUBMISSION_FILE, index=False)\n    print(f\"Submission Saved: {len(sub)} rows.\")\n\nif __name__ == \"__main__\":\n    try:\n        run_inference()\n    except Exception as e:\n        print(f\"Global Crash: {e}\")\n        pd.DataFrame({'sample_id': [], 'target': []}).to_csv(\"submission.csv\", index=False)\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-23T16:53:51.347338Z","iopub.execute_input":"2025-12-23T16:53:51.348036Z","iopub.status.idle":"2025-12-23T16:54:59.638318Z","shell.execute_reply.started":"2025-12-23T16:53:51.347993Z","shell.execute_reply":"2025-12-23T16:54:59.637552Z"}},"outputs":[],"execution_count":null}]}