血管分割比赛冠军解决方案（代码已更新）

首先，我们要感谢Kaggle和主办方举办了如此精彩的比赛。同时感谢@hengck23的精彩帖子，@junkoda的指标实现，以及所有其他参与者分享他们的实验。

概述

我们的最终提交是两个2.5D convnext tiny unet模型的集成，输入为3个通道。这两个模型唯一的区别在于数据增强和训练轮数。实际上，得分最高的提交不是所选的集成模型，而是集成中的一个单模型，在私有排行榜上的得分为0.835。

数据准备

我们使用了全部训练数据包括kidney_1_voi。

• 多视角切片（x, y, z）
• 归一化：无归一化，仅使用image = image / 65535.0
• 使用完整切片而非分块，所有切片调整大小或裁剪至1536x1536
• 数据增强：

  A.Compose([
      A.HorizontalFlip(p=0.5),
      A.VerticalFlip(p=0.5),
      A.Transpose(p=0.5),
      A.Affine(scale={"x":(0.7, 1.3), "y":(0.7, 1.3)}, translate_percent={"x":(0, 0.1), "y":(0, 0.1)}, rotate=(-30, 30), shear=(-20, 20), p=0.5),
      A.RandomBrightnessContrast(brightness_limit=0.4, contrast_limit=0.4, p=0.5),
      A.OneOf([
          A.Blur(blur_limit=3, p=0.2),
          A.MedianBlur(blur_limit=3, p=0.2),
      ], p=1.0),
      A.OneOf([
          A.ElasticTransform(alpha=1, sigma=50, alpha_affine=10, border_mode=1, p=0.5),
          A.GridDistortion(num_steps=5, distort_limit=0.1, border_mode=1, p=0.5)
      ], p=0.4),
      A.OneOf([
          A.Resize(1536, 1536, cv2.INTER_LINEAR, p=1),
          A.Compose([
              RandomResize(1536, 1536, scale_limit_x=0.5, scale_limit_y=0.5, p=1),
              A.PadIfNeeded(1536, 1536, position="random", border_mode=cv2.BORDER_REPLICATE, p=1.0),
              A.RandomCrop(1536, 1536, p=1.0)
          ], p=1.0),
      ], p=1.0),
      A.GaussNoise(var_limit=0.05, p=0.2),
  ])

• 随机3D旋转以获得不一定与坐标轴平行的切片。得分最高的提交使用了随机3D旋转增强并训练了更多轮次。
  

建模与训练

• 我们使用SMP中的unet，采用convnext tiny骨干网络，将BatchNorm和ReLU替换为GroupNorm和GELU，并添加了额外的卷积stem。所有模型的输入尺寸为3x1536x1536。

  self.extra_stem = nn.Sequential(
      nn.Conv2d(in_channels, out_channels, 3, 2, 1),
      LayerNorm2d(out_channels),
  )

• 损失函数：使用1.0 focal loss、1.0 dice loss、0.01边界损失（boundary loss）和1.0自定义损失。该自定义损失受@hengck23的帖子和@junkoda的指标实现启发。

  class CustomLoss(nn.Module):
      def __init__(self):
          super().__init__()
          power = 2**np.arange(0, 8).reshape(1, 1, 2, 2, 2).astype(np.float32)
          area = create_table_neighbour_code_to_surface_area((1, 1, 1)).astype(np.float32)
          self.power = nn.Parameter(torch.from_numpy(power), requires_grad=False)
          self.kernel = nn.Parameter(torch.ones(1, 1, 2, 2, 2), requires_grad=False)
          self.area = nn.Parameter(torch.from_numpy(area), requires_grad=False)
          
      def forward(self, preds, targets):
          """
          preds: 形状为[bs, 1, d, h, w]的张量
          targets: 形状为[bs, 1, d, h, w]的张量
          """
          bsz = preds.shape[0]
  
          # 体素logits转换为立方体logits
          foreground_probs = F.conv3d(F.logsigmoid(preds), self.kernel).exp().flatten(1)
          background_probs = F.conv3d(F.logsigmoid(-preds), self.kernel).exp().flatten(1)
          surface_probs = 1 - foreground_probs - background_probs
  
          # 真实标签
          with torch.no_grad():
              cubes_byte = F.conv3d(targets, self.power).to(torch.int32)
              gt_area = self.area[cubes_byte.reshape(-1)].reshape(bsz, -1)
              gt_foreground = (cubes_byte == 255).to(torch.float32).reshape(bsz, -1)
              gt_background = (cubes_byte == 0).to(torch.float32).reshape(bsz, -1)
              gt_surface = (gt_area > 0).to(torch.float32).reshape(bsz, -1)
          
          # dice计算
          foreground_dice = 2 * (foreground_probs*gt_foreground).sum(-1) / (foreground_probs.sum(-1)+gt_foreground.sum(-1)).clamp(1e-6)
          background_dice = 2 * (background_probs*gt_background).sum(-1) / (background_probs.sum(-1)+gt_background.sum(-1)).clamp(1e-6)
          surface_dice = 2 * (surface_probs*gt_area).sum(-1) / ((surface_probs+gt_surface)*gt_area).sum(-1).clamp(1e-6)
          dice = (foreground_dice + background_dice + surface_dice) / 3
          return 1 - dice.mean()

• 优化：使用AdamW和CosineAnnealingLR（从1e-4到0，带warmup）。除使用3D切片旋转增强的模型训练30轮外，其余模型均训练20轮，批大小为8，梯度累积步数为4。

推理

• 在3个轴上进行推理，使用8倍TTA（测试时增强）
• 尝试了不同的推理 resize 方法。对于得分最高的提交，所有切片简单 resize 到3072x3072；对于所选提交，使用动态缩放因子使(h*scale)*(w*scale)=3200*3200
• 提交使用的阈值为0.4，基于CV和LB的最优阈值约为0.4~0.5
• torch.compile()提供了约2倍加速，使我们能够使用高分辨率和TTA进行推理

无效尝试

• 3D模型
• 外部数据和伪标签
• Transformers
• 为2.5D模型堆叠更多切片（>3）

结果