如何在 PyTorch 中确保图像与掩码同步应用数据增强变换

在医学图像分割（如 u-net 训练）中，图像与对应掩码必须经受**完全相同的几何变换**（如旋转、翻转），否则标签错位将导致模型学习失效；本文提供基于 `torchvision.transforms.v2` 的可靠同步增强方案，并详解关键实现细节。

在使用 PyTorch 进行语义分割任务（如乳腺肿块分割）时，一个常见却极易被忽视的陷阱是：对图像和掩码分别调用同一 transform 实例，并不保证二者经历完全一致的随机变换。原因在于：torchvision.transforms.v2（以及旧版 v1）中的随机变换（如 RandomRotation、RandomHorizontalFlip）每次调用都会重新采样随机参数（如旋转角度、是否翻转）。因此，self.transform(image) 和 self.transform(mass_mask) 实际上执行了两次独立的随机操作，导致图像与掩码空间错位——这正是你观察到的图中掩码“漂移”或形变不匹配的根本原因。

✅ 正确解法：将图像与掩码沿通道维度拼接后统一变换，再拆分
该方法强制两者共享同一组随机参数，从源头保障几何一致性。以下是修正后的 __getitem__ 关键代码（适配 torchvision.transforms.v2）：

def __getitem__(self, index):
    dict_path = os.path.join(self.dict_dir, self.data[index])
    patient_dict = torch.load(dict_path)
    image = patient_dict['image'].unsqueeze(0)  # shape: [1, H, W]
    mass_mask = patient_dict['mass_mask'].unsqueeze(0)  # shape: [1, H, W]
    mass_mask[mass_mask > 1.0] = 1.0

    if self.transform is not None:
        # ✅ 关键：拼接 → 变换 → 拆分
        # 注意：需确保 image 和 mask 均为 float 类型（v2 要求）
        image = image.float()
        mass_mask = mass_mask.float()

        # 沿 channel 维度（dim=0）拼接：[2, H, W]
        combined = torch.cat([image, mass_mask], dim=0)

        # 单次调用 transform，确保同步
        transformed = self.transform(combined)

        # 拆分回原结构
        image = transformed[0:1, ...]      # 第 0 个通道 → 图像
        mass_mask = transformed[1:2, ...]  # 第 1 个通道 → 掩码

    return image, mass_mask

⚠️ 重要注意事项：

• 数据类型：torchvision.transforms.v2 默认要求输入为 float 张量（非 uint8），务必在拼接前调用 .float()，否则可能报错或行为异常；

• 填充策略：RandomRotation(fill=...) 中的 fill 值需同时适用于图像与掩码。对于二值掩码，推荐 fill=0.0（背景值）；若图像使用 fill=255.0，需确认掩码背景也为 0.0，避免旋转后引入非法灰度值；

  Face++旷视

  Face⁺⁺ AI开放平台

• 插值方式：图像通常用双线性插值（默认），而掩码应使用最近邻插值以保持像素类别完整性。v2 中可通过 InterpolationMode.NEAREST 显式指定（需自定义组合）：

  train_transform = T.Compose([
      T.RandomRotation(degrees=35, expand=True, fill=(0.0, 0.0)),  # (img_fill, mask_fill)
      T.RandomHorizontalFlip(p=0.5),
      T.RandomVerticalFlip(p=0.5),
      # ⚠️ 注意：v2 的 RandomRotation 不直接支持 per-channel interpolation，
      # 如需严格控制，建议改用 albumentations（见下文备选方案）
  ])

• 备选方案：Albumentations（更推荐用于分割）
  若需更精细控制（如为图像/掩码指定不同插值模式），albumentations 是行业首选：

  import albumentations as A
  from albumentations.pytorch import ToTensorV2
  
  train_transform = A.Compose([
      A.Rotate(limit=35, p=1.0, border_mode=cv2.BORDER_CONSTANT, value=0),  # mask-safe fill
      A.HorizontalFlip(p=0.5),
      A.VerticalFlip(p=0.5),
      ToTensorV2()  # 自动转为 torch.Tensor
  ])
  
  # 在 __getitem__ 中：
  # transformed = train_transform(image=image.numpy(), mask=mass_mask.numpy())
  # image = torch.from_numpy(transformed["image"])
  # mass_mask = torch.from_numpy(transformed["mask"])

? 总结：同步增强的核心原则是「一次采样，共同变换」。避免分别调用随机变换，坚持使用通道拼接法或专业分割库（如 Albumentations）。这不仅是技术细节，更是分割任务精度的基石——错位的掩码会向模型注入不可修复的噪声，直接拖垮 Dice 系数与临床可用性。