{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":117682,"databundleVersionId":15062069,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":14266755,"sourceType":"datasetVersion","datasetId":8766236},{"sourceId":14351778,"sourceType":"datasetVersion","datasetId":9163946},{"sourceId":14407816,"sourceType":"datasetVersion","datasetId":9201745}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import sys\n\nDATASET_NAME = \"imagecodecs\"\n\n!{sys.executable} -m pip install --no-index --find-links /kaggle/input/{DATASET_NAME} imagecodecs -q","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T03:40:34.317604Z","iopub.execute_input":"2026-01-06T03:40:34.318324Z","iopub.status.idle":"2026-01-06T03:40:37.674815Z","shell.execute_reply.started":"2026-01-06T03:40:34.318295Z","shell.execute_reply":"2026-01-06T03:40:37.673790Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 导入必要的库\nimport os\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nimport matplotlib.pyplot as plt\nfrom sklearn.model_selection import train_test_split\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-06T03:40:37.676751Z","iopub.execute_input":"2026-01-06T03:40:37.677053Z","iopub.status.idle":"2026-01-06T03:40:37.682334Z","shell.execute_reply.started":"2026-01-06T03:40:37.676995Z","shell.execute_reply":"2026-01-06T03:40:37.681639Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 设置设备\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T03:40:37.683268Z","iopub.execute_input":"2026-01-06T03:40:37.683623Z","iopub.status.idle":"2026-01-06T03:40:37.700956Z","shell.execute_reply.started":"2026-01-06T03:40:37.683592Z","shell.execute_reply":"2026-01-06T03:40:37.700337Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 1. 数据加载与预处理\n# -----------------------------\nimport os\nimport numpy as np\nimport torch\nfrom torch.utils.data import Dataset\nfrom torchvision import transforms\nfrom skimage.transform import resize\n\n# -----------------------------\n# 新增：从 3D ink label 生成 surface z-map\n# -----------------------------\ndef compute_surface_map(ink_label_3d):\n    \"\"\"\n    将 3D ink label (Z, H, W) 转换为 2D surface map (H, W)\n    对每个 (x,y)，返回第一个 ink=1 的 z 索引（从上往下）\n    若无 ink，设为 Z（可后续 clip）\n    \"\"\"\n    Z, H, W = ink_label_3d.shape\n    # 初始化为 Z（表示无墨水）\n    surface_map = np.full((H, W), Z, dtype=np.float32)\n    \n    # 沿 z 轴找第一个 ink=1 的位置\n    # 注意：argmax 在全零时返回 0，所以要用 mask\n    has_ink = ink_label_3d.any(axis=0)  # (H, W)\n    first_ink_z = np.argmax(ink_label_3d, axis=0).astype(np.float32)  # (H, W)\n    \n    surface_map[has_ink] = first_ink_z[has_ink]\n    return surface_map\n\nclass VesuviusDataset(Dataset):\n    def __init__(self, image_dir, label_dir, fragment_ids, z_start=28, z_size=3, img_size=320, transform=None):\n        self.image_paths = [os.path.join(image_dir, f\"{fid}.npy\") for fid in fragment_ids]\n        self.label_paths = [os.path.join(label_dir, f\"{fid}.npy\") for fid in fragment_ids]\n        self.z_start = z_start\n        self.z_size = z_size\n        self.img_size = img_size\n        self.transform = transform if transform is not None else transforms.ToTensor()\n\n    def __len__(self):\n        return len(self.image_paths)\n\n    def __getitem__(self, idx):\n        # Load 3D volume: shape (Z, H, W)\n        volume = np.load(self.image_paths[idx])  # (Z, H, W)\n        ink_label_3d = np.load(self.label_paths[idx])  # (Z, H, W) —— 这是 ink label!\n\n        # === Step 1: 生成 surface z-map (H, W) ===\n        surface_map = compute_surface_map(ink_label_3d)  # (H, W), float32\n\n        # 可选：只关注 z_start 附近的区域，clip 到合理范围\n        # 例如，我们只关心 z ∈ [z_start - 5, z_start + 10]\n        # 但通常直接使用原始 z 值即可\n        # surface_map = np.clip(surface_map, 0, volume.shape[0] - 1)\n\n        # === Step 2: Slice input volume ===\n        z_end = self.z_start + self.z_size\n        if z_end > volume.shape[0]:\n            z_end = volume.shape[0]\n            z_start_adj = max(0, z_end - self.z_size)\n            image_slice = volume[z_start_adj:z_end, :, :]  # (C, H, W)\n        else:\n            image_slice = volume[self.z_start:z_end, :, :]\n\n        # Transpose to (H, W, C) for resize and transform\n        image_slice = np.transpose(image_slice, (1, 2, 0))  # (H, W, C)\n        surface_map_orig = surface_map.copy()  # 用于 resize\n\n        # === Step 3: Resize if needed ===\n        if image_slice.shape[0] != self.img_size or image_slice.shape[1] != self.img_size:\n            # Resize image with anti-aliasing\n            image_slice = resize(\n                image_slice,\n                (self.img_size, self.img_size),\n                anti_aliasing=True,\n                preserve_range=True\n            )\n            # Resize surface map with nearest neighbor (order=0)\n            surface_map = resize(\n                surface_map_orig,\n                (self.img_size, self.img_size),\n                order=0,  # nearest neighbor for indices\n                anti_aliasing=False,\n                preserve_range=True\n            )\n\n        # === Step 4: Convert to tensors ===\n        image_tensor = self.transform(image_slice.astype(np.float32))\n        # 注意：surface_map 是 float32，不需要 unsqueeze(0) 如果模型输出单通道\n        label_tensor = torch.from_numpy(surface_map.astype(np.float32))  # (H, W)\n\n        return image_tensor, label_tensor\n\n# 示例用法\ntrain_dataset = VesuviusDataset(\n    image_dir=\"/path/to/images\",\n    label_dir=\"/path/to/labels\",\n    fragment_ids=[\"fragment1\", \"fragment2\"], # 示例ID列表\n    z_start=28,\n    z_size=3,\n    img_size=320,\n    transform=transforms.Compose([\n        transforms.ToTensor(),\n        # 可以添加其他变换\n    ])\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T03:40:37.701821Z","iopub.execute_input":"2026-01-06T03:40:37.702097Z","iopub.status.idle":"2026-01-06T03:40:37.718068Z","shell.execute_reply.started":"2026-01-06T03:40:37.702071Z","shell.execute_reply":"2026-01-06T03:40:37.717400Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 2. 构建训练和验证数据集\n# -----------------------------\n\nimport os\nimport pandas as pd\nfrom sklearn.model_selection import train_test_split\nfrom torchvision import transforms\n\n# 使用你实际添加的数据集路径\nDATA_ROOT = \"/kaggle/input/vesuvius-npy\"\ntrain_image_dir = os.path.join(DATA_ROOT, \"train_images\")\ntrain_label_dir = os.path.join(DATA_ROOT, \"train_labels\")\n\n# 读取 train.csv 获取所有样本 ID\ntrain_csv = pd.read_csv(os.path.join(DATA_ROOT, \"train.csv\"))\nids = train_csv[\"id\"].tolist()  # 列表形式，如 ['1004283650', '1006462223', ...]\n\n# 检查是否所有 ID 都存在\nmissing_ids = []\nfor id_str in ids:\n    img_path = os.path.join(train_image_dir, f\"{id_str}.npy\")\n    lbl_path = os.path.join(train_label_dir, f\"{id_str}.npy\")\n    if not os.path.exists(img_path) or not os.path.exists(lbl_path):\n        missing_ids.append(id_str)\n        print(f\"缺失文件: {id_str}\")\n\nif missing_ids:\n    print(f\"找到 {len(missing_ids)} 个缺失的样本，已跳过\")\n    # 过滤掉缺失的 ID\n    valid_ids = [id_str for id_str in ids if id_str not in missing_ids]\nelse:\n    valid_ids = ids\n\nprint(f\"成功加载 {len(valid_ids)} 个有效训练样本\")\n\n# 划分训练集和验证集\ntrain_ids, val_ids = train_test_split(\n    valid_ids, \n    test_size=0.2, \n    random_state=42\n)\n\n# 定义图像变换\ntransform = transforms.Compose([\n    transforms.ToTensor(),\n])\n\n# 创建训练和验证数据集实例\ntrain_dataset = VesuviusDataset(\n    image_dir=train_image_dir,\n    label_dir=train_label_dir,\n    fragment_ids=train_ids,\n    z_start=28,\n    z_size=3,\n    img_size=320,\n    transform=transform\n)\n\nval_dataset = VesuviusDataset(\n    image_dir=train_image_dir,\n    label_dir=train_label_dir,\n    fragment_ids=val_ids,\n    z_start=28,\n    z_size=3,\n    img_size=320,\n    transform=transform\n)\n\n# 输出数据集大小确认\nprint(f\"训练集大小: {len(train_dataset)}\")\nprint(f\"验证集大小: {len(val_dataset)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T03:40:37.719727Z","iopub.execute_input":"2026-01-06T03:40:37.719977Z","iopub.status.idle":"2026-01-06T03:40:38.378836Z","shell.execute_reply.started":"2026-01-06T03:40:37.719956Z","shell.execute_reply":"2026-01-06T03:40:38.378078Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 3. 定义 U-Net 模型\n# -----------------------------\n\nclass DoubleConv(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(DoubleConv, self).__init__()\n        self.conv = nn.Sequential(\n            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass Down(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(Down, self).__init__()\n        self.maxpool_conv = nn.Sequential(\n            nn.MaxPool2d(2),\n            DoubleConv(in_channels, out_channels)\n        )\n\n    def forward(self, x):\n        return self.maxpool_conv(x)\n\n\nclass Up(nn.Module):\n    def __init__(self, in_channels, out_channels, bilinear=True):\n        super(Up, self).__init__()\n        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) if bilinear else nn.ConvTranspose2d(in_channels // 2, in_channels // 2, kernel_size=2, stride=2)\n        self.conv = DoubleConv(in_channels, out_channels)\n\n    def forward(self, x1, x2):\n        x1 = self.up(x1)\n        diffY = x2.size()[2] - x1.size()[2]\n        diffX = x2.size()[3] - x1.size()[3]\n        x1 = nn.functional.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2])\n        x = torch.cat([x2, x1], dim=1)\n        return self.conv(x)\n\n\nclass OutConv(nn.Module):\n    def __init__(self, in_channels, out_channels):\n        super(OutConv, self).__init__()\n        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)\n\n    def forward(self, x):\n        return self.conv(x)\n\n\nclass UNet(nn.Module):\n    def __init__(self, n_channels=3, n_classes=1):\n        super(UNet, self).__init__()\n        self.n_channels = n_channels\n        self.n_classes = n_classes\n\n        self.inc = DoubleConv(n_channels, 64)\n        self.down1 = Down(64, 128)\n        self.down2 = Down(128, 256)\n        self.down3 = Down(256, 512)\n        self.down4 = Down(512, 512)\n        self.up1 = Up(1024, 256)\n        self.up2 = Up(512, 128)\n        self.up3 = Up(256, 64)\n        self.up4 = Up(128, 64)\n        self.outc = OutConv(64, n_classes)\n\n    def forward(self, x):\n        x1 = self.inc(x)\n        x2 = self.down1(x1)\n        x3 = self.down2(x2)\n        x4 = self.down3(x3)\n        x5 = self.down4(x4)\n        x = self.up1(x5, x4)\n        x = self.up2(x, x3)\n        x = self.up3(x, x2)\n        x = self.up4(x, x1)\n        logits = self.outc(x)\n        # 移除 sigmoid！\n        return logits  # 直接输出浮点 z 值\n\n# 初始化模型（不变）\nmodel = UNet(n_channels=3, n_classes=1).to(device)\n\n# 修改优化器和损失函数\noptimizer = optim.Adam(model.parameters(), lr=1e-4)\ncriterion = nn.L1Loss()  # 或 nn.SmoothL1Loss()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T03:40:38.379841Z","iopub.execute_input":"2026-01-06T03:40:38.380123Z","iopub.status.idle":"2026-01-06T03:40:38.494107Z","shell.execute_reply.started":"2026-01-06T03:40:38.380100Z","shell.execute_reply":"2026-01-06T03:40:38.493312Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 3.5 创建 DataLoader\n# -----------------------------\n\nfrom torch.utils.data import DataLoader\n\nbatch_size = 8  # 可根据显存调整（如 4, 16）\nnum_workers = 0  # 调试阶段设为 0，避免多进程问题\n\ntrain_loader = DataLoader(\n    train_dataset,\n    batch_size=batch_size,\n    shuffle=True,\n    num_workers=num_workers,\n    pin_memory=True if device.type == \"cuda\" else False\n)\n\nval_loader = DataLoader(\n    val_dataset,\n    batch_size=batch_size,\n    shuffle=False,\n    num_workers=num_workers,\n    pin_memory=True if device.type == \"cuda\" else False\n)\n\nprint(f\"Train loader: {len(train_loader)} batches\")\nprint(f\"Val loader: {len(val_loader)} batches\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T03:40:38.495102Z","iopub.execute_input":"2026-01-06T03:40:38.495411Z","iopub.status.idle":"2026-01-06T03:40:38.501992Z","shell.execute_reply.started":"2026-01-06T03:40:38.495371Z","shell.execute_reply":"2026-01-06T03:40:38.501464Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 调试：检查一个 batch 的 shape\nfor data, target in train_loader:\n    print(\"Data shape:\", data.shape)     # [B, 3, H, W]\n    print(\"Target shape:\", target.shape) # [B, 1, H, W]\n    break","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T03:40:38.502901Z","iopub.execute_input":"2026-01-06T03:40:38.503497Z","iopub.status.idle":"2026-01-06T03:40:41.684276Z","shell.execute_reply.started":"2026-01-06T03:40:38.503468Z","shell.execute_reply":"2026-01-06T03:40:41.683578Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 4. 训练循环\n# -----------------------------\n\n#def train_epoch(model, dataloader, optimizer, criterion, device):\n#    model.train()\n#    total_loss = 0\n#    for batch_idx, (data, target) in enumerate(dataloader):\n#        data, target = data.to(device), target.to(device)  # target: (B, H, W)\n#        optimizer.zero_grad()\n#        output = model(data)  # (B, 1, H, W)\n#        output = output.squeeze(1)  # (B, H, W) to match target\n#        loss = criterion(output, target)\n#        loss.backward()\n#        optimizer.step()\n#        total_loss += loss.item()\n#    return total_loss / len(dataloader)\n\n#def validate_epoch(model, dataloader, criterion, device):\n#    model.eval()\n#    total_loss = 0\n#    with torch.no_grad():\n#        for data, target in dataloader:\n#            data, target = data.to(device), target.to(device)\n#            output = model(data).squeeze(1)\n#            loss = criterion(output, target)\n#            total_loss += loss.item()\n#    return total_loss / len(dataloader)\n\n# 训练（取消注释）\n#num_epochs = 10\n#for epoch in range(num_epochs):\n#    train_loss = train_epoch(model, train_loader, optimizer, criterion, device)\n#    val_loss = validate_epoch(model, val_loader, criterion, device)\n#    print(f\"Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}\")\n\n# 直接加载预训练权重\nmodel = UNet(n_channels=3, n_classes=1).to(device) \nmodel.load_state_dict(torch.load(\"/kaggle/input/model4/unet_surface (3).pth\", map_location=device))\nmodel.eval()\nprint(\"Pretrained model loaded.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T03:40:41.685188Z","iopub.execute_input":"2026-01-06T03:40:41.685420Z","iopub.status.idle":"2026-01-06T04:26:42.152724Z","shell.execute_reply.started":"2026-01-06T03:40:41.685399Z","shell.execute_reply":"2026-01-06T04:26:42.149768Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 4.5 保存模型（仅在开发时运行）\n# -----------------------------\n\n#torch.save(model.state_dict(), \"/kaggle/working/unet_surface.pth\")\n#print(\"Model saved to /kaggle/working/unet_surface.pth\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T04:26:42.162416Z","iopub.execute_input":"2026-01-06T04:26:42.162854Z","iopub.status.idle":"2026-01-06T04:26:42.335424Z","shell.execute_reply.started":"2026-01-06T04:26:42.162830Z","shell.execute_reply":"2026-01-06T04:26:42.334605Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 5. 测试数据集与推理\n# -----------------------------\n\nclass VesuviusTestDataset(Dataset):\n    def __init__(self, image_dir, fragment_ids, z_start=28, z_size=3, img_size=320, stride=160):\n        \"\"\"\n        用于滑动窗口推理（大图切片）\n        Args:\n            image_dir: 测试图像目录\n            fragment_ids: 测试片段ID列表（如 ['a', 'b']）\n            z_start, z_size: 同训练\n            img_size: 模型输入尺寸\n            stride: 切片步长（< img_size 以重叠）\n        \"\"\"\n        self.image_paths = [os.path.join(image_dir, f\"{fid}.npy\") for fid in fragment_ids]\n        self.fragment_ids = fragment_ids\n        self.z_start = z_start\n        self.z_size = z_size\n        self.img_size = img_size\n        self.stride = stride\n\n    def __len__(self):\n        return len(self.image_paths)\n\n    def __getitem__(self, idx):\n        volume = np.load(self.image_paths[idx])  # (Z, H, W)\n        z_end = self.z_start + self.z_size\n        if z_end > volume.shape[0]:\n            z_end = volume.shape[0]\n            z_start_adj = max(0, z_end - self.z_size)\n            image_slice = volume[z_start_adj:z_end, :, :]  # (C, H, W)\n        else:\n            image_slice = volume[self.z_start:z_end, :, :]\n        \n        # 转为 (H, W, C)\n        image_slice = np.transpose(image_slice, (1, 2, 0))\n        original_shape = image_slice.shape[:2]  # (H, W)\n        \n        # 切片成 patches\n        patches = []\n        positions = []\n        h, w = original_shape\n        for y in range(0, h, self.stride):\n            for x in range(0, w, self.stride):\n                y_end = min(y + self.img_size, h)\n                x_end = min(x + self.img_size, w)\n                patch = np.zeros((self.img_size, self.img_size, image_slice.shape[2]), dtype=np.float32)\n                patch[0:y_end-y, 0:x_end-x] = image_slice[y:y_end, x:x_end]\n                patches.append(patch)\n                positions.append((y, x, y_end, x_end))\n        \n        patches = np.stack(patches)  # (N, H, W, C)\n        patches = torch.from_numpy(patches.transpose(0, 3, 1, 2))  # (N, C, H, W)\n        return patches, positions, original_shape, self.fragment_ids[idx]\ndef predict_fragment(model, patches, positions, original_shape, device, img_size=320):\n    model.eval()\n    pred_mask = np.zeros(original_shape, dtype=np.float32)\n    count_mask = np.zeros(original_shape, dtype=np.float32)\n\n    with torch.no_grad():\n        for i in range(len(patches)):\n            patch = patches[i].unsqueeze(0).to(device)  # (1, C, H, W)\n            y, x, y_end, x_end = positions[i]\n            \n            output = model(patch)  # (1, 1, H, W)\n            prob = output.squeeze().cpu().numpy()  # (H, W)\n            \n            # 只取有效区域（非 padding 部分）\n            h_valid = y_end - y\n            w_valid = x_end - x\n            pred_mask[y:y_end, x:x_end] += prob[:h_valid, :w_valid]\n            count_mask[y:y_end, x:x_end] += 1\n\n    # 平均重叠区域\n    pred_mask = np.divide(pred_mask, count_mask, where=count_mask != 0)\n    return pred_mask","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T04:26:42.336476Z","iopub.execute_input":"2026-01-06T04:26:42.336768Z","iopub.status.idle":"2026-01-06T04:26:42.351026Z","shell.execute_reply.started":"2026-01-06T04:26:42.336737Z","shell.execute_reply":"2026-01-06T04:26:42.350215Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 6. 正确推理：输出 z-index map，非 binary mask\n# -----------------------------\n\nimport tifffile\nfrom pathlib import Path\nimport numpy as np\nimport os\nimport glob\n\nTEST_IMAGE_DIR = \"/kaggle/input/vesuvius-challenge-surface-detection/test_images\"\nSUBMISSION_DIR = Path(\"/kaggle/working/submission_tifs\")\nSUBMISSION_DIR.mkdir(exist_ok=True)\n\ntest_files = glob.glob(os.path.join(TEST_IMAGE_DIR, \"*.tif\"))\nprint(f\"Found {len(test_files)} test files to process.\")\n\nfor test_file in test_files:\n    fragment_id = os.path.splitext(os.path.basename(test_file))[0]\n    print(f\"Processing fragment: {fragment_id}\")\n    \n    volume = tifffile.imread(test_file)  # (Z, H, W) —— 注意：Kaggle .tif 是 (H, W, Z)!\n    # ⚠️ 重要：Kaggle 提供的 .tif 是 (H, W, Z)，需转置！\n    if volume.ndim == 3 and volume.shape[2] > volume.shape[0]:\n        volume = volume.transpose(2, 0, 1)  # (H, W, Z) -> (Z, H, W)\n    Z, H, W = volume.shape\n    print(f\" Volume shape: ({Z}, {H}, {W})\")\n\n    img_size = 320\n    stride = 160\n    pred_z = np.zeros((H, W), dtype=np.float32)\n    count_mask = np.zeros((H, W), dtype=np.float32)\n\n    for y in range(0, H, stride):\n        for x in range(0, W, stride):\n            y_end = min(y + img_size, H)\n            x_end = min(x + img_size, W)\n            \n            z_start = 28\n            z_size = 3\n            z_end = min(z_start + z_size, Z)\n            z_start_adj = max(0, z_end - z_size)\n            \n            patch_3d = volume[z_start_adj:z_end, y:y_end, x:x_end]  # (C, h, w)\n            patch_full = np.zeros((z_size, img_size, img_size), dtype=np.float32)\n            patch_full[:, :y_end-y, :x_end-x] = patch_3d\n            \n            patch_tensor = torch.from_numpy(patch_full.astype(np.float32)).unsqueeze(0).to(device)\n            \n            with torch.no_grad():\n                output = model(patch_tensor)  # (1, 1, H, W)\n                z_pred = output.squeeze().cpu().numpy()  # (H, W)\n            \n            pred_z[y:y_end, x:x_end] += z_pred[:y_end-y, :x_end-x]\n            count_mask[y:y_end, x:x_end] += 1\n\n    # Average overlapping regions\n    pred_z = np.divide(pred_z, count_mask, where=count_mask != 0)\n    \n    # Round to nearest integer and convert to uint8 (assuming z < 256)\n    pred_z = np.round(pred_z).clip(0, 255).astype(np.uint8)\n\n    # 保存为 z-index map，不是 binary mask！\n    output_path = SUBMISSION_DIR / f\"{fragment_id}.tif\"\n    tifffile.imwrite(output_path, pred_z, dtype=np.uint8)\n    print(f\" Saved prediction: {output_path}\")\n\nprint(\"All test files processed.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T04:26:42.352127Z","iopub.execute_input":"2026-01-06T04:26:42.352483Z","iopub.status.idle":"2026-01-06T04:26:42.849330Z","shell.execute_reply.started":"2026-01-06T04:26:42.352449Z","shell.execute_reply":"2026-01-06T04:26:42.848562Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 7. 打包所有 .tif 文件为 submission.zip\n# -----------------------------\n\nimport shutil\n\n# 创建 ZIP 包含所有 .tif 文件\nshutil.make_archive(\"/kaggle/working/submission\", 'zip', SUBMISSION_DIR)\n\nprint(\"Successfully created submission.zip with .tif masks!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-06T04:26:42.850864Z","iopub.execute_input":"2026-01-06T04:26:42.851450Z","iopub.status.idle":"2026-01-06T04:26:42.857887Z","shell.execute_reply.started":"2026-01-06T04:26:42.851422Z","shell.execute_reply":"2026-01-06T04:26:42.857202Z"}},"outputs":[],"execution_count":null}]}