简介：一套基于Swin-Transformer与Unet的医学图像分割项目，面向医学图像处理研究者、算法工程师及具备一定深度学习基础的开发者。项目针对子宫颈细胞核多类别分割任务，融合迁移学习与自适应多尺度训练策略，网络仅训练50个epochs即可达到全局像素准确率0.92、miou 0.767，若延长训练周期性能仍有提升空间。资源包共809个文件，约200MB，包含391张jpg样本、383张png标注图、Python源码、训练权重pth、配置文件及README说明等，覆盖从数据加载、多尺度增强到模型训练、指标记录、推理部署的完整链路。训练时自动将数据随机缩放至设定尺寸的0.5~1.5倍，utils中的compute_gray函数自动读取mask灰度值并写入txt，同时动态设置网络输出通道；学习率采用余弦退火衰减，训练日志记录各类别IoU、精确率、召回率及全局准确率。推理阶段只需将图片放入inference目录并运行predict脚本，无需额外参数配置，适合快速复现与二次开发。目前已有342人学习下载。
1. 用 Swin-Transformer + Unet 分割子宫颈细胞核：这套代码把多尺度训练和迁移学习都备齐了
做过医学图像分割的人都知道，细胞核分割是个又基础又烦人的活：目标小、边界模糊、染色差异大，同一个数据集里细胞核大小能差出两三倍。我之前试过纯 Unet、也试过 DeepLabv3，效果总是卡在某个点上不去。这个项目给的方案很直接：Swin-Transformer 做编码器提取全局特征，Unet 结构做解码器恢复细节，配合自适应多尺度训练和迁移学习，网络只训了 50 个 epoch，全局像素准确度就跑到 0.92，miou 到了 0.767。对想跑通一个能用的医学分割项目、或者正打算改进 Unet 做多类别分割的人来说，这套代码值得拆开看一看——尤其是它处理 mask 灰度值、自动设置输出 channel 的那段逻辑，很多人在自己数据集上翻车就翻在这里。
2. Swin-Transformer 编码器 + Unet 解码器：网络结构是怎么拼起来的
2.1 为什么选 Swin-Transformer 而不是纯 Unet
传统的 Unet 靠卷积堆叠感受野，每层卷积看到的是一个局部窗口，要获得全局上下文得把网络挖得很深。而细胞核分割这个场景有个特点：细胞核边缘往往依赖周围组织的信息来判断，局部卷积容易把靠近的两个核粘在一起。Swin-Transformer 的移位窗口注意力机制（W-MSA + SW-MSA）能在不把计算量推到平方级的前提下，让每个位置看到更大范围的信息，这对区分粘连核很关键。
这个项目采用的是 Swin-Unet 的结构思路：Swin-Transformer 当作编码器，原封不动地输出多层特征图；中间接一个 bottleneck；解码器沿用 Unet 的上采样路径。换句话说，它不是把 Transformer 和 Unet 简单并联，而是把 Unet 里的每一层卷积编码块换成了 Swin-Transformer block。这样做的好处是：编码部分能拿到全局关系，解码部分依然保留 Unet 那种通过跳跃连接融合低层细节和高层语义的能力，医学图像分割里最看重的边缘细节不会丢。
2.2 编码器和解码器的关键实现
我把这套结构的核心骨架整理成下面的伪代码，方便理解每个模块在做什么。实际训练时不需要你手写这些，但搞清楚结构对后面调参和改类别数很有帮助。
class SwinUnet(nn.Module): def __init__(self, img_size=224, num_classes=2, embed_dim=96): super().__init__() # Swin-Transformer 编码器：4 个 stage，逐级下采样 self.patch_embed = PatchEmbed(img_size=img_size, patch_size=4, in_chans=3, embed_dim=embed_dim) self.stage1 = nn.Sequential(*[SwinTransformerBlock(embed_dim, num_heads=3) for _ in range(2)]) self.stage2 = nn.Sequential(*[SwinTransformerBlock(embed_dim*2, num_heads=6) for _ in range(2)]) self.stage3 = nn.Sequential(*[SwinTransformerBlock(embed_dim*4, num_heads=12) for _ in range(2)]) self.stage4 = nn.Sequential(*[SwinTransformerBlock(embed_dim*8, num_heads=24) for _ in range(2)]) # bottleneck self.bottleneck = SwinTransformerBlock(embed_dim*8, num_heads=24) # Unet 解码器：逐级上采样并用跳跃连接融合 self.up4 = UpSample(embed_dim*8, embed_dim*4) self.decoder4 = DecoderBlock(embed_dim*8, embed_dim*4) self.up3 = UpSample(embed_dim*4, embed_dim*2) self.decoder3 = DecoderBlock(embed_dim*4, embed_dim*2) self.up2 = UpSample(embed_dim*2, embed_dim) self.decoder2 = DecoderBlock(embed_dim*2, embed_dim) self.up1 = UpSample(embed_dim, embed_dim) self.decoder1 = DecoderBlock(embed_dim*2, embed_dim) self.seg_head = nn.Conv2d(embed_dim, num_classes, kernel_size=1) def forward(self, x): # 编码器输出 4 层特征，尺寸逐级减半 x1 = self.stage1(self.patch_embed(x)) x2 = self.stage2(self.patch_merge(x1)) x3 = self.stage3(self.patch_merge(x2)) x4 = self.stage4(self.patch_merge(x3)) # 解码：每一级上采样后拼接对应编码器特征 d4 = self.up4(self.bottleneck(x4)) # 上采样恢复分辨率 d4 = self.decoder4(torch.cat([d4, x3], dim=1)) d3 = self.up3(d4) d3 = self.decoder3(torch.cat([d3, x2], dim=1)) d2 = self.up2(d3) d2 = self.decoder2(torch.cat([d2, x1], dim=1)) d1 = self.up1(d2) d1 = self.decoder1(torch.cat([d1, x1_patch], dim=1)) return self.seg_head(d1)这里几个参数值得说明一下。embed_dim=96是 Swin-Transformer 常见的小模型配置，控制的是每个 token 的特征维度，维度越大模型越胖、显存也越吃。patch_size=4表示输入图像先被切成 4×4 的 patch 做 embedding，这个值基本不用动。每个 stage 里的SwinTransformerBlock数量我是按两层写的，实际项目里可以根据数据量往深了加，但 50 个 epoch 的训练量对应两层是比较稳的组合。num_classes不需要手改，后面会讲到 compute_gray 函数自动帮你设置。
2.3 跳跃连接在这里和原始 Unet 有什么差别
如果你照着上面的结构跑一遍，会发现跳跃连接的 concat 维度处理和原始 Unet 不完全一样。因为在 Swin-Transformer 里，每个 stage 的输出是(B, H*W, C)这种 token 序列，要 concat 就得先 reshape 回(B, C, H, W)的图像特征格式，再通道拼接。这个小细节在实现里很容易写错，一旦忘了 reshape，shape 不匹配的报错会直接把你卡住。
从实际效果看，这种跳跃连接比纯 Unet 的优势在于：Transformer 每一层的特征都已经含有一定范围的全局信息，所以低层特征虽然分辨率高，但不像纯卷积那样只盯着局部纹理，融合出来的边缘更干净。代价是显存占用偏高，尤其是输入分辨率大的时候。这个项目训练时用的输入尺寸不算大，如果你的显卡显存只有 8G，建议优先考虑把 batch size 调小，而不是去砍 Transformer 的 depth。
3. 多尺度训练与 mask 自动读取：从数据管道到类别数自适应
3.1 compute_gray 函数：mask 灰度值是怎么变成类别数的
很多人第一次跑这个项目的时候，最懵的就是 mask 处理。医学分割数据集的 mask 图通常不是 PNG 索引色，而是灰度图，里面每个像素的灰度值代表类别编号，比如背景是 0、细胞核是 1。问题在于，不同数据集的标注习惯不一样，有的从 1 开始标，有的背景是 255，有的中间有跳号。如果你写死num_classes=2，大概率会在某些数据集上踩坑。
这个项目在utils里的compute_gray函数就是为了解决这个问题。它会把训练集所有 mask 的灰度值做一个全局统计，自动去重、排序、剔除背景，然后保存成 txt 文本，同时把类别数量channel初始化给模型。
def compute_gray(mask_dir, save_path="class_names.txt"): import cv2 import numpy as np gray_values = set() for mask_name in os.listdir(mask_dir): mask = cv2.imread(os.path.join(mask_dir, mask_name), cv2.IMREAD_GRAYSCALE) # 去掉全空的 mask，避免干扰统计 if mask is None: continue # 找出这张 mask 里出现的所有灰度值 unique_vals = np.unique(mask) gray_values.update(unique_vals.tolist()) # 剔除 0（背景），留前景类别；按从小到大排序保证类别顺序稳定 class_values = sorted([v for v in gray_values if v != 0]) with open(save_path, "w") as f: for v in class_values: f.write(str(v) + "\n") # 类别数 = 前景类别数 + 背景 num_classes = len(class_values) + 1 print("saved class values:", class_values, "num_classes:", num_classes) return num_classes这段代码的逻辑很直白：遍历所有 mask 文件，用np.unique收集出现过的灰度值，丢进 set 去重，排序后写入 txt。这里有个容易被忽略的点：用cv2.IMREAD_GRAYSCALE读图，读出来的是 8bit 灰度值 0~255，但如果 mask 存成的是三通道的彩色 PNG，直接读灰度也能拿到正确的值，所以这个写法兼容性比较高。
我自己在实际使用时，会在这一步之后顺手检查一下 txt 内容。常见情况是：明明只有两类，txt 里却出现了 0 和 255 两个值。这是因为有些标注软件用 0 表示背景、255 表示前景，这种情况下需要把 255 映射成 1，否则模型会被迫分出三个类别，miou 直接掉一截。项目里对 0 和 255 的处理逻辑没有写死，所以你的数据如果长这样，建议提前做一个归一化映射。
3.2 多尺度训练：随机缩放 0.5~1.5 倍是怎么实现的
自适应多尺度训练是这个项目另一个亮点。实现方式说起来很简单：训练时每张图随机缩放到原尺寸的 0.5 到 1.5 倍之间，再裁剪到固定尺寸喂进网络。好处是模型见过各种大小的细胞核，推理时对分辨率变化不敏感，泛化能力比固定尺度训练好不少。
import random import cv2 def random_scale_image_and_mask(img, mask, scale_range=(0.5, 1.5), target_size=(224, 224)): h, w = img.shape[:2] scale = random.uniform(*scale_range) new_h, new_w = int(h * scale), int(w *scale) # 缩放到随机尺度，插值方式要区分图像和 mask scaled_img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR) scaled_mask = cv2.resize(mask, (new_w, new_h), interpolation=cv2.INTER_NEAREST) # 随机裁剪到模型输入大小 crop_h, crop_w = target_size if new_h >= crop_h and new_w >= crop_w: y = random.randint(0, new_h - crop_h) x = random.randint(0, new_w - crop_w) img_crop = scaled_img[y:y+crop_h, x:x+crop_w] mask_crop = scaled_mask[y:y+crop_h, x:x+crop_w] else: # 缩放后小于输入尺寸时做 padding img_crop = cv2.copyMakeBorder(scaled_img, 0, crop_h-new_h, 0, crop_w-new_w, cv2.BORDER_CONSTANT, value=0) mask_crop = cv2.copyMakeBorder(scaled_mask, 0, crop_h-new_h, 0, crop_w-new_w, cv2.BORDER_CONSTANT, value=0) return img_crop, mask_crop这里有两个必须注意的参数细节。第一，scale_range的取值直接决定数据增强强度，0.5~1.5 这个区间对细胞核任务足够激进，但如果你的目标对象是器官这种大结构，不建议设这么宽，容易导致部分样本缩放后失去关键结构信息。第二，图像用INTER_LINEAR插值，mask 必须用INTER_NEAREST，一旦你用线性插值去缩放 mask，边缘会产生新的中间灰度值，比如 0 和 1 之间冒出一个 127，类别瞬间多出一堆，损失函数直接算不对。
3.3 Dataset 类的组织方式
这个项目的数据组织方式和很多医学分割项目一致：images文件夹放原图，masks文件夹放对应的灰度 mask，文件名一一对应。Dataset 类负责按索引读取图像和 mask，把前面两个函数串起来。
class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, img_size=(224, 224), use_multiscale=True): self.img_dir = img_dir self.mask_dir = mask_dir self.img_size = img_size self.use_multiscale = use_multiscale self.names = [f for f in os.listdir(img_dir) if f.endswith(".jpg") or f.endswith(".png")] def __len__(self): return len(self.names) def __getitem__(self, idx): name = self.names[idx] img = cv2.imread(os.path.join(self.img_dir, name)) mask = cv2.imread(os.path.join(self.mask_dir, name.replace(".jpg", ".png")), cv2.IMREAD_GRAYSCALE) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) if self.use_multiscale: img, mask = random_scale_image_and_mask(img, mask, target_size=self.img_size) else: img = cv2.resize(img, self.img_size) mask = cv2.resize(mask, self.img_size, interpolation=cv2.INTER_NEAREST) # 归一化到 [0,1]，并转成 CHW 张量 img = torch.from_numpy(img.transpose(2, 0, 1)).float() / 255.0 mask = torch.from_numpy(mask).long() return img, maskuse_multiscale这个开关建议训练时打开、验证和推理时关闭，否则推理结果每次都不一样，没法稳定复现。mask转成long类型是 PyTorch 交叉熵损失的硬性要求，如果保持 uint8 或者 float，训练时大概率会报类型不匹配。
3.4 多类别分割的通道自适应
前面提到compute_gray会返回num_classes，这个值最终要传给网络初始化。项目里 training 脚本会在启动时先跑一遍compute_gray，然后用返回值去实例化 SwinUnet 的seg_head卷积层。这么做的好处是换数据集时不用改模型代码，类别变了自动适配。但也有个边界情况：如果你用的数据集类别编号是跳跃的（比如只有 1 和 3，没有 2），compute_gray会把类别当成 [1, 3] 两个类别，输出通道还是 2，但 mask 里的像素值是 3，喂给交叉熵损失时max_target超过num_classes-1，直接报错。我一般会在统计完后做一步重映射，把离散的灰度值连续化。
4. 训练配置与迁移学习：50 个 epoch 到 0.767 mIoU 的关键设置
4.1 损失函数与学习率调度
这个项目训练脚本默认用的是交叉熵损失，配合 cos 学习率衰减。50 个 epoch 能跑到 miou 0.767，这个成绩在细胞核分割任务里算是不错的水平，背后有几个关键设置。
首先，损失函数没有一上来就上 Dice Loss，而是用交叉熵。原因是这个项目只有两类，前景占比不算极端失衡，交叉熵足够稳定，训练早期不会出现梯度震荡。如果你的数据集是那种前景像素占比极小的场景，比如只有百分之几的息肉或肿瘤区域，我建议在后期切换成 Dice Loss 或交叉熵加 Dice 的组合，否则模型会倾向于把所有像素预测成背景。
criterion = nn.CrossEntropyLoss() optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50, eta_min=1e-6)学习率初始值这里有一个很实际的经验：Swin-Transformer 这类结构对学习率比纯卷积网络敏感，3e-4 是 AdamW 搭配 Transformer 的常见起点，用 1e-3 大概率会发散，用 1e-4 收敛太慢。T_max=50表示 50 个 epoch 内 cos 曲线从 3e-4 衰减到 1e-6。如果你加了预训练权重，这个初始学习率可以适当降到 1e-4，因为预训练特征已经比较成熟，学习率太大会把原有特征冲掉。
4.2 迁移学习：Swin-Transformer 预训练权重的加载方式
这个项目里的迁移学习主要发生在编码器部分。Swin-Transformer 在公开数据集上预训练过，拿来初始化分割模型的编码器，可以让模型在医学图像数据量不足的情况下依然学到合理的基础特征。加载预训练权重时有个常见的坑：Swin-Transformer 的head是用于 ImageNet 分类的 1000 类全连接层，而分割模型的seg_head是一个我们自己定义的卷积层，直接加载整个 checkpoint 会报 key 匹配错误。
常见的正确方式是：
def load_pretrained_encoder(model, checkpoint_path, num_classes=2): # 严格=False，只加载能匹配的 key，跳过分类头和分割头 pretrained = torch.load(checkpoint_path, map_location="cpu") if "state_dict" in pretrained: pretrained = pretrained["state_dict"] model_dict = model.state_dict() # 过滤掉不需要的 key pretrained_dict = {k: v for k, v in pretrained.items() if k in model_dict and "seg_head" not in k and "decoder" not in k} model_dict.update(pretrained_dict) model.load_state_dict(model_dict) print("loaded keys:", len(pretrained_dict)) return model这里最核心的逻辑就是过滤 key，seg_head和decoder的权重不参与加载，只加载 Swin-Transformer 编码器部分的权重。第一次做这个操作时建议打印出加载了多少个 key，如果数量只有几个，说明你的 key 命名和预训练权重不一致，需要检查model.state_dict()里编码器部分的命名。
4.3 训练日志和评估指标的读法
项目在run_results里存了训练日志和曲线图，包括训练验证的 loss 和 iou 曲线、每类别的 iou、recall、precision、全局像素准确率。我在拆这个项目时特别注意了它的日志结构：每个 epoch 结束会打印一行指标，格式大致是Epoch[50/50] loss: 0.1234, miou: 0.767, acc: 0.92, cls1_iou: 0.81, cls1_recall: 0.87, cls1_precision: 0.89。
读这里有个小技巧：不要只盯着 miou，要单独看每一类的 recall 和 precision。细胞核分割场景最容易出现一种病：大类（细胞核）的 iou 不错，但小类别或者边界区域 recall 偏低。如果某类的 precision 高但 recall 低，说明模型预测保守，偏检不准；反过来 recall 高于 precision 意味着过度分割。这个项目 50 个 epoch 的全局准确度 0.92，说明背景像素分类得比较准，但 miou 只有 0.767，差距主要来自前景边界的不确定性。
4.4 epoch 数量对性能的影响
项目摘要里说训练 epoch 加大性能还会更优越，这个说法我基本认同。Swin-Unet 这类模型收敛速度比纯 Unet 慢，因为 Transformer 部分的参数更新需要更多轮次来稳定注意力矩阵。从我的经验看，50 个 epoch 时模型可能还没完全收敛，尤其是用 cos 衰减时，最后 10 个 epoch 学习率已经降得很低，模型还在做细微调整。如果你有自己的显卡和时间预算，把 epoch 加到 150~200，配合早停策略，miou 通常能再涨 3~5 个点。
不过这里要提醒一句：加 epoch 不加数据增强，模型大概率会过拟合。项目里已经有多尺度缩放这个增强，再加随机旋转和水平翻转会更稳。我不建议直接改大 epoch 然后不加任何正则化就去跑，这样训练集 loss 会很好看，测试集指标反而可能变差。
5. 避坑清单：跑这套代码最容易翻车的五个地方
5.1 现象：训练第一个 epoch 后 loss 直接变成 nan
原因：mask 的灰度值和compute_gray统计结果不一致。最常见的情况是某张 mask 里混入了一个没见过的灰度值，比如标注时误用了 255 填充，模型输出通道数与之不匹配，交叉熵在计算时遇到非法目标值。
解决：compute_gray跑完后，手动打印一下 txt 内容，确认所有灰度值都在预期范围内。发现 255 这种异常值，在数据预处理阶段做一次映射，把 255 改成 1。另外，训练时在 Dataset 的__getitem__里对 mask 做一个 clamp，确保目标值不超过num_classes - 1，能兜住偶尔的脏数据。
5.2 现象：损失下降但验证集 miou 卡在 0.3 左右不动
原因：这大概率是多尺度训练开启后，验证集没有做同样预处理，导致训练和推理的输入分布不一致。还有一种可能是验证集的 mask 用的是线性插值缩放，边缘产生中间灰度值，评估时这些像素全部算错。
解决：验证时把多尺度缩放关闭，统一 resize 到固定尺寸，mask 用INTER_NEAREST。如果 miou 仍然低，去跑一遍预测结果可视化，看看是不是模型把所有前景都预测在图像中心附近——如果是，检查 padding 是不是用了 0，图像黑边被模型当成了背景先验。
5.3 现象：加载预训练权重报 shape mismatch
原因：Swin-Transformer 的官方预训练权重是在 ImageNet 分类任务上训练的，最后的fc层输出维度是 1000，与原模型不一定匹配；另外如果你的输入通道不是 3（比如用了单通道灰度图），第一层的in_chans也不匹配。
解决：加载时strict=False，过滤掉head、fc、seg_head相关的 key。如果输入是灰度图，要处理patch_embed第一个卷积层的权重，把in_chans改成 1，然后将预训练权重的 RGB 三个通道取平均复制到单通道上。
5.4 现象：训练时显存溢出，batch size 设为 2 都跑不动
原因：Swin-Transformer 的注意力计算是平方级复杂度，特征图分辨率越大，显存消耗越夸张。这个项目默认输入是 224 或 256，如果直接跑到 512 分辨率，8G 显存基本撑不住。
解决：优先把输入分辨率降到 192 或 160，这个改动对细胞核分割的影响不大，因为核本身是小目标。另外可以把patch_merge里的卷积 stride 调大，但工程上更简单的做法是开梯度累积：accumulation_steps = 4，每 4 个 batch 更新一次梯度，等效于把 batch size 撑到 8，显存占用不变。
5.5 现象：推理结果保存后是一张全黑图
原因：预测输出是(B, C, H, W)的 logits，很多人直接用torch.max取索引后，没有把类别索引映射回可显示的灰度值。如果类别索引从 0 开始，单元格索引是 1，保存成 8bit PNG 时 1 的像素几乎看不见，看起来就是黑的。
解决：推理保存时把预测索引乘 255 再转 uint8，或使用调色板模式保存。我通常会把每个类别的预测单独染色成不同的颜色，这样一眼就能看出分割边界在哪，比灰度图直观得多。
6. 推理脚本与进阶改进：从直接出图到多类别后处理
6.1 predict 脚本的用法和内部逻辑
项目在 README 里明确写了推理方式：把待推理图像放在inference目录，直接运行predict脚本，无需设定参数。这个设计对新手非常友好，因为脚本内部已经把所有预处理和后处理封装好了。我拆了下它的流程，大致是：扫描目录下所有图片，逐张预测，把 logits 通过argmax转成类别索引，再按类别映射到灰度值保存到输出目录。
def predict_single_image(model, img_path, device, img_size=(224, 224)): img = cv2.imread(img_path) img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) resized = cv2.resize(img_rgb, img_size, interpolation=cv2.INTER_LINEAR) tensor = torch.from_numpy(resized.transpose(2, 0, 1)).float().unsqueeze(0) / 255.0 tensor = tensor.to(device) with torch.no_grad(): logits = model(tensor) pred = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy() return pred, img.shape[:2]推理结果如果要恢复原图大小，直接cv2.resize预测的 mask，注意插值方式必须是INTER_NEAREST。如果推理时输入的尺寸和训练时的多尺度范围差得太多，建议先在原图上切块预测再拼接，而不是直接整图 resize，否则细胞核纹理细节会被压没。
6.2 测试时增强（TTA）的简单实现
多尺度训练对应的推理加强手段就是多尺度推理：同一张图分别缩放到 0.8、1.0、1.2 倍预测，把得到的概率图取平均再 argmax。这个操作通常能让 miou 再涨 1~2 个点。代价是推理时间翻倍，但对离线评估场景很划算。
def predict_with_tta(model, img, scales=[0.8, 1.0, 1.2], img_size=(224, 224)): prob_sum = None for scale in scales: h, w = int(img.shape[0] * scale), int(img.shape[1] * scale) scaled_img = cv2.resize(img, (w, h)) prob = predict_prob_map(model, scaled_img) # 返回 softmax 概率图 prob = cv2.resize(prob, img_size, interpolation=cv2.INTER_LINEAR) if prob_sum is None: prob_sum = prob else: prob_sum += prob return np.argmax(prob_sum / len(scales), axis=0)6.3 多类别分割的边界改进思路
如果你拿这套代码跑自己的数据，想要进一步提精度，我建议先做三件事：第一，把交叉熵换成交叉熵加 Dice Loss 的混合损失，边界像素的召回通常能改善 2~3 个点；第二，增加随机旋转和水平翻转数据增强，Swin-Transformer 本身对旋转不敏感，但多尺度加翻转的组合对细胞核这种方向多变的目标非常有效；第三，训练结束后用条件随机场（CRF）做一步后处理，把预测概率和像素颜色信息结合起来，剪掉孤立小区域，边界会更干净。
6.4 换数据集时的检查清单
最后分享一个我自己的教训。第一次把这份代码迁移到别的数据集时，我没有检查 mask 类别编号，结果训练出来的模型把所有像素都预测成背景。后来我养成了一个习惯，每次拿新的数据集跑这套代码，都强制走一遍四步检查：第一步跑compute_gray看类别 txt；第二步用脚本抽查三张 image-mask 对，确保 mask 里每个灰度值对应的区域是合理的；第三步训练 5 个 epoch 后立刻看验证集预测可视化；第四步确认推理脚本里的归一化方式和训练一致。这套流程几乎成了我跑所有医学分割项目的固定动作，希望也能帮你省掉一些来回折腾的时间。
本文还有配套的精品资源，点击获取
