{"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,"sourceType":"competition"}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport sys\nimport gc\nimport random\nimport pickle\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport timm \nimport xgboost as xgb\nfrom tqdm.auto import tqdm\nfrom sklearn.model_selection import StratifiedGroupKFold\nfrom sklearn.preprocessing import StandardScaler\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as T\nfrom PIL import Image\n\n# ====================================================\n# CONFIGURATION\n# ====================================================\nclass Config:\n    seed = 42\n    img_size = 518 \n    model_name = 'vit_large_patch14_dinov2'\n    \n    n_folds = 5\n    batch_size = 8 \n    num_workers = 4 \n    \n    data_root = '/kaggle/input/csiro-biomass'\n    train_csv_path = os.path.join(data_root, 'train.csv')\n    output_dir = 'training_artifacts'\n    \n    target_cols = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n    target_weights = np.array([0.1, 0.1, 0.1, 0.2, 0.5])\n    \n    height_col = 'Height_Ave_cm'\n    ndvi_col = 'Pre_GSHH_NDVI'\n\nos.makedirs(Config.output_dir, exist_ok=True)\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(seed)\n    torch.backends.cudnn.deterministic = True\n\nseed_everything(Config.seed)\n\ndef weighted_r2_score(y_true, y_pred, weights):\n    W = np.repeat(weights[None, :], y_true.shape[0], axis=0)\n    y_true_f = y_true.flatten()\n    y_pred_f = y_pred.flatten()\n    W_f = W.flatten()\n    y_w_mean = np.sum(W_f * y_true_f) / np.sum(W_f)\n    ss_res = np.sum(W_f * (y_true_f - y_pred_f)**2)\n    ss_tot = np.sum(W_f * (y_true_f - y_w_mean)**2)\n    return 1 - (ss_res / (ss_tot + 1e-6))\n\ndef remove_artifacts(image_np):\n    hsv = cv2.cvtColor(image_np, cv2.COLOR_RGB2HSV)\n    lower_orange = np.array([0, 150, 150])\n    upper_orange = np.array([25, 255, 255])\n    mask = cv2.inRange(hsv, lower_orange, upper_orange)\n    kernel = np.ones((3,3), np.uint8)\n    mask = cv2.dilate(mask, kernel, iterations=2)\n    clean_img = cv2.inpaint(image_np, mask, 3, cv2.INPAINT_TELEA)\n    return clean_img\n\ndef get_color_stats(image_np):\n    \"\"\"\n    Extracts simple color statistics from the image.\n    Input: RGB numpy image\n    Output: np array of shape (12,) -> [R_mean, R_std, G_mean, ... V_std]\n    \"\"\"\n    # RGB Stats\n    r_mean, r_std = np.mean(image_np[:,:,0]), np.std(image_np[:,:,0])\n    g_mean, g_std = np.mean(image_np[:,:,1]), np.std(image_np[:,:,1])\n    b_mean, b_std = np.mean(image_np[:,:,2]), np.std(image_np[:,:,2])\n    \n    # HSV Stats (Crucial for Green vs Dead)\n    hsv = cv2.cvtColor(image_np, cv2.COLOR_RGB2HSV)\n    h_mean, h_std = np.mean(hsv[:,:,0]), np.std(hsv[:,:,0])\n    s_mean, s_std = np.mean(hsv[:,:,1]), np.std(hsv[:,:,1])\n    v_mean, v_std = np.mean(hsv[:,:,2]), np.std(hsv[:,:,2])\n    \n    # Normalize roughly to 0-1 range (Pixel values are 0-255)\n    stats = np.array([\n        r_mean, r_std, g_mean, g_std, b_mean, b_std,\n        h_mean, h_std, s_mean, s_std, v_mean, v_std\n    ]) / 255.0\n    \n    return stats\n\ndef preprocess_train_data(csv_path):\n    df = pd.read_csv(csv_path)\n    if set(Config.target_cols).issubset(df.columns): return df\n    print(\"Converting Train Data from Long to Wide format...\")\n    meta_cols = ['image_path', 'Sampling_Date', 'State', 'Species', Config.ndvi_col, Config.height_col]\n    df_wide = df[meta_cols].drop_duplicates(subset=['image_path']).reset_index(drop=True)\n    targets_pivot = df.pivot(index='image_path', columns='target_name', values='target').reset_index()\n    df_wide = df_wide.merge(targets_pivot, on='image_path', how='left')\n    for col in Config.target_cols:\n        if col not in df_wide.columns: df_wide[col] = 0.0\n    return df_wide\n\n# ====================================================\n# FEATURE EXTRACTOR\n# ====================================================\nclass FeatureExtractor(torch.nn.Module):\n    def __init__(self, model_name, pretrained=True):\n        super().__init__()\n        self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0, global_pool='')\n        self.backbone.eval()\n        \n    def forward(self, x):\n        b, t, c, h, w = x.shape\n        x = x.view(b * t, c, h, w)\n        features = self.backbone(x)\n        features = features.mean(dim=1) \n        features = features.view(b, t, -1)\n        features = torch.mean(features, dim=1)\n        return features\n\ndef get_transforms():\n    return T.Compose([\n        T.Resize((Config.img_size, Config.img_size)),\n        T.ToTensor(),\n        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n    ])\n\n# ====================================================\n# DATASET\n# ====================================================\nclass BiomassDataset(Dataset):\n    def __init__(self, df, data_root, transforms=None, do_flip=False):\n        self.df = df\n        self.data_root = data_root\n        self.transforms = transforms\n        self.do_flip = do_flip \n        self.image_paths = [os.path.join(data_root, p) for p in df['image_path']]\n        \n    def __len__(self):\n        return len(self.df)\n    \n    def __getitem__(self, idx):\n        img_path = self.image_paths[idx]\n        try:\n            # Open with CV2 for Color Stats\n            image_cv = cv2.imread(img_path)\n            image_cv = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)\n            image_cv = remove_artifacts(image_cv)\n            \n            if self.do_flip:\n                image_cv = cv2.flip(image_cv, 1)\n                \n            # Extract Color Stats (12 features)\n            color_stats = get_color_stats(image_cv)\n            \n            # Convert to PIL for DINO transforms\n            image = Image.fromarray(image_cv)\n        except:\n            image = Image.new('RGB', (Config.img_size, Config.img_size))\n            color_stats = np.zeros(12, dtype=np.float32)\n\n        w, h = image.size\n        # Tiling logic\n        left_img = image.crop((0, 0, w // 2, h))\n        right_img = image.crop((w // 2, 0, w, h))\n        \n        if self.transforms:\n            left_t = self.transforms(left_img)\n            right_t = self.transforms(right_img)\n        \n        images = torch.stack([left_t, right_t])\n        \n        height = self.df.iloc[idx].get(Config.height_col, 0.0)\n        ndvi = self.df.iloc[idx].get(Config.ndvi_col, 0.0)\n        if pd.isna(height): height = 0.0\n        if pd.isna(ndvi): ndvi = 0.0\n            \n        meta = torch.tensor([height / 30.0, ndvi], dtype=torch.float32)\n        targets = self.df.iloc[idx][Config.target_cols].values.astype(np.float32)\n        \n        return images, torch.tensor(color_stats, dtype=torch.float32), meta, torch.tensor(targets)\n\ndef extract_features_with_color(df, data_root, device):\n    backbone = FeatureExtractor(Config.model_name, pretrained=True)\n    backbone.to(device)\n    torch.save(backbone.state_dict(), os.path.join(Config.output_dir, 'backbone_weights.pth'))\n    torch.backends.cudnn.benchmark = True\n    \n    # Helper for extraction loop\n    def run_pass(do_flip):\n        print(f\"Extracting {'Flipped' if do_flip else 'Original'}...\")\n        ds = BiomassDataset(df, data_root, transforms=get_transforms(), do_flip=do_flip)\n        dl = DataLoader(ds, batch_size=Config.batch_size, shuffle=False, num_workers=Config.num_workers, pin_memory=True)\n        \n        feats, colors, metas, targs = [], [], [], []\n        for imgs, col_stat, meta, targ in tqdm(dl):\n            imgs = imgs.to(device)\n            with torch.amp.autocast('cuda'):\n                with torch.no_grad(): f = backbone(imgs)\n            \n            feats.append(f.float().cpu().numpy())\n            colors.append(col_stat.numpy())\n            metas.append(meta.numpy())\n            targs.append(targ.numpy())\n            \n        return np.vstack(feats), np.vstack(colors), np.vstack(metas), np.vstack(targs)\n\n    # Two passes\n    X_o, C_o, M_o, Y_o = run_pass(False)\n    X_f, C_f, M_f, Y_f = run_pass(True)\n        \n    del backbone\n    torch.cuda.empty_cache()\n    gc.collect()\n    \n    return (X_o, C_o, M_o, Y_o), (X_f, C_f, M_f, Y_f)\n\n# ====================================================\n# TRAINING EXECUTION\n# ====================================================\ndef run_training():\n    print(\"--- CSIRO TRAINING (COLOR-ENHANCED XGBOOST) ---\")\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    \n    train_df = preprocess_train_data(Config.train_csv_path)\n    train_df['group_key'] = train_df['Sampling_Date'].astype(str) + '_' + train_df['State']\n    train_df['stratify_bin'] = pd.qcut(train_df['Dry_Total_g'], q=10, labels=False)\n    \n    sgkf = StratifiedGroupKFold(n_splits=Config.n_folds)\n    train_df['fold'] = -1\n    for fold, (train_idx, val_idx) in enumerate(sgkf.split(train_df, train_df['stratify_bin'], groups=train_df['group_key'])):\n        train_df.loc[val_idx, 'fold'] = fold\n        \n    # Extract DINO + Color + Meta\n    (X_o, C_o, M_o, Y_o), (X_f, C_f, M_f, Y_f) = extract_features_with_color(train_df, Config.data_root, device)\n    \n    # CONCATENATE ALL FEATURES\n    # Features = [DINO(1024), Color(12), Meta(2)] -> Total 1038 features\n    F_o = np.hstack([X_o, C_o, M_o])\n    F_f = np.hstack([X_f, C_f, M_f])\n    \n    oof_preds = np.zeros_like(Y_o)\n    \n    print(\"\\nStarting Training with Color Features...\")\n    \n    for fold in range(Config.n_folds):\n        train_mask = train_df['fold'] != fold\n        val_mask = train_df['fold'] == fold\n        \n        # Train on Orig + Flip\n        X_train = np.vstack([F_o[train_mask], F_f[train_mask]])\n        y_train = np.vstack([Y_o[train_mask], Y_f[train_mask]])\n        \n        X_val = F_o[val_mask]\n        y_val_true = Y_o[val_mask]\n        \n        fold_dir = os.path.join(Config.output_dir, f'fold_{fold}')\n        os.makedirs(fold_dir, exist_ok=True)\n        \n        fold_preds = []\n        \n        for t_idx in range(5):\n            # Target-Specific Hyperparameters (Same as before, trusted)\n            if t_idx in [1, 2]: # Dead, Clover\n                current_params = {\n                    'objective': 'reg:squarederror', 'n_estimators': 2000, 'learning_rate': 0.01,\n                    'max_depth': 2, \n                    'subsample': 0.6, 'colsample_bytree': 0.2,\n                    'reg_alpha': 10.0, 'reg_lambda': 10.0,\n                    'n_jobs': -1, 'early_stopping_rounds': 50, 'tree_method': 'hist'\n                }\n            elif t_idx == 4: # Total\n                current_params = {\n                    'objective': 'reg:squarederror', 'n_estimators': 3000, 'learning_rate': 0.015,\n                    'max_depth': 4, \n                    'subsample': 0.8, 'colsample_bytree': 0.4,\n                    'reg_alpha': 0.1, 'reg_lambda': 1.0,\n                    'n_jobs': -1, 'early_stopping_rounds': 50, 'tree_method': 'hist'\n                }\n            else: # Green, GDM\n                current_params = {\n                    'objective': 'reg:squarederror', 'n_estimators': 3000, 'learning_rate': 0.015,\n                    'max_depth': 3, \n                    'subsample': 0.7, 'colsample_bytree': 0.3,\n                    'reg_alpha': 1.0, 'reg_lambda': 5.0,\n                    'n_jobs': -1, 'early_stopping_rounds': 50, 'tree_method': 'hist'\n                }\n\n            model = xgb.XGBRegressor(**current_params)\n            model.fit(X_train, y_train[:, t_idx], eval_set=[(X_val, y_val_true[:, t_idx])], verbose=False)\n            model.save_model(os.path.join(fold_dir, f'xgb_target_{t_idx}.json'))\n            fold_preds.append(model.predict(X_val))\n            \n        fold_preds = np.column_stack(fold_preds)\n        fold_preds = np.clip(fold_preds, 0, None)\n        oof_preds[val_mask] = fold_preds\n        \n        score = weighted_r2_score(y_val_true, fold_preds, Config.target_weights)\n        print(f\"Fold {fold} | R2: {score:.4f}\")\n        \n    overall_score = weighted_r2_score(Y_o, oof_preds, Config.target_weights)\n    print(f\"\\n>>> OVERALL CV SCORE: {overall_score:.4f} <<<\")\n    \n    # Save meta stats for inference (Still needed for Height/NDVI imputation)\n    meta_stats = {'height_mean': np.mean(M_o[:, 0]), 'ndvi_mean': np.mean(M_o[:, 1])}\n    with open(os.path.join(Config.output_dir, 'meta_stats.pkl'), 'wb') as f:\n        pickle.dump(meta_stats, f)\n        \n    print(f\"Training complete. Artifacts saved in '{Config.output_dir}'\")\n\nif __name__ == \"__main__\":\n    run_training()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}