ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

1200张手术器械语义分割数据集实战:从加载到训练避坑指南

1200张手术器械语义分割数据集实战:从加载到训练避坑指南 简介面向医学图像分割研究者与算法工程师这份数据集聚焦手术器械图像的语义分割任务共包含约1200张已处理图片及对应标签覆盖32个类别适合用于训练、验证和评估各类分割网络。数据已按标准划分训练集约860张验证集约370张均提供images原图与masks模板目录可直接接入UNet、SwinUNet、TransUNet等模型流程。压缩包共2000个文件以png图像及标签、jpg原图为主附带1个txt说明文件和1个可视化脚本py。脚本可随机抽取一张样本将原始图像、GT掩膜及GT在原图上的叠加结果保存为对比图便于快速检查标注质量与分割效果。资源包整体大小约80.79MB目录和文件名组织清晰适合入门及进阶实验使用。目前已有104人学习下载。除数据与脚本外资源描述中还提供相关医学图像分割网络及改进方向的内容索引方便结合数据集进一步探索模型实现细节。1. 1200张手术器械图像语义分割数据集值不值得拿来训练模型手里捏着约1200张已处理好的手术器械语义分割数据第一反应是先别急着喂给模型。这个体量在内镜手术视频的序列帧场景里算中等偏小但关键优势是已处理——标签和图像对齐过、类别编码统一、没有多余的裁剪和噪声省掉了整个项目中收益最低、消耗最大的标注清洗阶段。对做医疗影像、机器人辅助手术或器械追踪的团队来说拿它跑通语义分割算法、验证流程、做预训练时间成本从以周计降到以天计。这篇文章按我实际做这类数据集的习惯来写先讲数据本身该怎么读、怎么验货再给出一套能直接跑的PyTorch加载与增强方案然后落到模型选型和训练参数最后把最容易翻车的五个坑摊开讲末尾附一个能稳定提分的技巧组合。适合刚上手语义分割的研究生和要快速搭建分割pipeline的工程团队。2. 拆开数据集看门道目录结构、标签编码与类别分布2.1 拿到手先别训练花十分钟验货手术器械语义分割数据的交付形态五花八门但已处理四个字不代表可以直接进DataLoader。我一般先做三件事看目录树、查标签通道数、统计每个类别的像素占比。目录结构常见做法如下dataset/ ├── images/ │ ├── 0001.png │ ├── 0002.png │ └── ... ├── masks/ │ ├── 0001.png │ ├── 0002.png │ └── ... └── train_val_split.txt注意 masks 这层很多数据集的标签图是 RGB 索引图不是单通道灰度图。RGB 索引图的意思是每个类别对应一种固定 RGB 颜色视觉上很直观但直接读成三通道放进模型会让类别数莫名膨胀到几十甚至上百个这是新手最容易中招的地方。我拿到手会先写一段脚本统计标签图里的唯一像素值确认到底是单通道还是三通道from PIL import Image import numpy as np from collections import Counter mask_path dataset/masks/0001.png mask np.array(Image.open(mask_path)) print(shape:, mask.shape, dtype:, mask.dtype) if len(mask.shape) 3: # 三通道索引图把 RGB 压缩成单一整数标签 h, w, c mask.shape flattened mask.reshape(-1, c).sum(axis1) unique_vals np.unique(flattened) print(RGB索引图, 有效标签数(去重后):, len(unique_vals)) else: unique_vals np.unique(mask) print(单通道图, 像素值集合:, unique_vals[:20]) # 统计类别像素占比 pixel_counter Counter(mask.reshape(-1).tolist()) total_pixels h * w for cls_id, cnt in pixel_counter.most_common(): ratio cnt / total_pixels * 100 if cls_id 0: print(f类别 {cls_id}: {ratio:.2f}%)这段脚本的逻辑分三步走先确认 mask 的通道数和 dtype再判断它是 RGB 索引图还是单通道标注图最后统计背景和每个前景类别的像素占比。像素占比这一步至关重要——手术器械在画面里往往只占很小面积如果某个类别的占比长期低于 2%训练时损失函数会被背景完全淹没后面必须引入加权或Focal Loss来兜底。提示RGB 索引图压缩成单通道时别用 RGB 三个通道做 one-hot直接把三个通道值按位运算映射成整数或者把整张图替换成类别 ID 映射表注意别把不同类别映射到同一个 ID。2.2 类别平衡性决定你要不要换损失函数这类数据的类别划分一般沿着背景 器械类别展开常见拆法是抓钳、持针器、剪刀、吸引管、单极电凝钩等再加上背景一共五到八个类。每类器械的使用频率、在画面中的尺寸、遮挡程度都不同像素占比常常相差一个数量级。验证完标签结构后我会画一张类别像素占比的柱状图直接决定当前方案用 CrossEntropy 还是 Dice Loss 还是加权组合。判断标准很简单如果占比最小的类小于 5%直接硬训标准 CrossEntropy大概率出现背景 mIoU 很高、器械类基本全丢的假阳性结果。这时候一般做法是把训练目标改成带类别权重或直接换成复合损失。另外还要检查一个细节同一把器械是否横跨了多张连续帧。内镜视频按帧抽出来的数据前后帧里器械的形态和位置高度相似如果切分训练验证集时按文件名随机切验证集可能全是训练集里出现过的熟人mIoU 虚高到 90% 以上换到真实场景立刻崩掉。正确做法是按视频来源或按手术分组切分保证同一把器械只出现在一个集合里。3. 从文件夹到 DataLoader数据加载与增强的完整写法3.1 自定义 Dataset处理路径对齐、类别映射和缓存确认数据格式后下一步写 Dataset 类。这一步最核心的问题是图像和 mask 的文件名对齐。有的数据包是 images/ 和 masks/ 同名有的在文件名后面带 _mask 后缀有的是 CSV 里记录配对关系。别假设直接在 Dataset 初始化时把配对检查做掉import os import cv2 import torch from torch.utils.data import Dataset class InstrumentSegDataset(Dataset): def __init__(self, image_dir, mask_dir, class_names, transformNone): self.image_dir image_dir self.mask_dir mask_dir self.transform transform self.class_names class_names # 1. 建立文件名配对 self.samples [] for fname in sorted(os.listdir(image_dir)): if not fname.lower().endswith((.png, .jpg, .jpeg)): continue mask_path os.path.join(mask_dir, self._get_mask_name(fname)) if not os.path.exists(mask_path): print(f[警告] {fname} 对应的 mask 不存在跳过) continue self.samples.append(( os.path.join(image_dir, fname), mask_path )) # 2. 校验一下配对数量 print(f成功配对 {len(self.samples)} 张图像) if len(self.samples) 0: raise RuntimeError(没有找到任何配对数据检查目录结构) def _get_mask_name(self, image_name): # 兼容两种常见命名同名 和 _mask 后缀 stem os.path.splitext(image_name)[0] for candidate in [ stem .png, stem _mask.png, stem .npy ]: if os.path.exists(os.path.join(self.mask_dir, candidate)): return candidate return not_found.png def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, mask_path self.samples[idx] image cv2.imread(img_path) image cv2.cvtColor(image, cv2.COLOR_BGR2RGB) mask cv2.imread(mask_path, cv2.IMREAD_UNCHANGED) # 如果 mask 是三通道索引图压缩成单通道 if len(mask.shape) 3: mask self._rgb_to_class_id(mask) # mask 像素值可能不连续重映射为 0..num_classes-1 mask self._remap_labels(mask) if self.transform is not None: transformed self.transform(imageimage, maskmask) image transformed[image] mask transformed[mask] image torch.from_numpy(image).permute(2, 0, 1).float() / 255.0 mask torch.from_numpy(mask).long() return image, mask_get_mask_name这一段解决的是命名不统一的磕绊避免在数据加载阶段被卡住。真正的核心在_remap_labels手术器械数据集的类别 ID 常因为标注软件的历史原因不连续比如标签文件里只有 0、1、4 三个值中间缺了 2 和 3。直接拿去训练CrossEntropy 会把 4 当成第五类类别数对不上就崩了。统一重映射成 0 开头的连续 ID 是必要的兜底操作。图像归一化放到了最后除以 255 转 float避免在数据集内部反复算。注意mask 读入时用cv2.IMREAD_UNCHANGED防止 16 位深度标签被强制转成 8 位丢掉高值类。很多采集软件保存的标签 PNG 是 16 位深度默认读法很容易把 300 多的类别值压回 44这类问题查起来非常隐蔽。3.2 增强策略为什么普通翻转不够需要弹性形变和亮度抖动手术器械分割里图像增强的定位是用二十张图模拟出两百张图的变化。内镜手术视频的光照变化剧烈器械表面有反光组织颜色深浅不一单靠水平翻转和随机裁剪完全不够。我一般用 albumentations 库配置这套组合import albumentations as A from albumentations.pytorch import ToTensorV2 def get_train_transforms(crop_size(512, 512)): return A.Compose([ A.RandomResizedCrop( heightcrop_size[0], widthcrop_size[1], scale(0.5, 1.0), ratio(0.75, 1.33), p0.8 ), A.HorizontalFlip(p0.5), A.Rotate(limit15, border_mode0, p0.6), A.ElasticTransform( alpha1, sigma8, alpha_affine8, border_mode0, p0.4 ), A.OneOf([ A.ColorJitter(brightness0.2, contrast0.2, p0.5), A.RandomBrightnessContrast(brightness_limit0.15, p0.5), ], p0.7), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ], additional_targets{mask: mask}) def get_val_transforms(crop_size(512, 512)): return A.Compose([ A.Resize(heightcrop_size[0], widthcrop_size[1]), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ], additional_targets{mask: mask})albumentations 的好处是Compose里同时传入 image 和 mask它会在内部保证几何变换翻转、旋转、裁剪、弹性形变对两者施加完全相同的参数不会出现图像转了 mask 没转的错位。这个同步问题如果用 torchvision 手动实现要写额外的 transform 绑定逻辑一不小心 mask 就错位了。参数上的两个选择要解释一下Rotate的border_mode0表示旋转出界区域补黑边而不是用边缘像素外扩填充。在手术器械场景里补黑边等于引入一个完全不属于真实环境的黑色未知类但代价是可以避免外扩填充把器械边缘模糊成一大片伪影。实测下来补黑边配合 RandomResizedCrop 对最终 IoU 的影响更小。ElasticTransform的 sigma 控制在 8 左右再大器械形态会被拉得不像器械模型学到的是扭曲形状而不是真实形状泛化反而变差。4. 模型选型与训练参数U-Net、DeepLabV3 和损失函数怎么定4.1 三个候选模型在小数据集上的表现差异1200 张图的规模模型的选择逻辑和十万张级别完全不一样。大模型不一定赢关键是收敛速度和过拟合风险。我在这类数据集上常用三个候选按推荐顺序排模型参数量适合原因需要警惕的点U-Net (ResNet34 encoder)约24M经典分割骨架1200张图足够收敛代码生态最成熟感受野有限对细长器械的断裂处敏感DeepLabV3 (MobileNetV3)约15MASPP模块对多尺度器械效果好推理速度快小目标边缘细节弱于U-NetSegFormer (MiT-B0)约3.7MTransformer全局建模遮挡场景鲁棒数据量偏少容易欠拟合需要更强增强个人经验是第一次跑通用 U-Net ResNet34 encoder。理由很简单这个组合在医学图像分割里的可复现性最高PyTorch 生态里用 segmentation-models-pytorch 几行就能拉起来踩坑的人最少后期换 DeepLabV3 只需要改一行模型名。4.2 Loss 组合与关键超参Dice 打底、CrossEntropy 保梯度损失函数的选择直接决定这次训练是能跑还是能打。标准 CrossEntropy 在像素占比悬殊时的表现上面说过会偏向背景。纯 Dice Loss 又有个问题当某个器械类别在 Batch 里完全不存在时Dice 的梯度为零这一轮更新对那个类没有任何学习信号。所以在这类数据上常见的做法是组合两个 Loss加权的比例我一般这样设import torch import torch.nn.functional as F from torch import nn class CombinedLoss(nn.Module): def __init__(self, num_classes, ce_weight0.4, dice_weight0.6, class_weightsNone, smooth1.0): super().__init__() self.num_classes num_classes self.ce_weight ce_weight self.dice_weight dice_weight self.smooth smooth # class_weights 由类别像素占比的反比计算得到 self.class_weights torch.tensor(class_weights, dtypetorch.float32) \ if class_weights is not None else None def forward(self, logits, targets): # logits: (B, C, H, W), targets: (B, H, W) 像素值为类别ID B, C, H, W logits.shape # CrossEntropy 部分 ce_loss F.cross_entropy( logits, targets, weightself.class_weights.to(logits.device) ) # Dice 部分对每个类别单独计算再平均 probs F.softmax(logits, dim1) # (B, C, H, W) targets_onehot F.one_hot(targets, num_classesC) # (B, H, W, C) targets_onehot targets_onehot.permute(0, 3, 1, 2).float() # (B, C, H, W) dice_total 0.0 for c in range(1, C): # 背景类不参与Dice否则背景占比太大会稀释信号 num (probs[:, c] * targets_onehot[:, c]).sum(dim(1, 2)) * 2 den probs[:, c].sum(dim(1, 2)) targets_onehot[:, c].sum(dim(1, 2)) dice (num self.smooth) / (den self.smooth) dice_total dice.mean() dice_loss 1.0 - dice_total / (C - 1) return self.ce_weight * ce_loss self.dice_weight * dice_loss这段代码的三个关键设计点说明一下。class_weights取各类像素占比的倒数并归一化作用是让占比 20% 的剪刀不至于在梯度里压制占比 1% 的吸引管。Dice 只从类别 1 开始计算跳过背景因为背景面积常常占 80% 以上把背景放进 Dice 会让器械类的梯度被严重稀释。最后加权比例上 CE 取 0.4、Dice 取 0.6是实测平衡性比较好的配置如果发现训练后期某些小类别仍然不敏感可以把 Dice 权重提到 0.7。训练超参上固定这几个值能少走弯路输入分辨率 512×512初始学习率 1e-4batch size 8配 AdamW 优化器权重衰减 1e-4。学习率用余弦退火从 1e-4 衰减到 1e-6总 epoch 数 40 到 60 之间。1200 张图配 512 分辨率单张 V100 上大概两小时能跑完 60 个 epoch发现 mIoU 在验证集上连续 8 个 epoch 不涨就提前停止。5. 训手术器械分割最容易翻车的 5 个坑与排查办法5.1 标签图被读成三通道时类别数莫名其妙翻倍现象训练脚本打印的num_classes和传入模型的分类头数量对不上报维度错误或者模型没报错但 mIoU 一直低得离谱。原因cv2.imread(mask_path)没加IMREAD_UNCHANGED时三通道 RGB 索引图被完整读成(H, W, 3)而后续逻辑把它当成(H, W)用。更隐蔽的是三通道里每个通道各自有 5 个像素值按通道取唯一值再组合时类别数从 6 膨胀到 60。解决在 Dataset 里加三通道检测len(mask.shape) 3时先查一遍这张图的 RGB 映射表把每个像素的 RGB 组合映射成一个类别 ID映射表从 0 开始重新编号。同时把归一化和类别数校验放到数据加载阶段提前暴露问题而不是等到 loss 崩了再倒查。5.2 类别不平衡导致网络只学背景器械全丢现象验证集 Dice 在背景类上高达 0.99但每个器械类别的 IoU 都低于 0.1可视化预测图时器械区域完全黑掉。原因像素占比差距超过 20 倍时标准 CrossEntropy 的梯度几乎被背景像素垄断。网络发现把所有像素预测为背景就能拿到 80% 以上的准确率于是停在了一个看起来很漂亮、实际没有用的局部最优。解决换用上面的组合 Loss并确认class_weights来自训练集全部图像的类别像素分布统计而不是抽样几眼。另一个并行手段是把训练时的混淆矩阵单独打印出来每次验证时看各类别的 IoU 明细别只看平均 mIoU。5.3 数据增强后 mask 和图像错位现象训练 loss 反复震荡降不下去打开可视化发现器械边缘像被平移过器械头部出现在组织上。原因用了torchvision.transforms手动对 image 和 mask 分别做随机翻转和旋转随机种子没有对齐两组操作产生了不同的参数。elastic deformation 类变换如果自己写更容易错位。解决统一换用 albumentations它内部保证几何变换同步。如果坚持用 torchvision就用torchvision.transforms.v2里的RandomResizedCrop等成对接口或者靠同一个torch.Generator传种子。我自己踩过一次自己写的同步函数排了半天发现是random.random()被调用了两次导致翻转概率不一致从此不再手动实现随机变换。5.4 细长器械被下采样直接切断IoU 归零现象预测图里持针器的弯曲部分断裂成两三截本应连续的器械轮廓中间出现空洞。原因器械在图像中占的像素少但形状细长。随机裁剪到 512×512 时器械整体可能被切出视野或者 Resize 到小尺寸时细长区域被降采样抹平边缘信息丢失。解决数据加载阶段保持原始分辨率训练输入尺度里设置随机裁剪范围比如scale(0.5, 1.0)让部分样本是原图高分辨率裁剪部分样本是全图缩放变相模拟多尺度。推理阶段不要直接 Resize 到固定尺寸用全卷积方式在原图上滑窗预测再把结果拼回原始分辨率。如果实在要 Resize最小边不要低于 480。5.5 train/val 同源泄漏验证 mIoU 虚高 15 个点以上现象训练时验证 mIoU 冲到 0.9 以上看起来模型已经收敛得很好但直接部署到新录的手术视频上表现掉到 0.65跌得人发懵。原因内镜视频抽帧的时序相关性很强相邻两帧里器械的形态和位置几乎相同。如果按文件名简单 8:2 随机划分验证集里可能混进了训练序列的相邻帧。模型等于见过答案验证曲线从头到尾都是自嗨。解决这类数据必须按视频源分组切分。如果数据包没提供视频来源标注就用文件名前缀分组同一前缀通常是同一段采集视频全部划到同一个集合。如果连前缀都没有至少检查相邻文件名是否落在同一集合里。分组切分之后mIoU 通常会掉 5 到 10 个点别慌这才是真实水平。6. 想再提点分类别加权、TTA 与 CRF 后处理的组合模型训出来的基线 mIoU 在 0.8 左右时再往上走靠的是推理阶段的细节打磨。第一个推荐做的是测试时增强对每个验证图像做水平翻转推理把两次的 softmax 概率取平均再 argmax。这个操作几乎零成本对器械分割稳定提升 1 到 2 个点 mIoU。TTA 的代码实现比想象中短def predict_with_tta(model, image_tensor, device): model.eval() with torch.no_grad(): pred torch.softmax(model(image_tensor.unsqueeze(0).to(device)), dim1) flipped torch.flip(image_tensor, dims[2]) pred_flip torch.softmax(model(flipped.unsqueeze(0).to(device)), dim1) pred_flip torch.flip(pred_flip, dims[2]) pred_avg (pred pred_flip) / 2.0 return pred_avg.squeeze(0)第二个推荐是 CRF 后处理。手术器械的预测图常有零星的孤立噪点器械边缘容易被错误分类成背景或反光组织。传统方法里全连接 CRF 能把空间上相邻、颜色相似的像素拉成同一个类别消除椒盐噪声的同时修齐边缘。对推理速度不敏感的场景CRF 值得加但对实时器械追踪后处理的开销可能超过模型推理时间本身需要根据场景取舍。注意TTA 和 CRF 本质是用计算换精度先确认自己的部署环境对单帧延迟的要求。离线标注辅助场景随便加实时分割场景建议只保留 TTA。第三个投入产出比较高的动作是硬负样本挖掘。把验证集上 IoU 最低的几十张图单独挑出来检查是剧烈运动模糊、器械出画、还是强反光造成的。这类样本在 1200 张数据里可能只有三五十张手动把它们复制三份并配以更强增强放进训练集比单纯加大 epoch 数更直接。我习惯在每次训练结束后输出 validation 集逐图的 IoU 排序从中挑出最差的 5% 做针对性补充这个习惯帮我处理过好几次平均指标不错、但某几类器械稳定丢的问题省下不少重新标注的功夫。另外说一下验证指标的选择。手术器械分割场景里别只盯 mIoU把每类的 IoU 单独打印重点关注小类别。平均 mIoU 会被优势类别拉高隐藏掉小类别的退化。我在训练脚本里每次验证都会输出一张混乱矩阵观察误分类集中在哪两个类之间——最常见的混乱发生在持针器和抓钳之间两者形态相近且光照下有镜面反光。发现类似情况后我会给分析目标加上这两个类的通道注意力或类别权重比无脑调学习率有效。轮到自己做这个数据集时最想留给你的一句实话是1200 张图足够你跑通和验证一整套分割流程但别指望它能撑起一个直接上线的医疗级模型。把这批数据当成流程验证和模型选型的跳板跑顺了之后再去扩数据或做增量标注才是这个体量的正确用法。希望这些踩坑记录能帮你少走几趟弯路。本文还有配套的精品资源点击获取
返回列表