{"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":"none","dataSources":[{"sourceId":112509,"databundleVersionId":14254895,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":634472,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":478366,"modelId":494244}],"dockerImageVersionId":31192,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# run_inference.py  (modified to use TwoStreamBiomassModel)\nimport os\nos.environ['CUDA_VISIBLE_DEVICES'] = '0'\nimport sys\nimport torch\nimport pandas as pd\nimport numpy as np\nfrom PIL import Image\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nfrom tqdm import tqdm\nfrom pathlib import Path\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport timm\n\n\n\nclass TwoStreamBiomassModel(nn.Module):\n\n    def __init__(self, \n                 backbone_name='convnext_tiny',\n                 pretrained=True,\n                 dropout=0.3):\n        \"\"\"\n        Args:\n            backbone_name: timm model name\n            pretrained: Use ImageNet pretrained weights\n            dropout: Dropout rate\n        \"\"\"\n        super().__init__()\n        \n        # Shared backbone\n        self.backbone = timm.create_model(\n            backbone_name,\n            pretrained=pretrained,\n            num_classes=0,\n            global_pool='avg'\n        )\n        \n        # Feature dimensions\n        self.feat_dim = self.backbone.num_features\n        self.combined_dim = self.feat_dim * 2\n        \n        print(f\"[Model V3] Backbone: {backbone_name}\")\n        print(f\"[Model V3] Feature dim: {self.feat_dim}\")\n        print(f\"[Model V3] Combined dim: {self.combined_dim}\")\n        \n        # SIMPLE heads like in example - NO complex activations\n        self.head_total = self._make_head(dropout)\n        self.head_gdm = self._make_head(dropout)\n        self.head_green = self._make_head(dropout)\n    \n    def _make_head(self, dropout):\n        \"\"\"Create simple MLP head - exactly like example\"\"\"\n        return nn.Sequential(\n            nn.Linear(self.combined_dim, self.combined_dim // 2),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(self.combined_dim // 2, 1)\n        )\n    \n    def forward(self, image):\n        \"\"\"\n        Simple forward pass - like example.\n        \n        Args:\n            image: [B, 3, H, W] stereo image (W = 2H)\n        \n        Returns:\n            dict with predictions (NO constraints, simple regression)\n        \"\"\"\n        B, C, H, W = image.shape\n        assert W == 2 * H, f\"Expected stereo image W=2H, got H={H}, W={W}\"\n        \n        # Split stereo image\n        mid_point = W // 2\n        img_left = image[:, :, :, :mid_point]\n        img_right = image[:, :, :, mid_point:]\n        \n        # Extract features\n        feat_left = self.backbone(img_left)\n        feat_right = self.backbone(img_right)\n        \n        # Concatenate\n        combined = torch.cat([feat_left, feat_right], dim=1)\n        \n        # Predict with each head (simple linear output)\n        out_total = self.head_total(combined).squeeze(-1)\n        out_gdm = self.head_gdm(combined).squeeze(-1)\n        out_green = self.head_green(combined).squeeze(-1)\n        \n        # Derive remaining targets (like example)\n        out_clover = torch.clamp(out_gdm - out_green, min=0)\n        out_dead = torch.clamp(out_total - out_gdm, min=0)\n        \n        return {\n            'Dry_Total_g': out_total,\n            'GDM_g': out_gdm,\n            'Dry_Green_g': out_green,\n            'Dry_Clover_g': out_clover,\n            'Dry_Dead_g': out_dead,\n        }\n\nclass TestDataset(Dataset):\n    def __init__(self, image_dir, transform=None, image_size=512):\n        self.image_dir = Path(image_dir)\n        self.transform = transform\n        self.image_size = image_size\n\n        self.image_files = sorted([\n            f for f in os.listdir(image_dir)\n            if f.lower().endswith(('.png', '.jpg', '.jpeg', '.tif', '.tiff'))\n        ])\n\n        # Default transform: resize to (H, W=2*H) for stereo model\n        if self.transform is None:\n            self.transform = transforms.Compose([\n                transforms.Resize((image_size, image_size * 2)),\n                transforms.ToTensor(),\n                transforms.Normalize(\n                    mean=[0.485, 0.456, 0.406],\n                    std=[0.229, 0.224, 0.225]\n                )\n            ])\n\n    def __len__(self):\n        return len(self.image_files)\n\n    def _ensure_stereo_image(self, img: Image.Image) -> Image.Image:\n\n        w, h = img.size\n        # Already stereo-shaped (approx)\n        if w >= 2 * h:\n            return img\n        # Square single view -> duplicate horizontally\n        if w == h:\n            new = Image.new('RGB', (w * 2, h))\n            new.paste(img, (0, 0))\n            new.paste(img, (w, 0))\n            return new\n\n        if w < h:\n            # scale image to height h and paste duplicated to width=2*h\n            img_resized = img.resize((h, h))\n            new = Image.new('RGB', (h * 2, h))\n            new.paste(img_resized, (0, 0))\n            new.paste(img_resized, (h, 0))\n            return new\n        return img\n\n    def __getitem__(self, idx):\n        image_file = self.image_files[idx]\n        image_path = self.image_dir / image_file\n\n        image = Image.open(image_path).convert('RGB')\n\n        # Ensure stereo format (width = 2 * height) or fall back to duplication\n        image = self._ensure_stereo_image(image)\n\n        if self.transform:\n            image = self.transform(image)\n\n        image_name = Path(image_file).stem\n        return image, image_name\n\n\ndef load_models(model_paths, device='cuda'):\n    models = []\n    for model_path in model_paths:\n        checkpoint = torch.load(model_path, map_location=device)\n        # retrieve training config if present\n        if 'config' in checkpoint:\n            training_config = checkpoint['config']\n            backbone = training_config.get('backbone', 'convnext_tiny')\n            dropout = training_config.get('dropout', 0.3)\n        else:\n            backbone = 'convnext_tiny'\n            dropout = 0.3\n\n        cfg = {\n            'backbone_name': backbone,\n            'pretrained': False,\n            'dropout': dropout\n        }\n\n        print(f\"Loading Two-Stream model from: {model_path}\")\n        print(f\"  - backbone_name: {cfg['backbone_name']}, dropout: {cfg['dropout']}\")\n\n        model = TwoStreamBiomassModel(\n            backbone_name=cfg['backbone_name'],\n            pretrained=cfg['pretrained'],\n            dropout=cfg['dropout']\n        )\n\n        # load state dict (handle both 'model_state_dict' wrappers and plain state_dict saves)\n        if 'model_state_dict' in checkpoint:\n            state_dict = checkpoint['model_state_dict']\n        else:\n            state_dict = checkpoint\n\n        # fix keys if saved from DataParallel (optional)\n        try:\n            model.load_state_dict(state_dict)\n        except RuntimeError:\n            # try removing 'module.' prefix if present in keys\n            new_state = {}\n            for k, v in state_dict.items():\n                new_key = k.replace('module.', '') if k.startswith('module.') else k\n                new_state[new_key] = v\n            model.load_state_dict(new_state)\n\n        model = model.to(device)\n        model.eval()\n        models.append(model)\n    return models\n\n\ndef predict_with_ensemble(models, dataloader, device='cuda'):\n    # target names expected by TwoStreamBiomassModel (V3)\n    target_names = ['Dry_Green_g', 'Dry_Dead_g', 'Dry_Clover_g', 'GDM_g', 'Dry_Total_g']\n    all_predictions = {name: [] for name in target_names}\n    all_image_names = []\n\n    with torch.no_grad():\n        for images, image_names in tqdm(dataloader, desc=\"Processing images\"):\n            images = images.to(device)\n            batch_size = images.size(0)\n\n            batch_preds = {name: torch.zeros(batch_size, device=device) for name in target_names}\n\n            for model in models:\n                preds = model(images)  # TwoStreamBiomassModel.forward returns dict\n                for target_name in target_names:\n                    batch_preds[target_name] += preds[target_name]\n\n            # average\n            for target_name in target_names:\n                batch_preds[target_name] /= len(models)\n\n            # collect\n            for target_name in target_names:\n                all_predictions[target_name].extend(batch_preds[target_name].cpu().numpy().tolist())\n\n            all_image_names.extend(image_names)\n\n    predictions_df = pd.DataFrame({'image_name': all_image_names})\n    for target_name in target_names:\n        predictions_df[target_name] = all_predictions[target_name]\n\n    return predictions_df\n\n\ndef create_submission(predictions_df, output_path='sample_submission.csv'):\n    # submission ordering (use same order you expect downstream)\n    target_names = ['Dry_Clover_g', 'Dry_Dead_g', 'Dry_Green_g', 'Dry_Total_g', 'GDM_g']\n\n    submission_rows = []\n    for _, row in predictions_df.iterrows():\n        image_name = row['image_name']\n        for target_name in target_names:\n            sample_id = f\"{image_name}__{target_name}\"\n            # ensure non-negative\n            target_value = float(max(0.0, row[target_name]))\n            submission_rows.append({'sample_id': sample_id, 'target': target_value})\n\n    submission_df = pd.DataFrame(submission_rows)\n    submission_df.to_csv(output_path, index=False)\n    return submission_df\n\n\ndef run_inference(test_image_dir, model_paths, output_path='sample_submission.csv',\n                  batch_size=16, image_size=512, num_workers=4, device='cuda'):\n\n    if device == 'cuda' and not torch.cuda.is_available():\n        device = 'cpu'\n\n    test_dataset = TestDataset(image_dir=test_image_dir, image_size=image_size)\n\n    test_loader = DataLoader(\n        test_dataset, batch_size=batch_size, shuffle=False,\n        num_workers=num_workers, pin_memory=True if device == 'cuda' else False\n    )\n\n    models = load_models(model_paths, device=device)\n    predictions_df = predict_with_ensemble(models, test_loader, device=device)\n    submission_df = create_submission(predictions_df, output_path=output_path)\n    return submission_df\n\n\nif __name__ == \"__main__\":\n    TEST_IMAGE_DIR = \"/kaggle/input/csiro-biomass/test\"\n    # update folder_name if needed\n    folder_name = 'two_stream_v3_convnext_tiny_view1000'  # example name\n    import glob\n    MODEL_PATHS = glob.glob('/kaggle/input/two-stream-v3-convnext-tiny-view1000/pytorch/two-stream-v3-convnext-tiny-view1000-all/1/*.pth')\n\n    OUTPUT_PATH = \"/kaggle/working/submission.csv\"\n\n    submission_df = run_inference(\n        test_image_dir=TEST_IMAGE_DIR,\n        model_paths=MODEL_PATHS,\n        output_path=OUTPUT_PATH,\n        batch_size=16,\n        image_size=1000,   # image height; width will be image_size*2\n        num_workers=4,\n        device='cuda'\n    )\n\n    print(f\"✅ Inference completed! Submission saved to: {OUTPUT_PATH}\")\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-08T18:34:24.985986Z","iopub.execute_input":"2025-11-08T18:34:24.986641Z","iopub.status.idle":"2025-11-08T18:35:04.605658Z","shell.execute_reply.started":"2025-11-08T18:34:24.986614Z","shell.execute_reply":"2025-11-08T18:35:04.604640Z"}},"outputs":[],"execution_count":null}]}