{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"sourceType":"competition"},{"sourceId":14541877,"sourceType":"datasetVersion","datasetId":9287943,"isSourceIdPinned":true}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# %% [markdown]\n\n# %%\nimport sys\nsys.path.insert(0, '/kaggle/input/timm-1-0-20/timm-1.0.20/')\n\n# %%\nimport os\nimport random\n\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom PIL import Image\nimport albumentations as A\n\n# %%\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\nimport cv2\n\nimport timm\nprint(timm.__version__)\n\ndef _make_scratch(in_shape, out_shape, groups=1, expand=False):\n    scratch = nn.Module()\n    out_shape1 = out_shape\n    out_shape2 = out_shape\n    out_shape3 = out_shape\n    out_shape4 = out_shape\n    scratch.layer1_rn = nn.Conv2d(in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)\n    scratch.layer2_rn = nn.Conv2d(in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)\n    scratch.layer3_rn = nn.Conv2d(in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)\n    scratch.layer4_rn = nn.Conv2d(in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups)\n    return scratch\n\nclass DPTHead(nn.Module):\n    def __init__(\n        self, \n        nclass,\n        in_channels, \n        features=256, \n        use_bn=False, \n        out_channels=[256, 512, 1024],\n    ):\n        super(DPTHead, self).__init__()\n        self.projects = nn.ModuleList([\n            nn.Conv2d(\n                in_channels=in_channels,\n                out_channels=out_channel,\n                kernel_size=1,\n                stride=1,\n                padding=0,\n            ) for out_channel in out_channels\n        ])\n        \n        self.scratch = _make_scratch(\n            out_channels,\n            features,\n            groups=1,\n            expand=False,\n        )\n        self.scratch.stem_transpose = None\n        self.scratch.output_conv = nn.Conv2d(features*4, nclass, kernel_size=1, stride=1, padding=0)  \n        self.proj = nn.ConvTranspose2d(features, features, 4, stride=4, padding=0, bias=False)\n\n    def forward(self, out_features):\n        out = []\n        for i, x in enumerate(out_features):\n            x = self.projects[i](x)\n            out.append(x)\n        \n        layer_1, layer_2, layer_3, layer_4 = out\n        layer_1_rn = self.scratch.layer1_rn(layer_1)\n        layer_2_rn = self.scratch.layer2_rn(layer_2)\n        layer_3_rn = self.scratch.layer3_rn(layer_3)\n        layer_4_rn = self.scratch.layer4_rn(layer_4)\n        layer_1_rn = self.proj(layer_1_rn)\n        target_hw = layer_1_rn.shape[-2:]  \n        layer_2_up = F.interpolate(layer_2_rn, size=target_hw, mode=\"bilinear\", align_corners=True)\n        layer_3_up = F.interpolate(layer_3_rn, size=target_hw, mode=\"bilinear\", align_corners=True)\n        layer_4_up = F.interpolate(layer_4_rn, size=target_hw, mode=\"bilinear\", align_corners=True)\n        fused = torch.cat([layer_1_rn, layer_2_up, layer_3_up, layer_4_up], dim=1)\n        out = self.scratch.output_conv(fused)\n        return out\n\nclass Net(nn.Module):\n    def __init__(self, device_id):\n        super().__init__()\n        self.model = timm.create_model('vit_huge_plus_patch16_dinov3_qkvb.lvd1689m', pretrained=False)\n        self.feature_indices = [12, 19, 25, 31]\n\n        self.head = DPTHead(3, self.model.embed_dim, 128, False, out_channels=[96, 192, 384, 768])\n\n        self.IMAGENET_DEFAULT_MEAN = torch.tensor([0.485, 0.456, 0.406]).to(device_id)\n        self.IMAGENET_DEFAULT_STD = torch.tensor([0.229, 0.224, 0.225]).to(device_id)\n\n        self.global_pool = nn.AdaptiveAvgPool2d((1, 1))\n        \n    def forward(self, x):\n        w = x.shape[3]\n        x = x/255.0\n        x = (x - self.IMAGENET_DEFAULT_MEAN[None,:, None, None])/self.IMAGENET_DEFAULT_STD[None,:, None, None]\n        x_l = x[:,:,:,:w//2]\n        x_r = x[:,:,:,w//2:]\n        outputs_l, features_l = self.model.forward_intermediates(x_l, indices=self.feature_indices)\n        outputs_r, features_r = self.model.forward_intermediates(x_r, indices=self.feature_indices)\n        features_l = self.head(features_l)\n        features_r = self.head(features_r)\n        main_logits_l = self.global_pool(features_l).flatten(1)\n        main_logits_r = self.global_pool(features_r).flatten(1)\n        main_logits = main_logits_l + main_logits_r\n        \n        clover_logits = main_logits[:,0]\n        dead_logits = main_logits[:,1]\n        green_logits = main_logits[:,2]\n        \n        return clover_logits, dead_logits, green_logits\n\n# %%\n# ======== Dataset (Albumentations, matches pipeline preprocessing) ========\nclass InferenceDataset(Dataset):\n    def __init__(self, df: pd.DataFrame, images_root: str, img_size: int):\n        self.df = df.reset_index(drop=True)\n        self.images_root = images_root\n        self.transform_1024 = A.Compose([\n            A.Resize(height=1024, width=1024, interpolation=cv2.INTER_CUBIC, p=1.0)\n        ])\n\n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx: int):\n        row = self.df.iloc[idx]\n        img_path = os.path.join(self.images_root, row[\"image_path\"])\n        image = Image.open(img_path).convert(\"RGB\")\n        img = np.asarray(image)\n        img = cv2.resize(img, (2000, 1000), interpolation=cv2.INTER_CUBIC)\n        img_left = img[:, :1000, :]\n        img_right = img[:, 1000:, :]\n        img_left_1024 = self.transform_1024(image=img_left)[\"image\"]\n        img_right_1024 = self.transform_1024(image=img_right)[\"image\"]\n        img_1024 = np.concatenate([img_left_1024, img_right_1024], axis=1)\n        img_1024 = torch.from_numpy(img_1024).permute(2, 0, 1).float()\n        return img_1024\n\n# %%\n# ======== Model (use eval.py-style per-target invocation) ========\n\ndef _to_numpy(t: torch.Tensor) -> np.ndarray:\n    return t.detach().cpu().numpy()\n\ndef _predict_all_targets(model, images):\n    with torch.inference_mode():\n        autocast_ctx = (\n            torch.autocast(device_type=\"cuda\", dtype=torch.float16)\n        )\n        y_pred_all = []\n        with autocast_ctx:\n            output_clover, output_dead, output_green = model(images)\n        output_clover = _to_numpy(output_clover)\n        output_dead = _to_numpy(output_dead)\n        output_green = _to_numpy(output_green)\n        for i in range(len(output_clover)):\n            y_pred_all.append(np.array([output_clover[i], output_dead[i], output_green[i]]))\n        y_pred_all = np.vstack(y_pred_all)\n        return y_pred_all\n\n# %%\n# ======== TTA + Utils ========\n\ndef tta_inference(model, images):\n    \"\"\"Average predictions across simple flips using eval.py-style per-target invocation.\"\"\"\n    preds = _predict_all_targets(model, images)\n    # preds_lr = _predict_all_targets(model, torch.flip(images, dims=[3]))\n    preds_mean = preds / 1.0\n    return preds_mean\n\n\ndef get_id(x: str) -> str:\n    return x.split('_')[0]\n\n# %%\n# ======== Parameters & Data ========\n# Adjust these if your weights are in a different dataset path\nimg_size = 1024\nbatch_size = 2\nnum_workers = 1\nfolds = [0, 1, 2, 3, 4]\n# Target names as used in training/configs\ntarget_names = [\"Dry_Clover_g\", \"Dry_Dead_g\", \"Dry_Green_g\", \"Dry_Total_g\", \"GDM_g\"]\n\n# Input paths (update if needed)\ntest_csv_path = \"/kaggle/input/csiro-biomass/test.csv\"\nimages_root = \"/kaggle/input/csiro-biomass\"\nweights_dir_0 = \"/kaggle/input/csiro-pipeline-15\"  # Directory containing 5 fold weight files\n# Load test dataframe\ntest_df = pd.read_csv(test_csv_path)\n# # Ensure unique image_path rows and prepare sample_id\ntest_df = test_df[~test_df['image_path'].duplicated()][['sample_id', 'image_path']].reset_index(drop=True)\ntest_df['sample_id'] = test_df['sample_id'].apply(get_id)\n\n# Dataset & DataLoader\ndataset = InferenceDataset(test_df, images_root=images_root, img_size=img_size)\ndataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)\n\n# %%\ndef _to_numpy(t) -> np.ndarray:\n    return t.detach().cpu().numpy()\n\n# %%\n# ======== Inference (2-way split, 2 GPUs, multiprocessing) ========\nimport contextlib\nimport torch.multiprocessing as mp\n\n# Helper to find a plausible weight file per fold\ndef find_weight_for_fold(base_dir: str, fold: int):\n    for fname in os.listdir(base_dir):\n        if f\"fold_{fold}_stage_0_42.pt\" in fname:\n            return os.path.join(base_dir, fname)\n        elif f\"checkpoint_{fold}_stage_0_42_26.pt\" in fname:\n            return os.path.join(base_dir, fname)\n    return None\n\n\ndef run_inference_split(split_id, split_df, split_indices, device_id, out_queue):\n    if len(split_df) == 0:\n        out_queue.put(\n            {\n                \"split_id\": split_id,\n                \"indices\": split_indices,\n                \"results_dict_0\": None,\n            }\n        )\n        return\n    if torch.cuda.is_available():\n        print(f\"Using GPU {device_id}\")\n        torch.cuda.set_device(device_id)\n        device = torch.device(f\"cuda:{device_id}\")\n    else:\n        device = torch.device(\"cpu\")\n\n    # Dataset & DataLoader for this split\n    split_dataset = InferenceDataset(split_df, images_root=images_root, img_size=img_size)\n    split_loader = DataLoader(\n        split_dataset,\n        batch_size=batch_size,\n        shuffle=False,\n        num_workers=num_workers,\n        pin_memory=torch.cuda.is_available(),\n    )\n\n\n    results_dict_0 = {}\n    results_dict_1 = {}\n    results_dict_2 = {}\n\n    for fold in folds:\n        weight_path_0 = find_weight_for_fold(weights_dir_0, fold)\n\n        if weight_path_0 is not None:\n            model_0 = Net(str(device))\n            model_0.to(device)\n            model_0.load_state_dict(torch.load(weight_path_0, map_location=\"cpu\"), strict=True)\n            model_0.eval()\n\n        results_0 = []\n        with torch.inference_mode():\n            autocast_ctx = (\n                torch.autocast(device_type=\"cuda\", dtype=torch.float16)\n                if torch.cuda.is_available()\n                else contextlib.nullcontext()\n            )\n\n            for batch in split_loader:\n                img_1024 = batch\n                img_1024 = img_1024.to(device)\n                if weight_path_0 is not None:\n                    preds_0 = tta_inference(model_0, img_1024)\n                    results_0.append(preds_0)\n\n        if weight_path_0 is not None:\n            results_dict_0[fold] = np.concatenate(results_0)\n\n    out_queue.put(\n        {\n            \"split_id\": split_id,\n            \"indices\": split_indices,\n            \"results_dict_0\": results_dict_0\n        }\n    )\n\n\nnum_samples = len(test_df)\nall_indices = np.arange(num_samples)\nsplit_indices = np.array_split(all_indices, 2)\n\n# Spawn two processes, one per GPU\ntry:\n    mp.set_start_method(\"fork\", force=True)\n    ctx = mp.get_context(\"fork\")\nexcept RuntimeError:\n    # Fall back if start method was already set\n    ctx = mp.get_context(\"fork\" if \"fork\" in mp.get_all_start_methods() else \"spawn\")\n\nout_queue = ctx.Queue()\nprocesses = []\n\nfor split_id, indices in enumerate(split_indices):\n    split_df = test_df.iloc[indices]\n    gpu_id = split_id  # GPU 0 for split 0, GPU 1 for split 1\n    p = ctx.Process(\n        target=run_inference_split,\n        args=(split_id, split_df, indices, gpu_id, out_queue),\n    )\n    p.start()\n    processes.append(p)\n\n# Collect results\nsplit_outputs = []\nfor _ in processes:\n    split_outputs.append(out_queue.get())\n\nfor p in processes:\n    p.join()\n\n# Merge split outputs back to original order\nresults_dict_0 = {}\nstate_results = np.zeros((num_samples, 4), dtype=np.float32)\n\nfor split_out in split_outputs:\n    indices = split_out[\"indices\"]\n    if split_out[\"results_dict_0\"] is None:\n        continue\n    for fold, preds in split_out[\"results_dict_0\"].items():\n        if fold not in results_dict_0:\n            results_dict_0[fold] = np.zeros((num_samples, 3), dtype=preds.dtype)\n        results_dict_0[fold][indices] = preds\n\n\n# %%\n# ======== Ensemble across folds ========\npreds_0 = np.mean([results_dict_0[fold] for fold in results_dict_0.keys()], axis=0)  # (N, 3): clover, dead, green\n\npreds = preds_0\n\nclover = preds[:, 0].astype(np.float32)\ndead = preds[:, 1].astype(np.float32)\ngreen = preds[:, 2].astype(np.float32)\n\n# Keep original global clipping maxima\nclover = np.clip(clover, 0, 71.7865)\ndead = np.clip(dead, 0, 83.8407)\ngreen = np.clip(green, 0, 157.9836)\n\n\n# Derived aggregates\nDry_Total_g = dead + clover + green\nGDM_g = clover + green\n\npreds = np.concatenate(\n    [clover[:, None], dead[:, None], green[:, None], Dry_Total_g[:, None], GDM_g[:, None]], axis=1\n)\n\n# %%\nresult_df = pd.DataFrame(preds, columns=[\"Dry_Clover_g\", \"Dry_Dead_g\", \"Dry_Green_g\", \"Dry_Total_g\", \"GDM_g\"])\nresult_df['sample_id'] = test_df['sample_id']\nresult_df = pd.melt(result_df, id_vars='sample_id', value_vars=[\"Dry_Clover_g\", \"Dry_Dead_g\", \"Dry_Green_g\", \"Dry_Total_g\", \"GDM_g\"], value_name='target')\nresult_df['sample_id'] = result_df['sample_id'] + '__' + result_df['variable']\n\nresult_df[['sample_id', 'target']].to_csv('submission.csv', index=False)\nprint(result_df)\n\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}