{"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":117682,"databundleVersionId":15062069,"sourceType":"competition"}],"dockerImageVersionId":31192,"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\nfrom PIL import Image\nimport zipfile\nimport cv2\n\nBASE_PATH = \"/kaggle/input/vesuvius-challenge-surface-detection\"\nTRAIN_IMAGES_DIR = os.path.join(BASE_PATH, \"train_images\")\nTEST_IMAGES_DIR = os.path.join(BASE_PATH, \"test_images\")\nTRAIN_LABELS_DIR = os.path.join(BASE_PATH, \"train_labels\")\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nclass Config:\n    batch_size = 8\n    learning_rate = 1e-3\n    epochs = 4\n    patch_size = 256\n    num_slices = 8\n\nconfig = Config()\n\nclass SimpleDataset(Dataset):\n    def __init__(self, image_paths, mask_paths=None, is_train=True):\n        self.image_paths = image_paths\n        self.mask_paths = mask_paths\n        self.is_train = is_train\n    \n    def __len__(self):\n        return len(self.image_paths)\n    \n    def load_volume_slice(self, path):\n        with Image.open(path) as img:\n            if img.n_frames == 1:\n                volume = np.array(img)[np.newaxis, ...]\n            else:\n                slices = []\n                for i in range(img.n_frames):\n                    img.seek(i)\n                    slices.append(np.array(img))\n                volume = np.stack(slices, axis=0)\n        \n        depth = volume.shape[0]\n        center = depth // 2\n        start = max(0, center - config.num_slices // 2)\n        end = min(depth, start + config.num_slices)\n        \n        if end - start < config.num_slices:\n            if start == 0:\n                end = config.num_slices\n            else:\n                start = depth - config.num_slices\n        \n        return volume[start:end]\n    \n    def __getitem__(self, idx):\n        volume_path = self.image_paths[idx]\n        volume_slice = self.load_volume_slice(volume_path)\n        \n        if self.mask_paths is not None:\n            mask_path = self.mask_paths[idx]\n            mask_slice = self.load_volume_slice(mask_path)\n            mask_slice = (mask_slice == 1).astype(np.float32)\n        else:\n            mask_slice = np.zeros_like(volume_slice, dtype=np.float32)\n        \n        height, width = volume_slice.shape[1:]\n        \n        h_start = max(0, (height - config.patch_size) // 2)\n        w_start = max(0, (width - config.patch_size) // 2)\n        \n        volume_patch = volume_slice[:, \n                                   h_start:h_start+config.patch_size, \n                                   w_start:w_start+config.patch_size]\n        \n        mask_patch = mask_slice[:, \n                               h_start:h_start+config.patch_size, \n                               w_start:w_start+config.patch_size]\n        \n        max_proj = np.max(volume_patch, axis=0)\n        mean_proj = np.mean(volume_patch, axis=0)\n        std_proj = np.std(volume_patch, axis=0)\n        \n        combined = np.stack([max_proj, mean_proj, std_proj], axis=-1)\n        \n        min_val = combined.min()\n        max_val = combined.max()\n        if max_val > min_val:\n            combined = (combined - min_val) / (max_val - min_val)\n        \n        combined = (combined * 255).astype(np.uint8)\n        \n        image_tensor = torch.from_numpy(combined).permute(2, 0, 1).float() / 255.0\n        mask_tensor = torch.from_numpy(mask_patch[0]).float()\n        \n        return image_tensor, mask_tensor\n\nclass SimpleModel(nn.Module):\n    def __init__(self):\n        super().__init__()\n        \n        self.encoder = nn.Sequential(\n            nn.Conv2d(3, 32, 3, padding=1),\n            nn.BatchNorm2d(32),\n            nn.ReLU(),\n            nn.MaxPool2d(2),\n            \n            nn.Conv2d(32, 64, 3, padding=1),\n            nn.BatchNorm2d(64),\n            nn.ReLU(),\n            nn.MaxPool2d(2),\n            \n            nn.Conv2d(64, 128, 3, padding=1),\n            nn.BatchNorm2d(128),\n            nn.ReLU(),\n            nn.MaxPool2d(2),\n            \n            nn.Conv2d(128, 256, 3, padding=1),\n            nn.BatchNorm2d(256),\n            nn.ReLU(),\n        )\n        \n        self.decoder = nn.Sequential(\n            nn.ConvTranspose2d(256, 128, 2, stride=2),\n            nn.BatchNorm2d(128),\n            nn.ReLU(),\n            \n            nn.ConvTranspose2d(128, 64, 2, stride=2),\n            nn.BatchNorm2d(64),\n            nn.ReLU(),\n            \n            nn.ConvTranspose2d(64, 32, 2, stride=2),\n            nn.BatchNorm2d(32),\n            nn.ReLU(),\n            \n            nn.Conv2d(32, 1, 1)\n        )\n    \n    def forward(self, x):\n        x = self.encoder(x)\n        x = self.decoder(x)\n        return x\n\ndef get_train_data():\n    train_files = []\n    mask_files = []\n    \n    available = []\n    for f in os.listdir(TRAIN_IMAGES_DIR):\n        if f.endswith('.tif'):\n            img_path = os.path.join(TRAIN_IMAGES_DIR, f)\n            mask_path = os.path.join(TRAIN_LABELS_DIR, f)\n            \n            if os.path.exists(mask_path):\n                try:\n                    with Image.open(img_path) as img:\n                        if img.n_frames > 0:\n                            available.append((img_path, mask_path))\n                except:\n                    continue\n    \n    if not available:\n        return None, None\n    \n    train_files, mask_files = zip(*available)\n    train_files = list(train_files)[:2]\n    mask_files = list(mask_files)[:2]\n    \n    dataset = SimpleDataset(train_files, mask_files, is_train=True)\n    dataloader = DataLoader(dataset, batch_size=config.batch_size, shuffle=True)\n    \n    return dataset, dataloader\n\ndef train_one_epoch(model, dataloader, optimizer):\n    model.train()\n    total_loss = 0\n    \n    for images, masks in dataloader:\n        images = images.to(device)\n        masks = masks.to(device).unsqueeze(1)\n        \n        optimizer.zero_grad()\n        \n        outputs = model(images)\n        \n        bce_loss = F.binary_cross_entropy_with_logits(outputs, masks)\n        \n        pred = torch.sigmoid(outputs)\n        pred_flat = pred.view(-1)\n        target_flat = masks.view(-1)\n        intersection = (pred_flat * target_flat).sum()\n        union = pred_flat.sum() + target_flat.sum()\n        dice_loss_val = 1 - (2. * intersection) / (union + 1e-8)\n        \n        loss = bce_loss + dice_loss_val\n        total_loss += loss.item()\n        \n        loss.backward()\n        optimizer.step()\n    \n    return model, total_loss / len(dataloader)\n\ndef predict_volume(model, volume_path):\n    model.eval()\n    \n    dataset = SimpleDataset([volume_path], is_train=False)\n    image, _ = dataset[0]\n    image = image.unsqueeze(0).to(device)\n    \n    with torch.no_grad():\n        output = model(image)\n        pred = torch.sigmoid(output)\n        pred_np = pred.squeeze().cpu().numpy()\n    \n    with Image.open(volume_path) as img:\n        if img.n_frames == 1:\n            original_shape = (1, np.array(img).shape[0], np.array(img).shape[1])\n        else:\n            original_shape = (img.n_frames, np.array(img).shape[0], np.array(img).shape[1])\n    \n    full_mask = np.zeros(original_shape, dtype=np.uint8)\n    \n    depth, height, width = original_shape\n    \n    h_start = max(0, (height - config.patch_size) // 2)\n    w_start = max(0, (width - config.patch_size) // 2)\n    \n    center_z = depth // 2\n    start_z = max(0, center_z - config.num_slices // 2)\n    end_z = min(depth, start_z + config.num_slices)\n    \n    binary_pred = (pred_np > 0.5).astype(np.uint8)\n    \n    patch_height = min(config.patch_size, height - h_start)\n    patch_width = min(config.patch_size, width - w_start)\n    \n    for z in range(start_z, end_z):\n        if h_start + patch_height <= height and w_start + patch_width <= width:\n            full_mask[z, \n                     h_start:h_start+patch_height, \n                     w_start:w_start+patch_width] = binary_pred[:patch_height, :patch_width]\n    \n    return full_mask\n\ndef save_tiff(mask, path):\n    mask_8bit = (mask * 255).astype(np.uint8)\n    \n    if mask_8bit.shape[0] == 1:\n        Image.fromarray(mask_8bit[0]).save(path)\n    else:\n        with Image.fromarray(mask_8bit[0]) as img:\n            images = [Image.fromarray(mask_8bit[i]) for i in range(1, mask_8bit.shape[0])]\n            img.save(path, save_all=True, append_images=images, compression=None)\n\ndef create_submission(model, test_df):\n    temp_files = []\n    \n    for idx, row in test_df.iterrows():\n        test_path = os.path.join(TEST_IMAGES_DIR, f\"{row['id']}.tif\")\n        \n        if os.path.exists(test_path):\n            prediction = predict_volume(model, test_path)\n        else:\n            prediction = np.zeros((64, 512, 512), dtype=np.uint8)\n            prediction[32, 256-32:256+32, 256-32:256+32] = 1\n        \n        temp_path = f\"{row['id']}.tif\"\n        save_tiff(prediction, temp_path)\n        temp_files.append(temp_path)\n    \n    with zipfile.ZipFile('submission.zip', 'w') as z:\n        for temp_path in temp_files:\n            z.write(temp_path)\n            os.remove(temp_path)\n\ntrain_dataset, train_loader = get_train_data()\n\nif train_loader is None:\n    test_df = pd.read_csv(os.path.join(BASE_PATH, \"test.csv\"))\n    dummy_model = SimpleModel().to(device)\n    create_submission(dummy_model, test_df)\n\nmodel = SimpleModel().to(device)\noptimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)\n\nfor epoch in range(config.epochs):\n    model, epoch_loss = train_one_epoch(model, train_loader, optimizer)\n\ntest_df = pd.read_csv(os.path.join(BASE_PATH, \"test.csv\"))\ncreate_submission(model, test_df)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-26T18:37:55.221874Z","iopub.execute_input":"2026-01-26T18:37:55.222126Z","iopub.status.idle":"2026-01-26T18:44:57.217102Z","shell.execute_reply.started":"2026-01-26T18:37:55.222106Z","shell.execute_reply":"2026-01-26T18:44:57.216266Z"}},"outputs":[],"execution_count":null}]}