{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"},{"sourceId":13900620,"sourceType":"datasetVersion","datasetId":8856212}],"dockerImageVersionId":31153,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport cv2\nimport timm\nimport torch\nimport numpy as np\nimport pandas as pd\nimport torch.nn as nn\nimport albumentations as A\nimport torch.nn.functional as F\n\nfrom tqdm import tqdm\nfrom albumentations.pytorch import ToTensorV2\nfrom torch.utils.data import Dataset, DataLoader\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nbatch_size = 4\ndropout = 0.30\nhidden_ratio = 0.25\ngrid = (2, 2)\n\ntest_csv = \"/kaggle/input/csiro-biomass/test.csv\"\ntest_img_dir = \"/kaggle/input/csiro-biomass/test\"\n\nmodel_dir = \"/kaggle/input/csiro-mvp-models\"\nmodel_paths = [os.path.join(model_dir, f\"model{i}.pth\") for i in range(1, 11)]\nexisting_models = [path for path in model_paths if os.path.exists(path)]\n\nCKPTS_A = existing_models[:5]\nCKPTS_B = existing_models[5:]\nW_A, W_B = 0.95, 0.075\n\ntargets = [\"Dry_Green_g\", \"Dry_Dead_g\", \"Dry_Clover_g\", \"GDM_g\", \"Dry_Total_g\"]\n\nbackbones = [\n    \"vit_base_patch14_reg4_dinov2\",\n    \"vit_base_patch14_reg4_dinov3\",\n    \"vit_base_patch14_dinov3\",\n]\n\nclass BiomassDataset(Dataset):\n    def __init__(self, df, transform, img_dir):\n        self.df = df.reset_index(drop=True)\n        self.transform = transform\n        self.img_dir = img_dir\n        self.paths = self.df[\"image_path\"].values\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        filename = os.path.basename(self.paths[idx])\n        full_path = os.path.join(self.img_dir, filename)\n        img = cv2.imread(full_path)\n        \n        if img is None:\n            img = np.zeros((1000, 2000, 3), dtype=np.uint8)\n        else:\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n        h, w, _ = img.shape\n        mid = w // 2\n        left_img = img[:, :mid]\n        right_img = img[:, mid:]\n\n        left_tensor = self.transform(image=left_img)[\"image\"]\n        right_tensor = self.transform(image=right_img)[\"image\"]\n        \n        return left_tensor, right_tensor\n\ndef get_input_size(model):\n    if hasattr(model, \"patch_embed\") and hasattr(model.patch_embed, \"img_size\"):\n        size = model.patch_embed.img_size\n        return int(size if isinstance(size, (int, float)) else size[0])\n    \n    if hasattr(model, \"img_size\"):\n        size = model.img_size\n        return int(size if isinstance(size, (int, float)) else size[0])\n    \n    cfg = getattr(model, \"default_cfg\", {}) or {}\n    input_size = cfg.get(\"input_size\", None)\n    \n    if input_size:\n        if isinstance(input_size, (tuple, list)) and len(input_size) >= 2:\n            return int(input_size[1])\n        return int(input_size if isinstance(input_size, (int, float)) else 224)\n    \n    arch = cfg.get(\"architecture\", \"\") or str(type(model))\n    return 518 if \"dinov2\" in arch.lower() or \"dinov3\" in arch.lower() else 224\n\ndef build_backbone(name):\n    model = timm.create_model(name, pretrained=False, num_classes=0)\n    features = model.num_features\n    input_size = get_input_size(model)\n    return model, features, input_size\n\nclass BaseDINO(nn.Module):\n    def __init__(self, backbone_name):\n        super().__init__()\n        self.backbone, feat_dim, input_size = build_backbone(backbone_name)\n        self.input_size = int(input_size)\n        self.feat_dim = feat_dim\n        self.combined_dim = feat_dim * 2\n        hidden_size = max(8, int(self.combined_dim * hidden_ratio))\n\n        def make_head():\n            return nn.Sequential(\n                nn.Linear(self.combined_dim, hidden_size),\n                nn.ReLU(inplace=True),\n                nn.Dropout(dropout),\n                nn.Linear(hidden_size, 1),\n            )\n\n        self.head_green = make_head()\n        self.head_clover = make_head()\n        self.head_dead = make_head()\n        self.softplus = nn.Softplus(beta=1.0)\n\n    def merge_features(self, left_feat, right_feat):\n        combined = torch.cat([left_feat, right_feat], dim=1)\n        green = self.softplus(self.head_green(combined))\n        clover = self.softplus(self.head_clover(combined))\n        dead = self.softplus(self.head_dead(combined))\n        gdm = green + clover\n        total = gdm + dead\n        return total, gdm, green\n\nclass PlainDINO(BaseDINO):\n    def forward(self, left_img, right_img):\n        left_feat = self.backbone(left_img)\n        right_feat = self.backbone(right_img)\n        return self.merge_features(left_feat, right_feat)\n\ndef split_dimension(length, parts):\n    step = length // parts\n    segments = []\n    start = 0\n    \n    for _ in range(parts - 1):\n        segments.append((start, start + step))\n        start += step\n    \n    segments.append((start, length))\n    return segments\n\nclass TiledDINO(BaseDINO):\n    def __init__(self, backbone_name):\n        super().__init__(backbone_name)\n        self.grid = grid\n\n    def process_tiles(self, x):\n        B, C, H, W = x.shape\n        rows, cols = self.grid\n        row_segments = split_dimension(H, rows)\n        col_segments = split_dimension(W, cols)\n        features = []\n        \n        for (rs, re) in row_segments:\n            for (cs, ce) in col_segments:\n                tile = x[:, :, rs:re, cs:ce]\n                if tile.shape[-2:] != (self.input_size, self.input_size):\n                    tile = F.interpolate(tile, size=(self.input_size, self.input_size), mode=\"bilinear\")\n                feat = self.backbone(tile)\n                features.append(feat)\n        \n        features = torch.stack(features, dim=0).permute(1, 0, 2)\n        return features.mean(dim=1)\n\n    def forward(self, left_img, right_img):\n        left_feat = self.process_tiles(left_img)\n        right_feat = self.process_tiles(right_img)\n        return self.merge_features(left_feat, right_feat)\n\nclass FiLM(nn.Module):\n    def __init__(self, feat_dim):\n        super().__init__()\n        hidden = max(32, feat_dim // 2)\n        self.mlp = nn.Sequential(\n            nn.Linear(feat_dim, hidden), nn.ReLU(inplace=True), nn.Linear(hidden, feat_dim * 2)\n        )\n\n    def forward(self, context):\n        gamma_beta = self.mlp(context)\n        return torch.chunk(gamma_beta, 2, dim=1)\n\nclass TiledFiLMDINO(BaseDINO):\n    def __init__(self, backbone_name):\n        super().__init__(backbone_name)\n        self.grid = grid\n        self.film_left = FiLM(self.feat_dim)\n        self.film_right = FiLM(self.feat_dim)\n\n    def extract_tile_features(self, x):\n        B, C, H, W = x.shape\n        rows, cols = self.grid\n        row_segments = split_dimension(H, rows)\n        col_segments = split_dimension(W, cols)\n        features = []\n        \n        for (rs, re) in row_segments:\n            for (cs, ce) in col_segments:\n                tile = x[:, :, rs:re, cs:ce]\n                if tile.shape[-2:] != (self.input_size, self.input_size):\n                    tile = F.interpolate(tile, size=(self.input_size, self.input_size), mode=\"bilinear\")\n                feat = self.backbone(tile)\n                features.append(feat)\n        \n        return torch.stack(features, dim=0).permute(1, 0, 2)\n\n    def process_stream(self, x, film_layer):\n        tiles = self.extract_tile_features(x)\n        context = tiles.mean(dim=1)\n        gamma, beta = film_layer(context)\n        modulated = tiles * (1 + gamma.unsqueeze(1)) + beta.unsqueeze(1)\n        return modulated.mean(dim=1)\n\n    def forward(self, left_img, right_img):\n        left_feat = self.process_stream(left_img, self.film_left)\n        right_feat = self.process_stream(right_img, self.film_right)\n        return self.merge_features(left_feat, right_feat)\n\ndef clean_state_dict(state_dict):\n    if not state_dict:\n        return state_dict\n    \n    cleaned_dict = {}\n    \n    for k, v in state_dict.items():\n        if k.startswith(\"module.\"):\n            k = k[7:]\n        \n        if k.startswith(\"student.\"):\n            k = k[8:]\n        \n        skip_prefixes = (\"txt_enc.\", \"img_proj.\", \"txt_film\", \"teacher.\", \"momentum_teacher.\")\n        if any(k.startswith(prefix) for prefix in skip_prefixes):\n            continue\n            \n        cleaned_dict[k] = v\n    \n    return cleaned_dict\n\ndef load_model(checkpoint_path):\n    print(f\"Loading model: {checkpoint_path}\")\n    \n    if not os.path.exists(checkpoint_path):\n        raise FileNotFoundError(f\"Checkpoint not found: {checkpoint_path}\")\n    \n    try:\n        raw_state = torch.load(checkpoint_path, map_location=device, weights_only=False)\n    except Exception as e:\n        print(f\"Error loading {checkpoint_path}: {e}\")\n        return None\n    \n    if isinstance(raw_state, dict):\n        if 'state_dict' in raw_state:\n            state_dict = raw_state['state_dict']\n        elif 'model' in raw_state:\n            state_dict = raw_state['model']\n        else:\n            state_dict = raw_state\n    else:\n        state_dict = raw_state\n    \n    state_dict = clean_state_dict(state_dict)\n    \n    if not state_dict:\n        print(f\"Empty state_dict for {checkpoint_path}\")\n        return None\n    \n    for backbone in backbones:\n        try:\n            model = TiledFiLMDINO(backbone)\n            result = model.load_state_dict(state_dict, strict=False)\n            \n            missing = [k for k in result.missing_keys if not k.startswith('backbone.pos_embed')]\n            \n            if len(missing) == 0:\n                print(f\"Loaded with backbone: {backbone}\")\n                model.to(device)\n                model.eval()\n                return model\n                \n        except Exception as e:\n            continue\n    \n    return None\n\ndef get_tta_transforms(img_size):\n    norm = [A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2()]\n    return [\n        A.Compose([A.Resize(img_size, img_size), *norm]),\n        A.Compose([A.HorizontalFlip(p=1.0), A.Resize(img_size, img_size), *norm]),\n        A.Compose([A.VerticalFlip(p=1.0), A.Resize(img_size, img_size), *norm]),\n        A.Compose([A.RandomRotate90(p=1.0), A.Resize(img_size, img_size), *norm]),\n    ]\n\n@torch.no_grad()\ndef predict_one_view(models, loader):\n    preds = []\n    use_amp = device.type == \"cuda\"\n    \n    for left_imgs, right_imgs in tqdm(loader, desc=\"Infer\", leave=False):\n        left_imgs = left_imgs.to(device, non_blocking=True)\n        right_imgs = right_imgs.to(device, non_blocking=True)\n        batch_preds = []\n        \n        with torch.amp.autocast(\"cuda\", enabled=use_amp):\n            for model in models:\n                total, gdm, green = model(left_imgs, right_imgs)\n                dead = torch.clamp(total - gdm, min=0.0)\n                clover = torch.clamp(gdm - green, min=0.0)\n                pred = torch.cat([green, dead, clover, gdm, total], dim=1)\n                batch_preds.append(pred.clamp(0.05, 400.0).cpu())\n        \n        preds.append(torch.stack(batch_preds).mean(dim=0).numpy())\n    \n    return np.concatenate(preds)\n\ndef run_inference_for_ckpts(checkpoint_paths, df, img_dir):\n    models = []\n    \n    for ckpt_path in checkpoint_paths:\n        model = load_model(ckpt_path)\n        if model is not None:\n            models.append(model)\n    \n    if not models:\n        raise ValueError(f\"No models loaded from {checkpoint_paths}\")\n    \n    input_size = models[0].input_size\n    \n    tta_preds = []\n    for transform in get_tta_transforms(input_size):\n        ds = BiomassDataset(df, transform, img_dir)\n        dl = DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=0, pin_memory=True)\n        tta_preds.append(predict_one_view(models, dl))\n    \n    return np.mean(tta_preds, axis=0)\n\ndef create_submission(final_preds, test_df, unique_df):\n    cols = [\"Dry_Green_g\", \"Dry_Dead_g\", \"Dry_Clover_g\", \"GDM_g\", \"Dry_Total_g\"]\n    \n    wide = pd.DataFrame({\"image_path\": unique_df[\"image_path\"]})\n    for i, col in enumerate(cols):\n        wide[col] = np.clip(final_preds[:, i], 0.05, 400.0)\n    \n    wide[\"GDM_g\"] = wide[\"Dry_Green_g\"] + wide[\"Dry_Clover_g\"]\n    wide[\"Dry_Total_g\"] = wide[\"GDM_g\"] + wide[\"Dry_Dead_g\"]\n    \n    wide[cols] = wide[cols].clip(0.05, 400.0)\n    \n    long_df = wide.melt(\n        id_vars=\"image_path\", \n        value_vars=cols, \n        var_name=\"target_name\", \n        value_name=\"target\"\n    )\n    \n    sub = test_df[[\"sample_id\", \"image_path\", \"target_name\"]].merge(\n        long_df, on=[\"image_path\", \"target_name\"], how=\"left\"\n    )[[\"sample_id\", \"target\"]]\n    \n    sub.to_csv(\"submission.csv\", index=False)\n    \n    return sub\n\ndef main():\n    test_df = pd.read_csv(test_csv)\n    unique_df = test_df.drop_duplicates(\"image_path\").reset_index(drop=True)\n    \n    pred_a = run_inference_for_ckpts(CKPTS_A, unique_df, test_img_dir)\n    pred_b = run_inference_for_ckpts(CKPTS_B, unique_df, test_img_dir)\n    \n    final_preds = W_A * pred_a + W_B * pred_b\n    \n    submission = create_submission(final_preds, test_df, unique_df)\n    \n    return submission\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T14:28:32.290199Z","iopub.execute_input":"2025-12-02T14:28:32.290480Z","iopub.status.idle":"2025-12-02T14:30:06.189678Z","shell.execute_reply.started":"2025-12-02T14:28:32.290459Z","shell.execute_reply":"2025-12-02T14:30:06.188868Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pd.read_csv('/kaggle/working/submission.csv').head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T14:30:06.190943Z","iopub.execute_input":"2025-12-02T14:30:06.191236Z","iopub.status.idle":"2025-12-02T14:30:06.208020Z","shell.execute_reply.started":"2025-12-02T14:30:06.191218Z","shell.execute_reply":"2025-12-02T14:30:06.207330Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}