{"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":"gpu","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":13619023,"sourceType":"datasetVersion","datasetId":8655075}],"dockerImageVersionId":31154,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# ============================================================================\n# CSIRO Image2Biomass Prediction - INFERENCE (Corrected with safe loader)\n# ============================================================================\nimport numpy as np\nimport pandas as pd\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\nimport cv2\nfrom tqdm import tqdm\nimport warnings\nimport os\nimport pickle\n\nwarnings.filterwarnings('ignore')\n\n# ============================================================================\n# CONFIGURATION\n# ============================================================================\nclass CFG:\n    # Paths\n    test_csv = '/kaggle/input/csiro-biomass/test.csv'\n    test_dir = '/kaggle/input/csiro-biomass/test'\n    model_dir = '/kaggle/input/csiro-models-awais'  # Directory containing model checkpoints\n    output_file = 'submission.csv'\n\n    # Model settings (MUST match training configuration)\n    model_name = 'tf_efficientnetv2_m'\n    img_size = 512\n    n_folds = 5  # Number of folds to ensemble\n\n    # Inference settings\n    batch_size = 32\n    num_workers = 4\n    use_tta = True\n    tta_steps = 5\n\n    # Target names (order matters!)\n    targets = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n\n    # Device\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nprint(f\"Using device: {CFG.device}\")\nprint(f\"Will ensemble up to {CFG.n_folds} models\")\nprint(f\"TTA enabled: {CFG.use_tta} ({CFG.tta_steps} augmentations)\" if CFG.use_tta else \"TTA disabled\")\n\n# ============================================================================\n# DATASET CLASS\n# ============================================================================\nclass BiomassTestDataset(Dataset):\n    \"\"\"\n    Dataset for test/inference data.\n    Loads images and applies tabular feature scaling.\n    \"\"\"\n    def __init__(self, df, img_dir, transform=None, tabular_scaler=None, has_metadata=True):\n        self.df = df.reset_index(drop=True)\n        self.img_dir = img_dir\n        self.transform = transform\n        self.has_metadata = has_metadata\n\n        # Prepare tabular features\n        if has_metadata and 'Pre_GSHH_NDVI' in df.columns and 'Height_Ave_cm' in df.columns:\n            tabular_data = df[['Pre_GSHH_NDVI', 'Height_Ave_cm']].fillna(0).values.astype(np.float32)\n            if tabular_scaler is not None:\n                # scaler is a sklearn StandardScaler (or similar)\n                try:\n                    self.tabular_features = tabular_scaler.transform(tabular_data)\n                except Exception:\n                    # if scaler was saved in a way that doesn't transform on CPU, still try\n                    self.tabular_features = tabular_data\n            else:\n                self.tabular_features = tabular_data\n        else:\n            # No metadata available - use zeros (scaled mean)\n            print(\"  ⚠ No metadata in test set - using zero values (scaled mean)\")\n            self.tabular_features = np.zeros((len(df), 2), dtype=np.float32)\n\n    def __len__(self):\n        return len(self.df)\n\n    def _resolve_image_path(self, row):\n        # Try multiple ways to find the image file robustly\n        if 'image_path' in row and isinstance(row['image_path'], str) and row['image_path'].strip() != '':\n            candidate = row['image_path']\n            if os.path.isabs(candidate) and os.path.exists(candidate):\n                return candidate\n            basename = os.path.basename(candidate)\n            joined = os.path.join(self.img_dir, basename)\n            if os.path.exists(joined):\n                return joined\n            if os.path.exists(candidate):\n                return candidate\n\n        if 'image_id' in row and not pd.isna(row['image_id']):\n            fname = str(row['image_id']).strip() + '.jpg'\n            joined = os.path.join(self.img_dir, fname)\n            if os.path.exists(joined):\n                return joined\n\n        if 'image_path' in row and isinstance(row['image_path'], str):\n            return os.path.join(self.img_dir, os.path.basename(row['image_path']))\n\n        raise FileNotFoundError(f\"No valid image path for row: {row.to_dict()}\")\n\n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n\n        img_path = self._resolve_image_path(row)\n        image = cv2.imread(img_path)\n        if image is None:\n            raise FileNotFoundError(f\"Failed to load image: {img_path}\")\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n\n        if self.transform:\n            augmented = self.transform(image=image)\n            image = augmented['image']\n\n        tabular = torch.tensor(self.tabular_features[idx], dtype=torch.float32)\n        return image, tabular\n\n# ============================================================================\n# AUGMENTATION TRANSFORMS\n# ============================================================================\ndef get_inference_transforms():\n    return A.Compose([\n        A.Resize(CFG.img_size, CFG.img_size),\n        A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n        ToTensorV2(),\n    ])\n\ndef get_tta_transforms():\n    transforms = [\n        A.Compose([A.Resize(CFG.img_size, CFG.img_size), A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]), ToTensorV2()]),\n        A.Compose([A.Resize(CFG.img_size, CFG.img_size), A.HorizontalFlip(p=1.0), A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]), ToTensorV2()]),\n        A.Compose([A.Resize(CFG.img_size, CFG.img_size), A.VerticalFlip(p=1.0), A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]), ToTensorV2()]),\n        A.Compose([A.Resize(CFG.img_size, CFG.img_size), A.Rotate(limit=(90,90), p=1.0), A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]), ToTensorV2()]),\n        A.Compose([A.Resize(CFG.img_size, CFG.img_size), A.RandomBrightnessContrast(brightness_limit=0.1, contrast_limit=0.1, p=1.0), A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]), ToTensorV2()]),\n    ]\n    return transforms\n\n# ============================================================================\n# MODEL (MATCHES TRAINING VERSION)\n# ============================================================================\nclass BiomassModel(nn.Module):\n    \"\"\"\n    Model architecture that mirrors the training model in your training script:\n    - timm backbone -> img features\n    - small tabular encoder (2 -> 128)\n    - fusion -> 200 -> head (5 outputs)\n    \"\"\"\n    def __init__(self, model_name, pretrained=False):\n        super().__init__()\n        self.backbone = timm.create_model(model_name, pretrained=pretrained, num_classes=0, global_pool='avg')\n\n        # get image feature size safely\n        img_features = getattr(self.backbone, 'num_features', None)\n        if img_features is None:\n            self.backbone.eval()\n            with torch.no_grad():\n                try:\n                    dummy = torch.randn(1, 3, min(224, CFG.img_size), min(224, CFG.img_size))\n                    feat = self.backbone(dummy)\n                    img_features = feat.shape[1]\n                except Exception:\n                    img_features = 1280  # reasonable default\n        self.img_features = img_features\n\n        # tabular encoder (same as training)\n        self.tabular_encoder = nn.Sequential(\n            nn.Linear(2, 64),\n            nn.BatchNorm1d(64),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n            nn.Linear(64, 128),\n            nn.BatchNorm1d(128),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n        )\n\n        fusion_dim = self.img_features + 128\n        self.fusion = nn.Sequential(\n            nn.Linear(fusion_dim, 200),\n            nn.BatchNorm1d(200),\n            nn.ReLU(),\n            nn.Dropout(0.4),\n            nn.Linear(200, 200),\n            nn.BatchNorm1d(200),\n            nn.ReLU(),\n            nn.Dropout(0.3),\n        )\n\n        # single head for 5 targets (matches training head)\n        self.head = nn.Linear(200, len(CFG.targets))\n\n    def forward(self, image, tabular):\n        img_feats = self.backbone(image)            # [B, img_features]\n        tab_feats = self.tabular_encoder(tabular)   # [B, 128]\n        combined = torch.cat([img_feats, tab_feats], dim=1)\n        fused = self.fusion(combined)\n        outputs = self.head(fused)\n        return outputs\n\n# ============================================================================\n# INFERENCE FUNCTIONS\n# ============================================================================\ndef predict_single_tta(model, dataset_obj, device, tta_transform):\n    dataset_tta = BiomassTestDataset(\n        dataset_obj.df,\n        dataset_obj.img_dir,\n        transform=tta_transform,\n        tabular_scaler=None,\n        has_metadata=(dataset_obj.tabular_features is not None)\n    )\n    dataset_tta.tabular_features = dataset_obj.tabular_features\n\n    loader = DataLoader(dataset_tta, batch_size=CFG.batch_size, shuffle=False, num_workers=CFG.num_workers, pin_memory=True)\n\n    model.eval()\n    predictions = []\n    with torch.no_grad():\n        for images, tabular in tqdm(loader, desc='tta_batch', leave=False):\n            images = images.to(device)\n            tabular = tabular.to(device)\n            outputs = model(images, tabular)\n            predictions.append(outputs.cpu().numpy())\n    if len(predictions) == 0:\n        return np.zeros((len(dataset_obj), len(CFG.targets)), dtype=np.float32)\n    return np.vstack(predictions)\n\ndef predict_with_tta(model, dataset_obj, device):\n    if not CFG.use_tta:\n        tta_transforms = [get_inference_transforms()]\n    else:\n        tta_transforms = get_tta_transforms()[:CFG.tta_steps]\n\n    all_tta_preds = []\n    for tta_idx, tta in enumerate(tta_transforms):\n        preds = predict_single_tta(model, dataset_obj, device, tta)\n        all_tta_preds.append(preds)\n    avg_preds = np.mean(all_tta_preds, axis=0)\n    return avg_preds\n\n# ============================================================================\n# CHECKPOINT LOADER (handles PyTorch 2.6+ safe-loading issue)\n# ============================================================================\ndef load_model_checkpoint(fold):\n    \"\"\"\n    Load checkpoint and return (model_state_dict, tabular_scaler, target_scaler)\n\n    This function first attempts standard torch.load(). If it fails due to\n    PyTorch safe-unpickling restrictions (e.g. because the checkpoint contains\n    sklearn objects like StandardScaler), it will attempt a secondary load that\n    allowlists StandardScaler. If that fails, it will try torch.load(..., weights_only=False).\n    \"\"\"\n    checkpoint_path = os.path.join(CFG.model_dir, f'best_model_fold{fold}.pth')\n    if not os.path.exists(checkpoint_path):\n        raise FileNotFoundError(f\"Model checkpoint not found: {checkpoint_path}\")\n\n    # Try the usual safe load first\n    try:\n        checkpoint = torch.load(checkpoint_path, map_location='cpu')\n    except Exception as e:\n        # If the error mentions \"StandardScaler\" or \"weights_only\", try allowlisting\n        print(\"⚠ Initial torch.load failed:\", repr(e))\n        print(\"  Attempting fallback load. Ensure you trust the checkpoint source.\")\n        checkpoint = None\n\n        # Try allowlisting sklearn.preprocessing.StandardScaler if available\n        try:\n            from sklearn.preprocessing import StandardScaler\n            # Prefer add_safe_globals if available (PyTorch >= 2.6)\n            if hasattr(torch.serialization, \"add_safe_globals\"):\n                print(\"  Using torch.serialization.add_safe_globals to allow StandardScaler...\")\n                torch.serialization.add_safe_globals([StandardScaler])\n                # re-run load (weights_only=False to allow non-tensor objects)\n                checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)\n            elif hasattr(torch.serialization, \"safe_globals\"):\n                # context manager style (if present)\n                print(\"  Using torch.serialization.safe_globals context manager to allow StandardScaler...\")\n                with torch.serialization.safe_globals([StandardScaler]):\n                    checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)\n            else:\n                # last resort: direct load with weights_only=False (less safe)\n                print(\"  add_safe_globals not found; falling back to torch.load(..., weights_only=False).\")\n                checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)\n        except Exception as e2:\n            # final fallback: try direct weights_only=False (if not already tried)\n            print(\"  Allowlist attempt failed:\", repr(e2))\n            print(\"  Final fallback: torch.load(..., weights_only=False). Only use if you trust the file.\")\n            checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)\n\n    # Interpret checkpoint formats\n    model_state = None\n    tabular_scaler = None\n    target_scaler = None\n\n    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:\n        model_state = checkpoint['model_state_dict']\n        tabular_scaler = checkpoint.get('tabular_scaler', None)\n        target_scaler = checkpoint.get('target_scaler', None)\n    elif isinstance(checkpoint, dict) and all(isinstance(v, torch.Tensor) for v in checkpoint.values()):\n        # likely raw state_dict\n        model_state = checkpoint\n    else:\n        # unknown object stored — maybe a full model object or custom wrapper\n        # Attempt to extract a state_dict attribute if present\n        if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:\n            model_state = checkpoint['state_dict']\n            tabular_scaler = checkpoint.get('tabular_scaler', None)\n            target_scaler = checkpoint.get('target_scaler', None)\n        else:\n            # Could be a saved full model or custom object; return it as-is (caller must handle)\n            model_state = checkpoint\n\n    return model_state, tabular_scaler, target_scaler\n\n# ============================================================================\n# MAIN INFERENCE PIPELINE\n# ============================================================================\ndef main():\n    print(\"=\"*70)\n    print(\"CSIRO BIOMASS PREDICTION - INFERENCE (Corrected with safe loader)\")\n    print(\"=\"*70)\n\n    # 1. Load test data\n    print(\"\\n[1/5] Loading test data...\")\n    test_df = pd.read_csv(CFG.test_csv)\n    print(f\"✓ Loaded {len(test_df)} test samples\")\n    print(f\"✓ Columns in test data: {test_df.columns.tolist()}\")\n\n    # Get unique images (test.csv may have multiple rows per image)\n    if 'image_path' in test_df.columns:\n        test_df_unique = test_df.drop_duplicates(subset=['image_path']).reset_index(drop=True)\n    elif 'image_id' in test_df.columns:\n        test_df_unique = test_df.drop_duplicates(subset=['image_id']).reset_index(drop=True)\n        test_df_unique['image_path'] = test_df_unique['image_id'].astype(str) + '.jpg'\n    else:\n        raise ValueError(\"Test CSV must contain either 'image_path' or 'image_id' column.\")\n\n    print(f\"✓ Found {len(test_df_unique)} unique images\")\n\n    # Check if metadata features are available\n    has_metadata = ('Pre_GSHH_NDVI' in test_df_unique.columns) and ('Height_Ave_cm' in test_df_unique.columns)\n\n    if not has_metadata:\n        test_df_unique['Pre_GSHH_NDVI'] = 0.0\n        test_df_unique['Height_Ave_cm'] = 0.0\n        print(\"⚠ Test data does NOT have metadata features; zeros will be used for tabular inputs\")\n\n    # 2. Verify model checkpoints exist\n    print(\"\\n[2/5] Checking model checkpoints...\")\n    available_folds = []\n    for fold in range(CFG.n_folds):\n        checkpoint_path = os.path.join(CFG.model_dir, f'best_model_fold{fold}.pth')\n        if os.path.exists(checkpoint_path):\n            available_folds.append(fold)\n            print(f\"✓ Found checkpoint for fold {fold}\")\n        else:\n            print(f\"✗ Missing checkpoint for fold {fold}\")\n\n    if len(available_folds) == 0:\n        raise FileNotFoundError(\"No model checkpoints found! Train models first.\")\n\n    print(f\"\\n✓ Will use {len(available_folds)} models for ensemble\")\n\n    # 3. Generate predictions from each fold\n    print(\"\\n[3/5] Generating predictions...\")\n    all_fold_predictions = []\n\n    for fold in available_folds:\n        print(f\"\\n--- Fold {fold} ---\")\n        model_state, tabular_scaler, target_scaler = load_model_checkpoint(fold)\n\n        # instantiate model (architecture must match training)\n        model = BiomassModel(CFG.model_name, pretrained=False)\n\n        # When model_state is a raw state_dict, load it; if it's a full model object, handle accordingly.\n        try:\n            if isinstance(model_state, dict) and all(isinstance(v, torch.Tensor) for v in model_state.values()):\n                model.load_state_dict(model_state)\n            elif isinstance(model_state, dict) and 'state_dict' in model_state:\n                model.load_state_dict(model_state['state_dict'])\n            else:\n                # model_state might be a saved full model object: try to load weights if possible\n                try:\n                    # Attempt to treat model_state as a full model and extract its state_dict attribute\n                    if hasattr(model_state, 'state_dict'):\n                        model.load_state_dict(model_state.state_dict())\n                    else:\n                        # last resort: attempt non-strict load (this may succeed if keys mostly match)\n                        model.load_state_dict(model_state, strict=False)\n                except Exception:\n                    # If it's truly a full model object, we can replace `model` with it (but move to device)\n                    if isinstance(model_state, nn.Module):\n                        model = model_state\n                    else:\n                        print(\"  Warning: unexpected checkpoint format; attempting non-strict load.\")\n                        model.load_state_dict(model_state, strict=False)\n        except Exception as e:\n            print(\"  Error loading state_dict into model:\", repr(e))\n            # try non-strict load\n            try:\n                model.load_state_dict(model_state, strict=False)\n                print(\"  Loaded with strict=False\")\n            except Exception as e2:\n                raise RuntimeError(f\"Failed to load model weights for fold {fold}: {e2}\")\n\n        model.to(CFG.device)\n        model.eval()\n        print(\"✓ Model loaded and moved to device\")\n\n        # create test dataset\n        test_dataset = BiomassTestDataset(\n            test_df_unique,\n            CFG.test_dir,\n            transform=get_inference_transforms(),\n            tabular_scaler=tabular_scaler,\n            has_metadata=has_metadata\n        )\n\n        # make predictions with TTA\n        fold_preds = predict_with_tta(model, test_dataset, CFG.device)\n\n        # inverse-transform predictions if a target scaler exists\n        if target_scaler is not None:\n            try:\n                fold_preds = target_scaler.inverse_transform(fold_preds)\n                print(\"✓ Predictions inverse-transformed to original target scale\")\n            except Exception:\n                print(\"⚠ target_scaler provided but inverse_transform failed; keeping scaled predictions\")\n\n        all_fold_predictions.append(fold_preds)\n        print(f\"✓ Fold {fold} predictions shape: {fold_preds.shape}\")\n        print(f\"  sample prediction (first): {fold_preds[0].tolist()}\")\n\n        # cleanup\n        del model\n        torch.cuda.empty_cache()\n\n    # 4. Ensemble predictions\n    print(\"\\n[4/5] Ensembling predictions from all folds...\")\n    final_predictions = np.mean(all_fold_predictions, axis=0)\n    print(f\"✓ Final predictions shape: {final_predictions.shape}\")\n    for idx, tname in enumerate(CFG.targets):\n        preds = final_predictions[:, idx]\n        print(f\"  {tname}: min={preds.min():.2f}, max={preds.max():.2f}, mean={preds.mean():.2f}, std={preds.std():.2f}\")\n\n    # 5. Create submission file\n    print(\"\\n[5/5] Creating submission file...\")\n    submission_rows = []\n    for idx, row in test_df_unique.iterrows():\n        image_path = row.get('image_path', None)\n        if image_path is None:\n            image_id = str(row.get('image_id', idx))\n        else:\n            image_id = os.path.basename(image_path).replace('.jpg', '')\n\n        for target_idx, target_name in enumerate(CFG.targets):\n            sample_id = f\"{image_id}__{target_name}\"\n            prediction = float(final_predictions[idx, target_idx])\n            prediction = max(0.0, prediction)\n            submission_rows.append({'sample_id': sample_id, 'target': prediction})\n\n    submission_df = pd.DataFrame(submission_rows)\n    submission_df.to_csv(CFG.output_file, index=False)\n\n    print(f\"\\n{'='*70}\")\n    print(\"✓ INFERENCE COMPLETE!\")\n    print(f\"✓ Submission file saved: {CFG.output_file}\")\n    print(f\"✓ Total predictions: {len(submission_df)} (expected {len(test_df_unique) * len(CFG.targets)})\")\n    print(\"\\nFirst few rows:\")\n    print(submission_df.head(10))\n    print(\"\\nLast few rows:\")\n    print(submission_df.tail(5))\n\n    # basic validation\n    expected_rows = len(test_df_unique) * len(CFG.targets)\n    if len(submission_df) == expected_rows:\n        print(f\"✓ Row count correct: {len(submission_df)} rows\")\n    else:\n        print(f\"⚠ Warning: Expected {expected_rows} rows, got {len(submission_df)}\")\n\n    if submission_df['target'].isna().sum() > 0:\n        print(f\"⚠ Warning: {submission_df['target'].isna().sum()} NaN values found\")\n    else:\n        print(\"✓ No NaN values\")\n\n    if (submission_df['target'] < 0).sum() > 0:\n        print(f\"⚠ Warning: {(submission_df['target'] < 0).sum()} negative values found\")\n    else:\n        print(\"✓ No negative values\")\n\n    print(\"\\n✓ Ready for submission!\")\n\nif __name__ == '__main__':\n    main()\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-05T06:26:51.290042Z","iopub.execute_input":"2025-11-05T06:26:51.290428Z","iopub.status.idle":"2025-11-05T06:27:13.835480Z","shell.execute_reply.started":"2025-11-05T06:26:51.290403Z","shell.execute_reply":"2025-11-05T06:27:13.834682Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}