{"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":4535,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":3327,"modelId":986},{"sourceId":711989,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":540972,"modelId":554164}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ====================================================\n# Imports\n# ====================================================\nimport os\nimport gc\nimport random\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torchvision.models import resnet18\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as T\nfrom transformers import AutoModel\n\nfrom sklearn.model_selection import KFold\nfrom sklearn.linear_model import Ridge\nfrom sklearn.preprocessing import RobustScaler\nfrom sklearn.multioutput import MultiOutputRegressor\nfrom sklearn.base import clone\n\nimport lightgbm as lgb\nimport xgboost as xgb\n\n# ====================================================\n# Config\n# ====================================================\nclass CFG:\n    DATA_PATH = '/kaggle/input/csiro-biomass/'\n    train_path = os.path.join(DATA_PATH, 'train.csv')\n    test_path = os.path.join(DATA_PATH, 'test.csv')\n    \n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    seeds = [46, 1337, 2024]\n    n_fold = 5\n    \n    img_size = 518\n    batch_size = 12\n    \n    target_cols = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'Dry_Total_g','GDM_g']\n    target_weights = [0.1, 0.1, 0.1, 0.2, 0.5]\n    \n    # Paths to DINO and SSL model checkpoints\n    dinov2_path = '/kaggle/input/dinov2/pytorch/large/1'\n    ssl_path = '/kaggle/input/resnet18backbone/pytorch/default/1/ssl_backbone_best.pth'  # replace with your path\n\n# ====================================================\n# Utilities\n# ====================================================\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n\ndef weighted_r2_metric(y_true, y_pred):\n    weights = np.array(CFG.target_weights)\n    if y_true.ndim == 1: y_true = y_true.reshape(-1,1)\n    if y_pred.ndim == 1: y_pred = y_pred.reshape(-1,1)\n    \n    ss_res = np.sum(weights * np.sum((y_true - y_pred)**2, axis=0))\n    global_mean = np.average(np.mean(y_true, axis=0), weights=weights)\n    ss_tot = np.sum(weights * np.sum((y_true - global_mean)**2, axis=0))\n    return 1 - (ss_res / (ss_tot + 1e-8))\n\ndef post_process_biomass(df_preds):\n    # Physical constraints\n    Y = df_preds[['Dry_Green_g','Dry_Clover_g','Dry_Dead_g','GDM_g','Dry_Total_g']].values.T\n    C = np.array([[1,1,0,-1,0],[0,0,1,1,-1]])\n    try:\n        inv_CCt = np.linalg.inv(C @ C.T)\n    except:\n        inv_CCt = np.linalg.pinv(C @ C.T)\n    P = np.eye(5) - C.T @ inv_CCt @ C\n    Y_reconciled = (P @ Y).T.clip(min=0)\n    df_preds[['Dry_Green_g','Dry_Clover_g','Dry_Dead_g','GDM_g','Dry_Total_g']] = Y_reconciled\n    df_preds['GDM_g'] = df_preds['Dry_Green_g'] + df_preds['Dry_Clover_g']\n    df_preds['Dry_Total_g'] = df_preds['GDM_g'] + df_preds['Dry_Dead_g']\n    return df_preds\n\nFEATURE_DIM = 128    # Must match the training\n\nclass SSLResNet(nn.Module):\n    def __init__(self, feature_dim=FEATURE_DIM):\n        super().__init__()\n        backbone = resnet18()\n        backbone.fc = nn.Identity()  # Remove original classification head\n        self.backbone = backbone\n        # Projection head (as used in training)\n        self.projection = nn.Sequential(\n            nn.Linear(512, 256),\n            nn.ReLU(inplace=True),\n            nn.Linear(256, 128)  # <- match your checkpoint\n        )\n\n    def forward(self, x):\n        h = self.backbone(x)\n        z = self.projection(h)\n        z = F.normalize(z, dim=1)\n        return z\n# ====================================================\n# Dataset\n# ====================================================\nclass BiomassDataset(Dataset):\n    def __init__(self, df, transform=None):\n        self.image_paths = df['image_path'].values\n        self.transform = transform\n    \n    def __len__(self):\n        return len(self.image_paths)\n    \n    def __getitem__(self, idx):\n        path = os.path.join(CFG.DATA_PATH, self.image_paths[idx])\n        try:\n            image = Image.open(path).convert(\"RGB\")\n        except:\n            image = Image.new(\"RGB\", (256,256))\n        if self.transform:\n            image = self.transform(image)\n        return image\n\n# ====================================================\n# Feature extraction (DINO + SSL)\n# ====================================================\ndef extract_features(df, model, ssl_model=None):\n    mean = [0.485,0.456,0.406]\n    std = [0.229,0.224,0.225]\n    \n    transforms_list = [\n        T.Compose([T.Resize(CFG.img_size), T.CenterCrop(CFG.img_size), T.ToTensor(), T.Normalize(mean,std)]),\n        T.Compose([T.Resize(CFG.img_size), T.CenterCrop(CFG.img_size), T.RandomHorizontalFlip(p=1.0), T.ToTensor(), T.Normalize(mean,std)]),\n        T.Compose([T.Resize(CFG.img_size), T.CenterCrop(CFG.img_size), T.RandomVerticalFlip(p=1.0), T.ToTensor(), T.Normalize(mean,std)])\n    ]\n    \n    feats_tta = []\n    \n    for i, transform in enumerate(transforms_list):\n        ds = BiomassDataset(df, transform=transform)\n        loader = DataLoader(\n                ds,\n                batch_size=CFG.batch_size,\n                shuffle=False,\n                num_workers=4,\n                pin_memory=True,\n                persistent_workers=True\n            )\n        model.eval()\n        ssl_features_list = []\n        din_features_list = []\n        with torch.no_grad():\n            for imgs in tqdm(loader, desc=f\"TTA {i+1}\"):\n                imgs = imgs.to(CFG.device)\n                \n                # DINO features\n                if hasattr(model, 'forward_features'):\n                    out = model.forward_features(imgs)\n                    cls_token = out[:,0,:]\n                    patch_tokens = out[:,1:,:]\n                    avg_pool = patch_tokens.mean(dim=1)\n                    max_pool = patch_tokens.max(dim=1)[0]\n                    din_feat = torch.cat([cls_token, avg_pool, max_pool], dim=1)\n                else:\n                    outputs = model(imgs)\n                    last_hidden = outputs.last_hidden_state  # (B, L, D)\n                    cls_token = last_hidden[:, 0, :]\n                    patch_tokens = last_hidden[:, 1:, :]\n                    avg_pool = patch_tokens.mean(dim=1)\n                    max_pool = patch_tokens.max(dim=1)[0]\n                    din_feat = torch.cat([cls_token, avg_pool, max_pool], dim=1)\n                \n                din_features_list.append(din_feat.cpu().numpy())\n                \n                # SSL backbone features\n                if ssl_model is not None:\n                    ssl_out = ssl_model(imgs)\n                    if ssl_out.dim() == 3:  # transformer output\n                        cls_token = ssl_out[:,0,:]\n                        patch_tokens = ssl_out[:,1:,:]\n                        avg_pool = patch_tokens.mean(dim=1)\n                        max_pool = patch_tokens.max(dim=1)[0]\n                        ssl_feat = torch.cat([cls_token, avg_pool, max_pool], dim=1)\n                    else:\n                        ssl_feat = ssl_out\n                    ssl_features_list.append(ssl_feat.cpu().numpy())\n        \n        # Combine TTA\n        din_feats = np.concatenate(din_features_list, axis=0)\n        if ssl_model is not None:\n            ssl_feats = np.concatenate(ssl_features_list, axis=0)\n            combined = np.concatenate([din_feats, ssl_feats], axis=1)\n        else:\n            combined = din_feats\n        \n        feats_tta.append(combined)\n    \n    final_feats = np.mean(feats_tta, axis=0)\n    feat_cols = [f\"feat_{i}\" for i in range(final_feats.shape[1])]\n    return pd.DataFrame(final_feats, columns=feat_cols)\n\n# ====================================================\n# Stacking Trainer\n# ====================================================\nclass StackingTrainer:\n    def __init__(self, n_folds=CFG.n_fold, seed=46):\n        self.n_folds = n_folds\n        self.seed = seed\n    \n    def get_base_models(self):\n        lgbm = lgb.LGBMRegressor(\n                n_estimators=300,\n                learning_rate=0.05,\n                num_leaves=31,\n                random_state=self.seed,\n                n_jobs=-1,\n                verbose=-1\n            )\n        xgb_model = xgb.XGBRegressor(\n                n_estimators=300,\n                learning_rate=0.05,\n                max_depth=6,             \n                subsample=0.8,\n                colsample_bytree=0.8,    \n                reg_alpha=1.0,\n                reg_lambda=1.0,\n                objective='reg:squarederror',\n                tree_method='hist',\n                device='cuda' if torch.cuda.is_available() else 'cpu',\n                random_state=self.seed\n            )\n        ridge = Ridge(alpha=1.0, random_state=self.seed)\n        return {\n            'lgbm': MultiOutputRegressor(lgbm),\n            'xgb': MultiOutputRegressor(xgb_model),\n            'ridge': MultiOutputRegressor(ridge)\n        }\n    \n    def train_predict(self, X, y, X_test):\n        kf = KFold(n_splits=self.n_folds, shuffle=True, random_state=self.seed)\n        base_models = self.get_base_models()\n        oof_preds = {name: np.zeros(y.shape) for name in base_models}\n        test_preds = {name: np.zeros((X_test.shape[0], y.shape[1])) for name in base_models}\n        \n        for fold, (tr_idx, val_idx) in enumerate(kf.split(X,y)):\n            X_tr, X_val = X[tr_idx], X[val_idx]\n            y_tr, y_val = y[tr_idx], y[val_idx]\n            \n            for name, model in base_models.items():\n                model_fold = clone(model)\n                model_fold.fit(X_tr, y_tr)\n                oof_preds[name][val_idx] = np.clip(model_fold.predict(X_val),0,None)\n                test_preds[name] += np.clip(model_fold.predict(X_test),0,None)/self.n_folds\n        \n        return oof_preds, test_preds\n    \n    def fit_stacking(self, oof_preds, y_true, test_preds_dict):\n        final_oof = np.zeros_like(y_true)\n        final_test = np.zeros_like(list(test_preds_dict.values())[0])\n        for i in range(y_true.shape[1]):\n            X_meta = np.column_stack([oof_preds['lgbm'][:,i], oof_preds['xgb'][:,i], oof_preds['ridge'][:,i]])\n            X_meta_test = np.column_stack([test_preds_dict['lgbm'][:,i], test_preds_dict['xgb'][:,i], test_preds_dict['ridge'][:,i]])\n            meta_model = Ridge(alpha=10.0, random_state=self.seed)\n            meta_model.fit(X_meta, y_true[:,i])\n            final_oof[:,i] = np.clip(meta_model.predict(X_meta),0,None)\n            final_test[:,i] = np.clip(meta_model.predict(X_meta_test),0,None)\n        return final_oof, final_test\n\n# ====================================================\n# Main\n# ====================================================\ndef main():\n    train_df = pd.read_csv(CFG.train_path)\n    test_df = pd.read_csv(CFG.test_path)\n    \n    train_wide = train_df.pivot(index='image_path', columns='target_name', values='target').reset_index()\n    train_wide = train_wide[['image_path'] + CFG.target_cols]\n    all_seed_preds = []\n    for seed in CFG.seeds:\n        print(f\"\\n Seed {seed}\")\n        seed_everything(seed)\n        \n        # Load models\n        dino_model = AutoModel.from_pretrained(\n        '/kaggle/input/dinov2/pytorch/large/1',  # folder containing the .bin & config.json\n        local_files_only=True\n        ).to(CFG.device).eval()\n\n        ssl_model = SSLResNet(feature_dim=FEATURE_DIM).to(CFG.device)\n        \n        state_dict = torch.load(CFG.ssl_path, map_location=CFG.device)\n        \n        # Handle possible 'module.' prefix\n        new_sd = {}\n        for k, v in state_dict.items():\n            kk = k\n            if k.startswith(\"module.\"):\n                kk = k[len(\"module.\"):]\n            new_sd[kk] = v\n        \n        ssl_model.load_state_dict(new_sd, strict=False)\n        ssl_model.eval()\n        \n        # Extract features\n    \n        X_test = extract_features(test_df, dino_model, ssl_model)\n        \n        # Scaling\n        scaler = RobustScaler()\n        X_train= extract_features(train_wide, dino_model, ssl_model) \n        X_train_scaled = scaler.fit_transform(X_train)\n        X_test_scaled = scaler.transform(X_test)\n        \n        # Targets\n        Y_train_log = np.log1p(train_wide[CFG.target_cols].values)\n        \n        # Stacking\n        trainer = StackingTrainer(n_folds=CFG.n_fold, seed=seed)\n        oof_preds_log, test_preds_log = trainer.train_predict(X_train_scaled, Y_train_log, X_test_scaled)\n        final_oof_log, final_test_log = trainer.fit_stacking(oof_preds_log, Y_train_log, test_preds_log)\n\n        # ... inside the seed loop ...\n        final_oof_log, final_test_log = trainer.fit_stacking(oof_preds_log, Y_train_log, test_preds_log)\n        \n        # --- ADD THESE LINES TO CALCULATE R2 ---\n        # We calculate it in log space because that's what the model was optimized for\n        seed_r2 = weighted_r2_metric(Y_train_log, final_oof_log)\n        print(f\"Seed {seed} | Weighted R2 (Log Space): {seed_r2:.4f}\")\n        \n        # Optional: Calculate for each target to see which one is failing\n        for i, col in enumerate(CFG.target_cols):\n            from sklearn.metrics import r2_score\n            col_r2 = r2_score(Y_train_log[:, i], final_oof_log[:, i])\n            print(f\"   -> {col}: {col_r2:.4f}\")\n        # ---------------------------------------\n\n        \n        # Back to original space\n        final_test = np.expm1(final_test_log)\n        all_seed_preds.append(final_test)\n        \n        del dino_model, ssl_model\n        gc.collect()\n        torch.cuda.empty_cache()\n    \n    # Average over seeds\n    final_preds = np.mean(all_seed_preds, axis=0)\n    \n    # Post-process\n    df_sub = test_df.copy()\n    df_sub[CFG.target_cols] = final_preds\n    df_sub = post_process_biomass(df_sub)\n    \n    df_unique_paths = df_sub.groupby('image_path', as_index=False)[CFG.target_cols].mean()\n    \n    submission_target_order = ['Dry_Green_g','Dry_Dead_g','Dry_Clover_g','Dry_Total_g','GDM_g']\n    \n    submission_rows = []\n    for _, row in df_unique_paths.iterrows():  # loop once per unique image\n        # Get the base filename without folders or extension\n        # e.g., 'images/test/123.jpg' -> '123'\n        base_id = os.path.basename(row['image_path']).split('.')[0]\n        \n        for col in CFG.target_cols:\n            submission_rows.append({\n                'sample_id': f\"{base_id}__{col}\", \n                'target': max(0, float(row[col])) # Ensure no negative values\n            })\n    \n    submission = pd.DataFrame(submission_rows)\n    print(submission.head(10))\n    submission.to_csv(\"/kaggle/working/submission.csv\", index=False)\n    print(\"✅ Submission saved successfully!\")\n\nif __name__=='__main__':\n    main()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-08T17:19:56.094657Z","iopub.execute_input":"2026-01-08T17:19:56.095005Z","iopub.status.idle":"2026-01-08T17:39:05.319314Z","shell.execute_reply.started":"2026-01-08T17:19:56.094970Z","shell.execute_reply":"2026-01-08T17:39:05.318493Z"}},"outputs":[],"execution_count":null}]}