Mask R-CNN多任务脊柱CT分割：椎体定位+病灶标注联合建模

简介本资源为Spark“数字人体”AI挑战赛——脊柱疾病智能诊断大赛的完整参赛源码包面向计算机、数学、电子信息等专业的本科生及研究生适用于算法竞赛备赛、医学AI项目实践与深度学习模型复现学习。压缩包共26个文件含16个核心Python脚本涵盖模型构建、训练、推理及可视化模块、7个编译缓存文件pyc、2个说明文档README.md与SECURITY.md及1个依赖清单requirements.txt整体体积仅138KB轻量紧凑且结构清晰便于快速部署与代码剖析。已有83人下载学习适合作为竞赛级医学影像分析项目的参考范例。读者可直接运行复现参赛方案深入理解Mask R-CNN在脊柱CT图像分割与病灶定位中的工程实现细节掌握数据预处理、模型配置、并行训练及结果可视化等关键环节并借鉴其模块化设计如utils、mrcnn、config等目录组织提升自身AI项目开发规范性。1. 这不是普通医学图像分割项目它用 Mask R-CNN 在脊柱 CT 上做多类椎体定位病灶标注联合建模你打开这个 ZIP 包第一眼看到mrcnn/目录和model.py里大量build_rpn_graph、proposal_layer的调用可能会下意识认为——“又一个基于 Matterport Mask R-CNN 的复现”。但实际跑通samples/tianchi/train.py后你会发现它没用标准 COCO 类别结构而是把 L1–L5、T1–T12、C1–C7 共 26 个椎体解剖位置当作独立类别更关键的是utils.py中SpineDataset类重写了load_mask()返回的 mask 不是单通道二值图而是 shape(H,W,27) 的张量——第 0 通道是背景第 1–26 通道分别对应每个椎体的像素级掩膜第 27 通道才是病灶如椎间盘突出、椎体压缩骨折、Schmorl结节的联合标注。这意味着模型必须同步完成解剖结构识别 病理区域定位两个任务且二者存在强空间约束病灶只能落在某一个椎体区域内。这种设计跳出了常规“先检测后分割”或“单任务分割”的框架直指临床真实需求——放射科医生看片时从来不是孤立判断“有没有突出”而是说“L4/5 椎间盘向后突出压迫硬膜囊”。资源适配的是天池平台提供的脊柱 CT 重建数据集DICOM→NIfTI→PNG 序列所有预处理脚本已封装在utils.py的preprocess_dicom_series()函数中支持自动窗宽窗位归一化与层厚插值。适合计算机视觉方向的本科生组队打医疗AI竞赛也适合想深入理解多任务医学图像建模边界的研究生拆解学习。2. 为什么选 Mask R-CNN 而非 U-Net 或 TransUNet从脊柱解剖先验到损失函数定制2.1 解剖结构建模决定网络架构选型脊柱疾病诊断的核心难点在于同一张 CT 图像中需同时识别多个相似解剖结构椎体形态高度同质化并精确定位微小病灶如 Schmorl 结节直径常3mm。U-Net 类分割模型虽在像素级精度上表现优异但缺乏显式的位置感知能力——它无法天然区分“T12 椎体下缘”和“L1 椎体上缘”而这两个区域在灰度分布上几乎无差别。TransUNet 引入 Transformer 编码器提升了长程依赖建模能力但其注意力机制对局部细微纹理变化不敏感易将椎体边缘的骨皮质伪影误判为病灶。相比之下Mask R-CNN 的 Region Proposal NetworkRPN强制模型先生成候选框再对每个框进行分类与掩膜预测天然契合“先定位椎体、再细化病灶”的临床阅片逻辑。源码中config.py的SpineConfig类明确设置了NUM_CLASSES 2726椎体1病灶且DETECTION_MIN_CONFIDENCE 0.7高于常规目标检测阈值0.3–0.5这是为了过滤掉因椎体形态微小差异导致的低置信度误检。2.2 多任务损失函数的耦合设计标准 Mask R-CNN 的损失由分类损失loss_class、边界框回归损失loss_bbox、掩膜损失loss_mask三部分加权组成。本项目在此基础上新增了loss_anatomical_constraint其实现位于model.py的compute_loss()函数中# model.py 第 1287 行起 def compute_loss(self, target_class_ids, target_bbox, target_mask, pred_class_logits, pred_deltas, pred_masks): # ... 原始 loss 计算 ... # 新增解剖约束损失强制病灶掩膜只存在于对应椎体掩膜内 anatomical_loss 0.0 for i in range(target_mask.shape[0]): # batch 内每张图 # 获取该图中所有椎体掩膜通道 1-26 vertebra_masks target_mask[i, ..., 1:27] # shape(H,W,26) lesion_mask target_mask[i, ..., 26] # shape(H,W) # 计算病灶区域是否超出椎体覆盖范围 vertebra_union np.max(vertebra_masks, axis-1) # 合并所有椎体区域 # 病灶区域外的像素若被预测为病灶则惩罚 outside_lesion lesion_mask * (1 - vertebra_union) anatomical_loss np.sum(outside_lesion) * 10.0 # 权重设为 10.0 return rpn_class_loss, rpn_bbox_loss, class_loss, bbox_loss, mask_loss, anatomical_loss提示该损失项未参与梯度反向传播的权重分配即未乘以self.loss_weights而是直接加到总 loss 中。这意味着模型在训练早期会优先学习“病灶不能漂移出椎体”的硬约束避免后期优化陷入局部最优。实际调试时可通过调整anatomical_loss的系数当前为 10.0控制约束强度——系数过大会抑制病灶细节学习过小则无法纠正空间错位。2.3 数据增强策略针对 CT 影像特性定制CT 图像具有高对比度、低噪声、固定 Hounsfield 单位HU范围的特点通用增强方法如随机旋转、色彩抖动反而会破坏解剖结构真实性。utils.py中SpineDataset.load_image()调用的augment_image()函数仅启用三项操作random_crop按比例裁剪保留中心 85% 区域模拟不同扫描视野random_flip仅水平翻转脊柱左右对称垂直翻转会破坏解剖逻辑adjust_contrast在 [-100, 400] HU 范围内线性拉伸窗宽窗位对应软组织窗而非随机 gamma 变换。# utils.py 第 421 行 def augment_image(self, image, mask): # image shape: (H, W), dtype: float32, range: [-1024, 3071] (典型CT) if random.random() 0.5: image np.fliplr(image).copy() mask np.fliplr(mask).copy() if random.random() 0.7: h, w image.shape crop_h, crop_w int(h*0.85), int(w*0.85) start_h random.randint(0, h-crop_h) start_w random.randint(0, w-crop_w) image image[start_h:start_hcrop_h, start_w:start_wcrop_w] mask mask[start_h:start_hcrop_h, start_w:start_wcrop_w, :] # 窗宽窗位调整模拟不同设备参数 window_center random.randint(40, 60) # 软组织窗中心 window_width random.randint(350, 450) # 窗宽 image np.clip((image - (window_center - window_width/2)) / window_width, 0, 1) return image, mask注意adjust_contrast实际执行的是 HU 值截断归一化而非传统图像增强中的亮度/对比度调节。代码中np.clip(..., 0, 1)确保输出始终在 [0,1] 范围避免后续卷积层输入溢出。此操作等效于 DICOM 标准中的 WW/WLWindow Width/Window Level调节是放射科医生阅片的基础操作。3. 从零启动训练环境配置、数据准备与关键参数调优3.1 环境依赖与 GPU 显存优化配置项目要求 Python 3.7、TensorFlow 1.15注意非 TF2.xCUDA 版本需匹配 NVIDIA 驱动。requirements.txt中关键依赖如下包名版本作用替换建议tensorflow-gpu1.15.0主框架必须TF2.x 不兼容 Matterport 代码opencv-python4.5.5图像读写≥4.5.0低于此版本cv2.resize对 float32 支持异常scikit-image0.19.3医学图像处理≥0.18.0用于measure.label分析连通域pycocotools2.0.6COCO API 兼容必须evaluate_coco()依赖提示若使用 RTX 3090/4090 等 Ampere 架构显卡需安装 CUDA 11.0 cuDNN 7.6.5否则tensorflow-gpu1.15.0会报Failed to get convolution algorithm错误。验证命令python -c import tensorflow as tf; print(tf.test.is_gpu_available())返回True。3.2 天池数据集预处理全流程原始数据为 DICOM 序列每例含 50–120 张横断面图像需转换为模型可读格式。samples/tianchi/prepare_data.py提供完整流程# 步骤1解压天池下载的 zip 文件 unzip tianchi_spine_data.zip -d raw_data/ # 步骤2运行预处理脚本自动完成 DICOM→NIfTI→PNG→mask 生成 python samples/tianchi/prepare_data.py \ --dicom_dir raw_data/ \ --output_dir processed_data/ \ --train_ratio 0.7 \ --val_ratio 0.15 \ --test_ratio 0.15该脚本核心逻辑调用pydicom读取每例 DICOM 元数据提取ImagePositionPatient和PixelSpacing计算真实物理尺寸使用nibabel将序列重采样为 1mm×1mm×1mm 各向同性体素通过itk的OtsuThresholdImageFilter自动分割椎体粗略区域再用scikit-image的morphology.remove_small_objects清除噪声最终生成processed_data/train/images/PNG 格式uint80–255和processed_data/train/masks/NPY 格式float32shape(H,W,27)。注意prepare_data.py中--train_ratio参数影响数据集划分但SpineDataset类默认从processed_data/train/加载全部样本。若需自定义划分需修改samples/tianchi/train.py第 89 行dataset_train.load_spine(dataset_dirprocessed_data/train, subsettrain)的subset参数并确保processed_data/下存在对应子目录。3.3 训练命令与超参数调优表启动训练使用samples/tianchi/train.py关键参数说明如下参数默认值推荐值说明--datasetprocessed_data/必填路径数据根目录需包含train/val/子目录--weightscocolast或imagenetcoco加载 COCO 预训练权重适用于迁移学习last从logs/spine2023xxx/加载断点--logslogs/自定义路径日志与模型保存目录建议设为绝对路径--epochs5030–80医疗数据量少通常 30 轮即可收敛超过 50 轮易过拟合--learning_rate0.0010.0005–0.002初始学习率0.001适合 COCO 迁移若从头训练建议0.0005# 示例使用 COCO 预训练权重在 2 块 V100 上训练 python samples/tianchi/train.py \ --dataset /path/to/processed_data/ \ --weights coco \ --logs /path/to/logs/ \ --epochs 40 \ --learning_rate 0.001 \ --gpu_count 2训练过程监控重点rpn_class_loss应在 10 轮内降至 0.1否则检查 RPN anchor 尺寸是否匹配椎体大小config.py中RPN_ANCHOR_SCALES (16, 32, 64, 128, 256)针对 512×512 输入图像mask_loss稳定在 0.15–0.25 区间若持续 0.3 说明病灶标注质量差或增强过度val_loss在第 25 轮后不再下降应停止训练并保存epoch025.h5。4. 模型推理与结果可视化如何验证椎体定位与病灶标注的临床合理性4.1 单图推理脚本的临床级输出解析samples/tianchi/inference.py提供端到端推理接口关键输出不仅是 bounding box 坐标更是结构化临床报告# inference.py 第 156 行 results model.detect([image], verbose0)[0] # results 字典包含 # rois: (N,4) 形状y1,x1,y2,x2 坐标 # class_ids: (N,) 形状类别索引1L1, 2L2,...,26L5,T1-T12,C1-C7,27lesion # scores: (N,) 形状置信度 # masks: (H,W,N) 形状二值掩膜 # vert_labels: [L1,L2,...,Lesion] # 解码后的标签列表运行命令python samples/tianchi/inference.py \ --model /path/to/logs/spine20230501T1200/mask_rcnn_spine_0030.h5 \ --image /path/to/processed_data/test/images/001.png \ --output_dir /path/to/output/输出文件包括001_pred.png叠加预测掩膜的原图椎体用不同颜色病灶用红色高亮001_report.json结构化 JSON含vertebrae数组每个元素含label、confidence、bbox_mm物理坐标和lesions数组含type根据掩膜形状分类、size_mm2、location_relative如“L4/5 椎间盘后缘”。4.2 可视化工具visualize.py的深度定制用法visualize.py不仅绘制热力图还支持解剖关系验证。例如检查 L4/5 椎间盘突出是否符合空间逻辑# visualize.py 第 289 行添加验证函数 def validate_lesion_location(masks, class_ids, vert_labels): lesion_idx np.where(class_ids 27)[0] # 病灶索引 if len(lesion_idx) 0: return No lesion detected lesion_mask masks[..., lesion_idx[0]] # 获取病灶质心 cy, cx ndimage.center_of_mass(lesion_mask) # 查找覆盖质心的椎体 vert_covered [] for i, label in enumerate(vert_labels[:26]): # 前26类为椎体 if masks[..., i][int(cy), int(cx)] 0.5: vert_covered.append(label) if len(vert_covered) 0: return fLesion at ({cx:.1f},{cy:.1f}) not covered by any vertebra elif len(vert_covered) 1: return fLesion overlaps multiple vertebrae: {vert_covered} else: return fLesion located in {vert_covered[0]} region # 调用示例 report validate_lesion_location(results[masks], results[class_ids], results[vert_labels]) print(report) # 输出Lesion located in L4 region提示该验证逻辑嵌入inference.py的display_instances()调用前可自动过滤掉空间逻辑错误的预测结果避免向医生展示违背解剖常识的假阳性。4.3 关键评估指标计算与临床意义映射项目未提供evaluate_coco()的完整实现但samples/tianchi/evaluate.py包含核心指标计算指标计算方式临床意义合格阈值Vertebra Localization AP0.5COCO 标准 APIoU≥0.5椎体定位准确性≥0.85Lesion Segmentation Dice2*A∩B/(Anatomical Consistency Rate病灶质心落入对应椎体掩膜的比例解剖逻辑可靠性≥0.98# evaluate.py 第 92 行 Dice 计算 def compute_dice(y_true, y_pred): intersection np.sum(y_true * y_pred) union np.sum(y_true) np.sum(y_pred) return 2. * intersection / (union 1e-6) # 防止除零 # 批量计算示例 dice_scores [] for i in range(len(val_dataset)): image, image_meta, gt_class_id, gt_bbox, gt_mask val_dataset.load_image(i) results model.detect([image], verbose0)[0] # 提取第27类病灶预测掩膜 lesion_pred results[masks][..., np.where(results[class_ids]27)[0][0]] lesion_true gt_mask[..., 26] # 真实病灶通道 dice_scores.append(compute_dice(lesion_true, lesion_pred)) print(fMean Dice: {np.mean(dice_scores):.3f})注意Dice 系数对小目标敏感若病灶面积 50 像素建议改用Hausdorff Distance评估边界精度scikit-image.measure.hausdorff_distance因其更能反映放射科医生关注的“边缘贴合度”。5. 竞赛提交技巧如何将本地训练模型适配天池平台限制条件5.1 模型轻量化与推理加速策略天池平台对单次推理耗时有严格限制≤3s/例而原始 Mask R-CNN 在 V100 上推理一张 512×512 图像约需 4.2s。model.py中build_detection_graph()函数支持两种加速模式# model.py 第 342 行 def build_detection_graph(self, config, modeinference): # ... 原始图构建 ... if mode inference and config.USE_PRUNING: # 启用通道剪枝移除冗余卷积核 from tensorflow_model_optimization.sparsity import keras as sparsity pruned_model sparsity.prune_low_magnitude( model, pruning_schedulesparsity.PolynomialDecay( initial_sparsity0.2, final_sparsity0.5, begin_step0, end_step10000 ) ) return pruned_model启用方式在config.py中设置USE_PRUNING True并在训练命令中添加--pruning参数。剪枝后模型体积减少 38%推理速度提升至 2.7s/例且 AP0.5 仅下降 0.012。5.2 提交包结构与 Dockerfile 编写规范天池要求提交submission.zip内部结构必须为submission/ ├── model/ │ ├── mask_rcnn_spine.h5 # Keras 模型权重 │ └── config.pkl # SpineConfig 实例 pickle ├── predict.py # 必须入口文件含 main() 函数 └── requirements.txt # 仅限 pip install 依赖predict.py必须实现def main(input_dir, output_dir): # 1. 加载 config.pkl 和 model.h5 # 2. 遍历 input_dir 下所有 PNG 文件 # 3. 对每张图调用 model.detect()保存 JSON 报告到 output_dir pass if __name__ __main__: import argparse parser argparse.ArgumentParser() parser.add_argument(--input_dir) parser.add_argument(--output_dir) args parser.parse_args() main(args.input_dir, args.output_dir)Dockerfile 示例适配天池 GPU 环境FROM nvidia/cuda:11.0-cudnn7-runtime-ubuntu18.04 RUN apt-get update apt-get install -y python3-pip python3-dev COPY requirements.txt . RUN pip3 install --no-cache-dir -r requirements.txt COPY submission/ /app/ WORKDIR /app CMD [python3, predict.py, --input_dir, /input, --output_dir, /output]提示天池平台挂载/input为测试数据目录/output为结果输出目录容器内路径必须严格匹配。requirements.txt中禁止出现tensorflow-gpu1.15必须精确指定tensorflow-gpu1.15.0否则平台安装失败。5.3 竞赛常见扣分点与规避方案扣分场景原因规避方法提交超时模型未剪枝或 batch_size1 未优化使用tf.function装饰detect()方法预编译计算图标签错误class_ids未映射到天池要求的 1–27 编号在predict.py中硬编码映射{1:L1,2:L2,...,26:C7,27:lesion}格式不符JSON 报告缺少lesion_type字段修改inference.py的format_result()函数添加lesion_type herniation if area50 else schmorl判断逻辑内存溢出加载大尺寸图像未 resize在predict.py的load_image()中强制cv2.resize(img, (512,512))验证提交包是否合规的本地命令# 模拟天池环境运行 docker run -v $(pwd)/test_input:/input -v $(pwd)/test_output:/output submission_img # 检查 test_output/ 下是否生成符合 schema 的 JSON 文件 jq .vertebrae[0].label test_output/001.json # 应输出 L1本文还有配套的精品资源点击获取