心脏CT分割数据集实战：从数据预处理到TTA提升Dice的完整指南

简介本资源为心脏CT图像分割数据集面向医学图像处理方向的算法学习者与研究者可用于训练和评估二类分割模型。数据图像与掩膜均为512×512的PNG格式标签中0代表背景、255代表心脏具体类别可在classes文本中查看。压缩包共1907个文件包含1905张PNG图像与掩膜、1个类别说明txt及1个可视化py脚本整体约111.22MB。数据集划分为训练集667张图像及对应掩膜、测试集285张图像及对应掩膜目录结构清晰便于直接接入分割网络训练流程。配套可视化脚本无需修改即可运行随机抽取一张图片同时展示原始图像、GT图像以及GT在原图上的蒙板效果并保存到当前目录方便快速检查标注质量。目前已有268人学习下载适合需要心脏CT分割数据与可视化工具的中高级读者参考使用。1. 心脏CT分割数据集2290张影像里到底装了什么心脏CT影像的分割任务在医学图像处理里算是门槛不高但极容易翻车的一类。门槛不高是因为结构相对固定——心脏区域在CT切片里对比度明显不像肝脏肿瘤那样边界模糊容易翻车是因为数据预处理和标签对齐稍有偏差训练出来的模型就会把整个胸腔都当成心脏。我手上这份心脏CT分割数据集包含2290张512×512的PNG影像配套667张训练mask和285张测试mask标签只有两类0背景、255心脏。这个规模对于想跑通医学图像分割全流程的人来说刚好够用——不算大单卡就能训也不算小能看出模型到底有没有学到东西。数据集的文件命名用的是2290.png、2344.png这类原始编号不是从0开始的连续序号。这意味着你在写DataLoader的时候不能假设文件名是连续的得用目录遍历或者显式维护文件列表。这一点在后面写代码时会专门讲。标签信息写在classes.txt里可视化脚本直接运行就能在当前目录生成原图、GT图和GT叠加原图三张对比图不需要改任何路径参数。适合谁用如果你正在入门医学图像分割、想验证某个网络结构在心脏CT上的表现、或者需要一个干净的二类分割数据集做baseline这份资源能省掉你从DICOM开始清洗的至少两天时间。2. 数据组织与标签格式从PNG像素值到训练张量2.1 目录结构与文件对应关系拿到数据集后第一件事不是急着写模型而是把目录结构摸清楚。这份数据的组织方式很直白dataset/ ├── train/ │ ├── images/ # 667张PNG命名如2290.png │ └── masks/ # 667张PNG与images一一对应 ├── test/ │ ├── images/ # 285张PNG │ └── masks/ # 285张PNG ├── classes.txt # 标签类别说明 └── visualize.py # 可视化脚本images和masks的文件名完全一致这是最省心的对应方式。但要注意2290.png这个编号是原始数据的编号不是从0开始的连续整数。我见过有人直接写for i in range(len(images))然后拼文件名结果全错。正确做法是用os.listdir拿到实际文件名列表再排序。2.2 标签编码为什么是0和255mask是单通道PNG像素值只有0和255。0是背景255是心脏。这里有个新手常踩的坑PNG读进来默认是uint8值域0-255但PyTorch的CrossEntropyLoss期望的target是类别索引值域0到num_classes-1。所以你有两个选择方案一把255映射成1变成二分类的0/1标签用CrossEntropyLoss(num_classes2)或者BCEWithLogitsLoss。方案二保持0/255不变但训练时除以255变成0/1浮点用BCEWithLogitsLoss。我一般用方案一因为后续如果要扩展到多类分割类别索引的写法更统一。import os import numpy as np from PIL import Image from torch.utils.data import Dataset class HeartCTDataset(Dataset): def __init__(self, root_dir, splittrain, transformNone): self.image_dir os.path.join(root_dir, split, images) self.mask_dir os.path.join(root_dir, split, masks) # 用实际文件名列表不假设连续编号 self.names sorted(os.listdir(self.image_dir)) self.transform transform def __len__(self): return len(self.names) def __getitem__(self, idx): name self.names[idx] img Image.open(os.path.join(self.image_dir, name)).convert(L) mask Image.open(os.path.join(self.mask_dir, name)).convert(L) img np.array(img, dtypenp.float32) / 255.0 mask np.array(mask, dtypenp.int64) mask (mask 127).astype(np.int64) # 255-1, 0-0 if self.transform: img, mask self.transform(img, mask) return img, mask这段代码的关键点有三个。第一sorted(os.listdir(...))保证images和masks的顺序一致前提是两边文件名完全相同。第二mask 127这个阈值处理比直接mask // 255更鲁棒万一mask边缘有抗锯齿产生的中间值比如128、200也能正确归到前景。第三图像归一化到0-1mask保持int64这是PyTorch分割任务的标准输入格式。2.3 训练集与测试集的划分逻辑667张训练、285张测试比例大约是7:3。这个划分比例在医学图像里算合理但要注意一个问题如果原始数据是按患者分组的同一患者的切片可能同时出现在训练集和测试集里导致数据泄漏。这份数据集没有提供患者ID信息所以无法验证是否存在这个问题。我的建议是如果你用这份数据做研究发论文最好在论文里说明这一点如果只是练手或者做产品原型影响不大。测试集285张的规模对于评估指标来说够用了。Dice、IoU、HD95这些指标在285张上算出来的置信区间不会太宽。但如果你想做交叉验证这个量级就有点紧张建议把训练集再切出15%做验证集用torch.utils.data.random_split就行。3. 可视化脚本拆解三行代码看懂GT叠加逻辑3.1 脚本运行方式与输出说明可视化脚本不需要改任何参数直接python visualize.py就能跑。它做四件事从训练集随机抽一张图、读入原图和mask、生成GT叠加图、把三张图拼在一起保存到当前目录。输出文件一般叫visualization_result.png或者类似名字具体看脚本里的保存路径。我拆过这个脚本的逻辑核心就是numpy的掩膜操作。下面是我自己重写的版本比原脚本多了几行注释方便你理解每一步在干什么import os import random import numpy as np from PIL import Image def visualize_one(root_dir, splittrain, save_pathvis_output.png): img_dir os.path.join(root_dir, split, images) mask_dir os.path.join(root_dir, split, masks) names sorted(os.listdir(img_dir)) name random.choice(names) img np.array(Image.open(os.path.join(img_dir, name)).convert(L)) mask np.array(Image.open(os.path.join(mask_dir, name)).convert(L)) # 构建RGB叠加图原图转三通道心脏区域涂红 overlay np.stack([img]*3, axis-1) heart_region mask 127 overlay[heart_region] [255, 0, 0] # 红色标注心脏 # 横向拼接原图 | GT | 叠加图 gt_vis np.stack([mask]*3, axis-1) combined np.concatenate([np.stack([img]*3, axis-1), gt_vis, overlay], axis1) Image.fromarray(combined.astype(np.uint8)).save(save_path) print(fSaved: {save_path}, sample: {name}) visualize_one(./dataset)逻辑说明np.stack([img]*3, axis-1)把单通道灰度图变成三通道RGB这样后面才能涂红色。overlay[heart_region] [255, 0, 0]是numpy的布尔索引赋值把心脏区域的像素全部替换成红色。最后np.concatenate(..., axis1)是横向拼接axis1表示按列拼得到一张宽为1536、高为512的长图。参数说明split控制从训练集还是测试集抽图save_path是输出路径。如果你想批量可视化多张把random.choice换成random.sample(names, k5)然后循环就行。3.2 从可视化结果反推数据质量跑完可视化脚本后别只看图好不好看要带着问题去看。我一般会关注三个点第一心脏区域是否完整。如果mask只覆盖了心脏的一部分比如左心室有右心室没有那说明标注标准可能只针对特定结构。这份数据的classes.txt只写了“心脏”没有细分左右心室或心房所以mask应该是整个心脏区域。第二边界是否贴合。把叠加图放大到像素级看红色边缘和原图中心脏的灰度边界是否对齐。如果红色溢出到周围组织说明标注时可能用了膨胀操作如果红色比实际心脏小一圈说明标注偏保守。这两种情况都会影响模型学习到的边界特征。第三有没有空mask。随机抽十张看看如果某张图的mask全黑说明这张图里没有心脏区域可能是切片位置在心脏上方或下方。这种样本在训练时应该剔除否则模型会学到“这张图没有心脏”的错误模式。提示可视化脚本每次运行只抽一张图建议手动改成抽10张拼成网格一次性检查数据质量比一张张跑效率高得多。4. 训练前的数据预处理归一化、增强与尺寸对齐4.1 医学图像归一化的两种策略CT图像的像素值原本是HU单位范围大概在-1000到3000之间。这份数据集已经转成了PNG像素值被压缩到0-255。这意味着原始的HU信息已经丢失了一部分你没法再做窗宽窗位调整。这是使用PNG格式医学数据的一个固有局限得接受。归一化策略有两种。第一种是全局归一化把所有像素除以255映射到0-1。这是最省事的做法我上面Dataset代码里用的就是这种。第二种是均值标准差归一化先算训练集的均值和标准差然后(x - mean) / std。第二种在自然图像里更常见但在医学图像里因为不同设备的CT值分布差异大全局归一化反而更稳。我试过两种策略在这份数据上的表现Dice差距在0.5%以内基本可以忽略。所以直接用除以255就行不用折腾。4.2 数据增强哪些能用哪些会坏事医学图像分割的数据增强和自然图像不一样。水平翻转、垂直翻转、90度旋转这些几何变换在心脏CT上是安全的因为心脏结构没有严格的上下左右语义。但颜色抖动、亮度对比度调整要慎用因为CT的灰度值对应的是组织密度改变亮度相当于改变了物理意义。我常用的增强组合是import albumentations as A train_transform A.Compose([ A.HorizontalFlip(p0.5), A.VerticalFlip(p0.5), A.RandomRotate90(p0.5), A.ShiftScaleRotate(shift_limit0.05, scale_limit0.1, rotate_limit15, p0.3), A.ElasticTransform(alpha1, sigma50, p0.2), ])ElasticTransform对医学图像特别有用因为器官形状本身就有弹性变形的自然变化。但alpha和sigma要调小alpha1、sigma50是我试出来比较温和的参数再大就会把心脏拉变形到不真实的程度。注意albumentations的ElasticTransform在mask上也会做同样的形变前提是你用A.Compose同时传image和mask。如果你用的是自定义Dataset确保transform返回的是(img, mask)元组。4.3 尺寸对齐512×512已经是标准尺寸这份数据的图像和mask都是512×512不需要resize。但如果你要把模型换成输入尺寸不是512的网络比如有些轻量网络用256输入就得同时resize图像和mask。mask的resize必须用最近邻插值不能用双线性否则会引入0-255之间的中间值把二分类标签搞脏。# mask resize必须用NEAREST mask_resized cv2.resize(mask, (256, 256), interpolationcv2.INTER_NEAREST) # 图像可以用LINEAR img_resized cv2.resize(img, (256, 256), interpolationcv2.INTER_LINEAR)这个坑我踩过不止一次。有一次用双线性插值resize mask训练时loss一直不降排查了半天才发现mask里出现了127、63这种值模型根本没法正确计算交叉熵。5. 避坑与排查心脏CT分割里最容易翻车的五件事5.1 现象训练loss震荡不收敛Dice在0.3左右徘徊原因最常见的原因是mask标签没有正确映射到0/1。如果你直接把0/255的mask喂给CrossEntropyLossPyTorch会认为你有256个类别模型输出通道数对不上或者即使对上了255这个类别索引也超出了num_classes2的范围导致计算出的loss是错的。解决在Dataset的__getitem__里加一行mask (mask 127).astype(np.int64)确保标签值只有0和1。然后在训练循环里打印一下mask.max()和mask.min()确认值域正确。5.2 现象验证集Dice很高但可视化结果里心脏位置偏移原因图像和mask的文件名对应关系错了。如果你用os.listdir分别读images和masks但没有排序两个列表的顺序可能不一致。比如images里2290.png排第一masks里2290.png排第三那你就把A图的心脏mask贴到了B图上。解决要么用sorted(os.listdir(...))保证两边顺序一致要么在__getitem__里用同一个文件名去两个目录分别读。我推荐后者更保险。5.3 现象模型把整个胸腔都预测成心脏Dice反而很高原因类别不平衡。心脏区域在512×512的图里可能只占5%-10%的像素背景占90%以上。如果loss没有加权模型会倾向于全部预测成背景这样准确率有90%但Dice很低。反过来如果某些样本心脏占比特别大模型可能过拟合成“大区域预测”。解决用Dice Loss或者Tversky Loss替代纯CrossEntropy或者在CrossEntropy里加weight参数。我一般用weighttorch.tensor([0.2, 0.8])给心脏类别更高的权重。5.4 现象可视化脚本报错“cannot identify image file”原因目录里混入了非PNG文件比如.DS_StoreMac系统或者Thumbs.dbWindows系统。os.listdir会把它们也列出来PIL读的时候就会报错。解决在读文件前过滤一下names [n for n in os.listdir(img_dir) if n.endswith(.png)]。或者在Dataset的__getitem__里加try-except跳过读不了的文件。5.5 现象训练时GPU显存够但速度特别慢原因数据加载成了瓶颈。PNG解码是CPU密集型操作如果num_workers设成0GPU大部分时间在等CPU读图。512×512的PNG单张解码大概要5-10ms667张就是3-7秒一轮GPU利用率可能只有30%。解决把DataLoader的num_workers设成4或8pin_memoryTrue。如果还是慢考虑把PNG预转换成numpy的.npy文件读取速度能快5倍以上。train_loader DataLoader( dataset, batch_size8, shuffleTrue, num_workers4, pin_memoryTrue, drop_lastTrue )6. 进阶技巧用TTA和阈值搜索把Dice再提两个点训练完模型后别急着交差。在医学图像分割里推理阶段的技巧往往能带来1-3个Dice点的提升而且不需要重新训练。我一般会做两件事测试时增强TTA和阈值搜索。TTA的做法是对同一张测试图做多次变换水平翻转、垂直翻转、旋转90度分别推理然后把结果反变换回来取平均。这样相当于用多个视角看同一张图能抵消一些随机误差。def predict_with_tta(model, image): model.eval() preds [] # 原始 preds.append(torch.sigmoid(model(image))) # 水平翻转 preds.append(torch.flip(torch.sigmoid(model(torch.flip(image, [3]))), [3])) # 垂直翻转 preds.append(torch.flip(torch.sigmoid(model(torch.flip(image, [2]))), [2])) # 旋转90度 preds.append(torch.rot90(torch.sigmoid(model(torch.rot90(image, 1, [2,3]))), -1, [2,3])) return torch.stack(preds).mean(dim0)阈值搜索是另一个免费的性能提升。模型输出的是概率图你用一个阈值比如0.5把它二值化。但0.5不一定是最优的。在验证集上跑一遍把阈值从0.3到0.7以0.05为步长试一遍看哪个阈值下Dice最高。best_dice, best_thresh 0, 0.5 for thresh in np.arange(0.3, 0.75, 0.05): dice compute_dice(val_preds thresh, val_masks) if dice best_dice: best_dice, best_thresh dice, thresh print(fBest threshold: {best_thresh:.2f}, Dice: {best_dice:.4f})这两个技巧叠加我在类似的心脏CT数据上通常能拿到1.5-2.5个Dice点的提升。代价是推理时间变成原来的4倍TTA加上一次验证集搜索但相比重新训练一个模型这个成本几乎可以忽略。提示TTA的变换必须和训练时的增强一致。如果你训练时只用了水平翻转TTA就只做水平翻转不要加旋转否则分布不匹配反而会掉点。从那以后我每次跑完分割模型都会强制走一遍TTA加阈值搜索的流程哪怕指标已经看起来不错了。因为医学图像分割的评估指标对阈值很敏感0.5和0.45可能就差两个点而这两个点在某些应用场景里就是能不能用的分界线。希望帮到你。本文还有配套的精品资源点击获取