{"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,"sourceType":"competition"},{"sourceId":4537,"sourceType":"modelInstanceVersion","modelInstanceId":3329,"modelId":986}],"dockerImageVersionId":31193,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as transforms\nfrom PIL import Image\nfrom tqdm import tqdm\nimport sys\nfrom typing import Dict, List, Tuple, Optional\n\n# Bypass transformers torch version check\nimport transformers.utils.import_utils\ntransformers.utils.import_utils.check_torch_load_is_safe = lambda: None\n\n# Kaggle paths (match training config)\nmodel_path = '/kaggle/input/dinov2/pytorch/giant/1'\nfrom transformers import Dinov2Model\n\n# Configuration Class (exact match with training)\nclass CFG:\n    seed = 42\n    batch_size = 32\n    img_size = 420\n    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n    root_dir = \"/kaggle/input/csiro-biomass\"  # Test images root\n    test_csv = \"/kaggle/input/csiro-biomass/test.csv\"  # Test CSV path\n    targets = [\"Dry_Clover_g\", \"Dry_Dead_g\", \"Dry_Green_g\", \"Dry_Total_g\", \"GDM_g\"]\n    MODEL_DIR = '/kaggle/input/trained-biomass-models/model'  # Trained models path\n    embed_dim = 1536  # DINOv2-giant embed dim (fixed)\n\n# Set seed\ndef seed_everything(seed: int) -> None:\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    np.random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\nseed_everything(CFG.seed)\n\n# ------------------------------\n# Model Architecture (exact match with training)\n# ------------------------------\nclass MultiHeadSelfAttention(nn.Module):\n    def __init__(self, dim: int, num_heads: int = 8, dropout: float = 0.1):\n        super().__init__()\n        assert dim % num_heads == 0, \"dim must be divisible by num_heads\"\n        \n        self.dim = dim\n        self.num_heads = num_heads\n        self.head_dim = dim // num_heads\n        self.scale = self.head_dim ** -0.5\n        self.dropout = nn.Dropout(dropout)\n        \n        self.qkv = nn.Linear(dim, dim * 3)\n        self.proj = nn.Linear(dim, dim)\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        B, N, C = x.shape\n        \n        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)\n        q, k, v = qkv.unbind(0)\n        \n        attn = (q @ k.transpose(-2, -1)) * self.scale\n        attn = attn.softmax(dim=-1)\n        attn = self.dropout(attn)\n        \n        x = (attn @ v).transpose(1, 2).reshape(B, N, C)\n        x = self.proj(x)\n        x = self.dropout(x)\n        \n        return x\n\nclass GatedResidualBlock(nn.Module):\n    def __init__(self, dim: int, expansion: int = 2, dropout: float = 0.2):\n        super().__init__()\n        self.norm1 = nn.LayerNorm(dim)\n        self.linear1 = nn.Linear(dim, dim * expansion)\n        self.act = nn.GELU()\n        self.dropout1 = nn.Dropout(dropout)\n        \n        self.norm2 = nn.LayerNorm(dim * expansion)\n        self.linear2 = nn.Linear(dim * expansion, dim)\n        self.dropout2 = nn.Dropout(dropout)\n        \n        self.gate = nn.Sequential(\n            nn.Linear(dim, dim),\n            nn.Sigmoid()\n        )\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        residual = x\n        x = self.norm1(x)\n        x = self.linear1(x)\n        x = self.act(x)\n        x = self.dropout1(x)\n        \n        x = self.norm2(x)\n        x = self.linear2(x)\n        x = self.dropout2(x)\n        \n        gate_weight = self.gate(residual)\n        x = gate_weight * x + (1 - gate_weight) * residual\n        return x\n\nclass FeatureAggregation(nn.Module):\n    def __init__(self, dim: int):\n        super().__init__()\n        self.attn_pool = nn.Sequential(\n            nn.Linear(dim, dim),\n            nn.Tanh(),\n            nn.Linear(dim, 1)\n        )\n        self.conv_pool = nn.Conv1d(dim, dim, kernel_size=3, padding=1)\n        self.norm = nn.LayerNorm(dim * 3)\n        self.proj = nn.Linear(dim * 3, dim)\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        B, N, C = x.shape\n        \n        attn_scores = self.attn_pool(x)\n        attn_weights = F.softmax(attn_scores, dim=1)\n        attn_feat = torch.sum(attn_weights * x, dim=1)\n        \n        max_feat = torch.max(x, dim=1)[0]\n        \n        conv_feat = self.conv_pool(x.transpose(1, 2)).transpose(1, 2)\n        conv_feat = torch.mean(conv_feat, dim=1)\n        \n        agg_feat = torch.cat([attn_feat, max_feat, conv_feat], dim=1)\n        agg_feat = self.norm(agg_feat)\n        agg_feat = self.proj(agg_feat)\n        \n        return agg_feat\n\nclass EnhancedBiomassMLP(nn.Module):\n    def __init__(\n        self, \n        input_dim: int = 1536, \n        hidden_dim: int = 768,\n        num_heads: int = 8,\n        num_blocks: int = 3,\n        dropout: float = 0.2\n    ):\n        super().__init__()\n        \n        self.input_proj = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.LayerNorm(hidden_dim),\n            nn.GELU(),\n            nn.Dropout(dropout)\n        )\n        self.noise_scale = nn.Parameter(torch.tensor(0.01))\n        \n        self.self_attn = MultiHeadSelfAttention(hidden_dim, num_heads, dropout)\n        \n        self.res_blocks = nn.Sequential(\n            *[GatedResidualBlock(hidden_dim, expansion=2, dropout=dropout) \n              for _ in range(num_blocks)]\n        )\n        \n        self.aggregation = FeatureAggregation(hidden_dim)\n        \n        self.output_head = nn.Sequential(\n            nn.Linear(hidden_dim, hidden_dim // 2),\n            nn.LayerNorm(hidden_dim // 2),\n            nn.GELU(),\n            nn.Dropout(dropout * 0.5),\n            nn.Linear(hidden_dim // 2, 1)\n        )\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        B, N, C = x.shape\n        \n        # No noise in inference\n        x = self.input_proj(x)\n        x_attn = self.self_attn(x)\n        x = x + x_attn\n        x = self.res_blocks(x)\n        feat = self.aggregation(x)\n        out = self.output_head(feat)\n        \n        return out.squeeze(-1)\n\n# ------------------------------\n# Utility Classes (match training)\n# ------------------------------\nclass ImageFeatureExtractor:\n    def __init__(self, config: CFG, model_path: str = None):\n        self.config = config\n        self.mean = [0.485, 0.456, 0.406]\n        self.std = [0.229, 0.224, 0.225]\n        self.transform = self._build_transform()\n        self.model = self._load_dinov2(model_path or 'facebook/dinov2-giant')\n        self.embed_dim = self.model.config.hidden_size\n\n    def _build_transform(self) -> transforms.Compose:\n        return transforms.Compose([\n            transforms.ToTensor(),\n            transforms.Resize((self.config.img_size, self.config.img_size)),\n            transforms.Normalize(self.mean, self.std)\n        ])\n\n    def _load_dinov2(self, model_name_or_path: str) -> Dinov2Model:\n        print(\"Loading DINOv2 Model...\")\n        try:\n            model = Dinov2Model.from_pretrained(model_name_or_path)\n        except OSError:\n            print(f\"Failed to load from {model_name_or_path}, trying to download...\")\n            model = Dinov2Model.from_pretrained('facebook/dinov2-giant')\n        \n        model.head = nn.Identity()\n        model.to(self.config.device)\n        model.eval()\n        return model\n\n    def extract_test_features(self, df: pd.DataFrame) -> Dict[str, torch.Tensor]:\n        \"\"\"Extract features for test set (match training logic)\"\"\"\n        class TestImageDataset(Dataset):\n            def __init__(self, df, root_dir, transform):\n                self.df = df\n                self.root_dir = root_dir\n                self.transform = transform\n            def __len__(self): return len(self.df)\n            def __getitem__(self, idx):\n                row = self.df.iloc[idx]\n                img_path = os.path.join(self.root_dir, row['image_path'])\n                try:\n                    image = Image.open(img_path).convert('RGB')\n                except:\n                    image = Image.new('RGB', (self.config.img_size, self.config.img_size))\n                return self.transform(image), row['sample_id']\n\n        batch_size = 8 if self.config.img_size <= 500 else 4\n        loader = DataLoader(\n            TestImageDataset(df, self.config.root_dir, self.transform),\n            batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True\n        )\n\n        feature_cache = {}\n        with torch.no_grad():\n            for imgs, ids in tqdm(loader, desc=\"Extracting Test Features\"):\n                imgs = imgs.to(self.config.device, non_blocking=True)\n                outputs = self.model(imgs)\n                dense_features = outputs.last_hidden_state[:, 1:, :].cpu()\n                for i, sample_id in enumerate(ids):\n                    feature_cache[sample_id] = dense_features[i].clone()\n        \n        del self.model\n        torch.cuda.empty_cache()\n        print(\"Test feature extraction complete.\")\n        return feature_cache\n\nclass TestFeatureDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, features_dict: Dict[str, torch.Tensor], input_size: int):\n        self.df = df\n        self.features_dict = features_dict\n        self.input_size = input_size\n\n    def __len__(self) -> int:\n        return len(self.df)\n\n    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, str]:\n        row = self.df.iloc[idx]\n        sample_id = row['sample_id']\n        feat = self.features_dict.get(sample_id, torch.zeros((1, self.input_size)))\n        return feat, sample_id\n\n# ------------------------------\n# Inference Functions\n# ------------------------------\ndef load_best_model(target_name: str, config: CFG) -> nn.Module:\n    \"\"\"Load best model for a target (auto-find best R2 model)\"\"\"\n    model_dir = os.path.join(config.MODEL_DIR, target_name)\n    if not os.path.exists(model_dir):\n        raise FileNotFoundError(f\"Model directory not found: {model_dir}\")\n    \n    # Find best model (highest R2 in filename)\n    model_files = [f for f in os.listdir(model_dir) if f.endswith('.pt')]\n    if not model_files:\n        raise FileNotFoundError(f\"No model files found in {model_dir}\")\n    \n    # Parse R2 from filename (format: foldX_R2_0.XXXX.pt)\n    def get_r2_from_filename(filename: str) -> float:\n        try:\n            return float(filename.split('_R2_')[1].split('.pt')[0])\n        except:\n            return -float('inf')\n    \n    best_model_file = max(model_files, key=get_r2_from_filename)\n    best_model_path = os.path.join(model_dir, best_model_file)\n    print(f\"Loading best model for {target_name}: {best_model_file}\")\n    \n    # Initialize model (exact same config as training)\n    model = EnhancedBiomassMLP(\n        input_dim=config.embed_dim,\n        hidden_dim=768,\n        num_heads=8,\n        num_blocks=3,\n        dropout=0.2\n    ).to(config.device)\n    \n    # Load state dict\n    model.load_state_dict(torch.load(best_model_path, map_location=config.device))\n    model.eval()\n    return model\n\ndef inference_target(target_name: str, df: pd.DataFrame, feature_cache: Dict[str, torch.Tensor], config: CFG) -> pd.Series:\n    \"\"\"Inference for single target\"\"\"\n    # Create dataset and loader\n    dataset = TestFeatureDataset(df, feature_cache, config.embed_dim)\n    loader = DataLoader(\n        dataset, batch_size=config.batch_size, shuffle=False, num_workers=2, pin_memory=True\n    )\n    \n    # Load model\n    model = load_best_model(target_name, config)\n    \n    # Inference\n    predictions = {}\n    with torch.no_grad():\n        for features, sample_ids in tqdm(loader, desc=f\"Inferencing {target_name}\"):\n            features = features.to(config.device, non_blocking=True)\n            outputs = model(features)\n            for idx, sample_id in enumerate(sample_ids):\n                predictions[sample_id] = outputs[idx].cpu().item()\n    \n    # Convert to Series (match sample order in df)\n    return pd.Series([predictions.get(row['sample_id'], 0.0) for _, row in df.iterrows()], index=df.index)\n\ndef main():\n    # 1. Load test data\n    print(\"Loading test data...\")\n    df_test = pd.read_csv(CFG.test_csv)\n    print(f\"Test data shape: {df_test.shape}\")\n    \n    # 2. Extract test features (same as training)\n    feature_extractor = ImageFeatureExtractor(CFG, model_path)\n    test_feature_cache = feature_extractor.extract_test_features(df_test)\n    \n    # 3. Inference for all targets\n    print(\"\\nStarting inference for all targets...\")\n    submission_df = df_test[['sample_id']].copy()\n    \n    for target in CFG.targets:\n        submission_df[target] = inference_target(target, df_test, test_feature_cache, CFG)\n    \n    # 4. Save submission (Kaggle required format)\n    submission_path = \"/kaggle/working/submission.csv\"\n    submission_df.to_csv(submission_path, index=False)\n    print(f\"\\nSubmission saved to: {submission_path}\")\n    print(f\"Submission shape: {submission_df.shape}\")\n    print(\"\\nSample submission:\")\n    print(submission_df.head())\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}