简介:基于论文复现的InDuDoNet模型Python实现源码,为需要深入理解或二次开发该医学图像处理模型的研究者与开发者提供完整参考。项目聚焦CT图像中器官与病变的准确分割,覆盖数据加载、训练、推理、可视化、评估等完整流程,并针对大鼠股骨、DeepLesion及临床金属伪影等不同数据场景提供对应配置与启动脚本,便于直接迁移或改造。压缩包共61个文件,以44个Python脚本为主体,辅以9个YAML配置文件、2个CSV数据文件及文本说明、MAT数据等;YAML集中管理训练与推理参数,CSV可用于记录实验结果,整体仅918KB,目录结构清晰,易于按模块检索。已有385人学习下载。该复现实现不仅还原了InDuDoNet及其Plus变体的网络结构,还配有先验网络、投影构建、评估指标等工具函数,并带.gitignore与开发日志,适合作为论文复现、算法对比、实验扩展的实用起点。
1. InDuDoNet 复现工程:不是一张网络图,而是一条能跑的CT图像处理流水线
决定复现 InDuDoNet 之前,我最担心的不是论文里的公式,而是这套源码能不能直接跑通。CT 金属伪影去除这个方向,代码通常散落在论文附录和 GitHub 旧仓库里,能拿到一份带完整训练、推理、评估脚本的工程并不常见。这份 Python 实现的 InDuDoNet 源码共61个文件,覆盖了从投影几何构建、FBP 重建、数据加载,到 WNet 网络、Trainer、YAML 配置和评价指标的整条链路。适合正在做 CT 图像分割、去伪影工作,或者想把展开迭代网络迁移到自己数据上的研究者和开发者。接下来我按「文件主线 → 训练链路 → 避坑 → 验证 → 迁移」的顺序把这份工程彻底拆开,重点讲哪些参数能改、哪些地方会翻车。
2. 拆开 61 个文件:执行主线与数据流向
2.1 文件总览:先分清核心代码和外围脚本
拿到压缩包第一步不是翻代码,而是先按目录把文件分类。这套工程里没有复杂的界面层,所有逻辑集中在几个命名清晰的模块里。从根目录结构看,network、trainers、datasets、utils 四个目录构成了主体,剩下的 test_module.py、train.py、train_plus.py 属于入口脚本,YAML 文件属于配置中心,FBP.py 和 build_gemotry.py 承担 CT 成像相关的预处理职责。
| 目录/文件 | 职责 | 关键文件 |
|---|---|---|
| network/ | 模型结构定义 | WNet.py, InDuDoNet.py, InDuDoNet_Plus.py, priornet.py, InDuDoNetRatFemur.py |
| trainers/ | 训练与推理封装 | RatFemurInDuDoNetTrainer.py, DeeplesionInDuDoNetTrainer.py, 各Inference文件 |
| utils/ 与根目录工具 | 工具集与预处理 | CommonUtils.py, LogUtils.py, Visualizer.py, FBP.py, build_gemotry.py |
| datasets/ | 数据集加载 | datasets_plus.py, dataloader.py |
| 入口脚本 | 训练与调试 | train.py, train_plus.py, test_deeplesion.py, test_clinic.py |
| YAML 配置 | 实验参数中心 | InDuDoNet.yaml, InDuDoNetPlus.yaml 等9份 |
这个分类很重要,因为很多人一上来就从 WNet.py 开始读,读到一半发现数据进不去网络,再回头找数据处理,方向就反了。我一般建议按「数据 → 配置 → 网络 → 训练」的顺序读,先把 datasets_plus.py 和 FBP.py 的关系弄清楚,再看 network 里的模型,最后回到 train.py 把整条链串起来。
network 目录里有多个模型文件,其中 WNet.py 是基础组件,InDuDoNet.py 是主模型,InDuDoNetPlus.py 是增强版本,还单独放了 priornet.py 做先验网络。这种拆分在论文复现工程里很常见,它对应了 InDuDoNet 论文里“双域展开”的结构:一个域处理投影数据,另一个域处理图像重建,两个子网络通过迭代互相校准。如果只把注意力放在主模型文件上,容易忽略 priornet.py 里那部分先验约束逻辑,后面训练时会发现 loss 曲线很奇怪,实际上就是先验分支没有被正确接入。Deeplesion、RatFemur 命名直接透露了不同文件对应的数据来源,这也是复现工程里最常见的组织方式——一个场景一套 Trainer 一套配置,互不污染。
2.2 数据流三件套:build_gemotry、FBP 与 datasets_plus
CT 数据输入网络之前要经历三个环节。build_gemotry.py 负责构建扫描几何,包括射线源到探测器的距离、角度采样数、探测器单元数量,这些参数会直接决定投影矩阵的形状。FBP.py 做滤波反投影,把正弦图(sinogram)重建回图像域。InDuDoNet 这类双域网络训练时同时用到投影域和图像域的数据,所以 FBP 这一步不能跳过,它输出的带伪影重建图就是图像域分支的输入。
datasets_plus.py 则把几样东西打包成一个训练样本:原始投影、FBP 重建的带伪影图像、金属掩膜、对应的干净参考图。以 RatFemur(大鼠股骨)场景为例,样本里必须带上金属掩膜,因为去除金属伪影的前提是先知道金属在图像里的位置。代码逻辑大致如下:
# datasets_plus.py 中常见的样本构建逻辑(示意) # 1. 从 .mat 或 .npy 读取投影数据 sinogram # 2. 用 build_gemotry 提供的几何参数生成投影矩阵 # 3. 用 FBP 把 sinogram 重建为初始图像 fb_image # 4. 将 sinogram、fb_image、metal_mask 组合为 dict 返回 sample = { "sinogram": sinogram_tensor, # 投影域输入,形状通常为 (角度数, 探测器数) "fb_image": fb_image_tensor, # 图像域输入,带金属伪影的 FBP 重建图 "metal_mask": mask_tensor, # 金属区域掩膜,用于约束损失计算 "target": clean_image_tensor # 干净参考图,训练时计算损失 }这里的核心参数是 build_gemotry.py 输出的探测器数量和角度数。如果角度数是 640,那么 sinogram 的形状通常是(角度数, 探测器单元数),也就是项目里 dataset_py_640geo.yaml 这个文件名中“640geo”的含义。常见做法是用扇形束扫描几何,角度步长等于 360 除以角度数,探测器单元间距决定了重建图像的分辨率。改数据集时最常翻车的就是这里:换了 CT 设备后角度数和探测器数量变了,但 YAML 里没同步更新,FBP 出来的图像就是模糊的,边缘还会出现条状伪影。
数据流到这里形成闭环:build_gemotry 定义几何,FBP 用几何做重建,datasets_plus 把重建结果和标签打包成 PyTorch Dataset,之后 dataloader.py 负责批量加载。读源码时先看这三个文件的调用顺序,就能理解为什么 train.py 里不需要手动调用 FBP——数据加载阶段已经完成了。如果你想验证自己的数据能不能用,优先检查这三个环节的输出形状是否匹配,比直接读模型代码更快。
2.3 YAML 配置中心:9 份文件各管一个场景
YAML 在这套工程里不是装饰,而是训练入口的配置文件。9 份 YAML 对应不同的实验场景,其中 InDuDoNet.yaml 是基础版主模型配置,InDuDoNetPlus.yaml 对应增强版,RatFemurInDuDoNetTrain.yaml 和 DeeplesionInDuDoNetTrain.yaml 分别服务大鼠股骨和 DeepLesion 两个公开数据集,RealRatFemurInDuDoNetInference.yaml 和 DeeplesionInDuDoNetInference.yaml 则用于推理阶段。JiGan.yaml 从命名看可能是“级联”或扩展实验的配置,复现时先不用管它。
配置内容通常分三块:数据路径、几何参数、训练超参数。数据路径最容易出错,Windows 下用反斜杠路径时,YAML 解析会把\d之类的内容误判成转义字符,我一般建议全部改成正斜杠或者用原始字符串。几何参数必须和 build_gemotry 保持一致,角度数量、探测器数量、图像尺寸这些值不匹配时,训练不一定会直接报错,但 loss 会异常高,因为网络输入和标签根本不在同一个坐标系里。
从实践来看,复现时最稳妥的做法是先看 YAML 里配的是哪个 Trainer,再去看 trainers 目录里对应 Trainer 的初始化参数。这套工程的配置集中管理思路很清晰,后续迁移到自己的数据时只需要新写一份 YAML,不用改主代码,这是它值得下载的一个重要原因。注意 YAML 里一般还有 device 字段,默认是 cuda:0,没有 GPU 的机器要手动改成 cpu,否则一启动就报设备错误。
3. 训练链路:从 train.py 到 Trainer 的调用关系
3.1 入口脚本:train.py 和 train_plus.py 的分工
工程根目录放了 train.py 和 train_plus.py 两个入口,对应基础版和 Plus 版模型。两者结构高度一致:先读取 YAML 配置,再通过 root_config.py 把配置加载成 Python 对象,然后根据配置选择对应的 Trainer,最后调用 Trainer.train()。直接跑训练的命令大概是:
python train.py --config InDuDoNet.yaml如果要用增强版,就换成 train_plus.py 和 InDuDoNetPlus.yaml:
python train_plus.py --config InDuDoNetPlus.yaml注意这里的 --config 参数可能不是标准 argparse 的必需参数,工程里也可能直接用固定路径读取 YAML,具体要看 train.py 开头几行的写法。我倾向于认为它保留了读取外部 YAML 的能力,因为 9 份配置文件都放在根目录,如果不支持 --config 指定,维护这么多份 YAML 就没有意义。
进入 root_config.py 后,配置会被解析成几类:数据配置(路径、几何参数)、模型配置(网络层数、通道数)、训练配置(batch_size、学习率、epoch)、日志配置(输出目录、保存频率)。这个模块是整个工程的“黑匣子”入口,很多复现者直接跳过它去看模型,结果发现怎么改 YAML 都不生效——大概率是 root_config 里把配置值硬编码覆盖了。遇到这种情况先改 root_config.py 里的默认值,再改 YAML,两处保持一致。
3.2 Trainer 分层:训练逻辑和网络解耦
trainers 目录里每个文件对应一个实验场景,这是这套工程最值得借鉴的设计。比如 RatFemurInDuDoNetTrainer 处理大鼠股骨数据,DeeplesionInDuDoNetTrainer 处理 DeepLesion 数据集,CLINIC_metal 目录里还有针对牙科金属伪影的预处理逻辑。每个 Trainer 内部会做四件事:加载数据集、构建优化器和学习率调度器、执行前向反向传播、周期性保存 checkpoint。
# Trainer 内部训练循环的典型步骤(示意) # 1. 从 dataloader 取一个 batch # 2. 把 sinogram、fb_image、metal_mask 送入网络 # 3. 向前传播,得到重建结果和分割结果 # 4. 用损失函数计算 loss,反向传播更新参数 # 5. 每 N 个 epoch 保存一次 checkpoint 到 result 目录 optimizer.zero_grad() output = model(sinogram, fb_image, metal_mask) loss = criterion(output, target) loss.backward() optimizer.step()这里的 loss 通常不是一个简单的 L1 或 MSE,而是图像域损失加上投影域损失的加权和。InDuDoNet 论文的核心贡献就在这里:它让两个域交替优化,投影域的信息帮助图像域去掉金属伪影,图像域的反馈又反过来修正投影域的误差。所以训练时如果只盯着一张输出图看效果,往往看不出网络在收敛,要把投影域的 loss 也打出来一起观察。
每种 Trainer 大概率还实现了 validate 方法,在训练过程中周期性地在验证集上跑指标。这一步很容易被新手忽略,因为 train.py 主流程里可能没有强制调用。复现时建议把验证逻辑打开,每 5 个或 10 个 epoch 评估一次 PSNR 和 SSIM,否则训练结束才发现过拟合,又没有中间结果的曲线,排查起来非常被动。
3.3 关键超参数与硬件配置
从 YAML 的常见写法来看,这套工程的超参数一般集中在:batch_size、learning_rate、num_epochs、checkpoint_dir、device。GPU 显存不够时,很多人第一个想到的是调小 batch_size,但要注意 batch_size 变化后 learning_rate 也要相应调整,否则收敛速度会明显变慢。另一个容易被忽略的参数是 num_workers,它控制 dataloader 的子进程数量,在 Windows 上设成大于 0 偶尔会有多进程启动报错,我一般直接设 0,慢一点但不影响训练。
实验环境的 Python 版本建议 3.8 以上,因为工程里用了较多 dataclass 和 pathlib 相关的写法,太老的 Python 版本会直接语法报错。依赖安装建议按 requirements.txt 走,其中 numpy 是绕不开的底层库,CT 数据基本都是数组操作,先确保 numpy、scipy、PyTorch 三个库的版本兼容,再装其他依赖。如果之前机器上装过老版本 PyTorch,建议先卸载干净再装,避免出现 libcudnn 版本对不上这类玄学问题。
4. 避坑排查:复现 InDuDoNet 时最容易翻车的五个点
4.1 .mat 数据加载失败:scipy 与 h5py 的选择
现象:运行 dataloader.py 加载 test_project.mat 时直接报错,提示 "Unable to read MAT-file"。
原因:旧版 MATLAB 保存的 .mat 文件可以用 scipy.io.loadmat 读取,但新版 MATLAB 默认保存为 HDF5 格式,scipy 读不了。
解决:先用下面这行代码判断文件头格式。
python -c "with open('test_project.mat','rb') as f: print(f.read(128))"如果文件头是 HDF5,就用 h5py 读取;如果是 MATLAB 的 MAT 格式,再用 scipy.io.loadmat。读取逻辑最好封装成一个函数,在 datasets_plus.py 里做兼容,以后换数据集也能复用。
4.2 YAML 路径转义翻车
现象:Windows 下配置的 data_path 写成D:\dataset\train,运行时报找不到文件,或路径包含乱码。
原因:YAML 解析时反斜杠是转义字符,\d、\t这些组合会被解析成特殊含义,路径就变了。
解决:改配置时统一用正斜杠D:/dataset/train,或者给路径加双引号并在 Python 里用 pathlib.Path 规范化。这是没有技术含量但出现频率最高的坑,几乎每个迁移数据集的同事都踩过一次。
4.3 FBP 几何参数不匹配导致图像反转
现象:训练正常跑,但把重建结果可视化后发现图像上下颠倒或左右翻转,指标却还不差。原因是网络学到了翻转后的分布,PSNR 可能还行,但图像的物理方向不对。
原因:build_gemotry 里定义的角度方向和 FBP 里投影的角度方向不一致,比如从顺时针改成逆时针,重建出来的图像就会翻转。
解决:先用一张标准测试图,比如一个明显偏置的方块,跑一遍 FBP,确认重建方向和原图一致再训练。项目里 ProjectionTest.py 大概率就是干这个的,复现时先跑它,不要跳过。
4.4 显存溢出与 batch_size 调参
现象:训练第一个 epoch 就报 CUDA out of memory,显存占用直接拉满。
原因:InDuDoNet 的双域结构在前向传播时同时保存了投影域和图像域的中间特征,显存需求比普通分割网络高一个量级。盲目调小 batch_size 后又发现收敛变慢。
解决:先按 batch_size=1 跑通完整数据链路,确认单样本占多少显存,再倍增 batch_size。同时把 dataloader 的 pin_memory 设为 True 可以稍微缓解带宽瓶颈,但真正治本的是减小输入图像尺寸和网络通道数。
4.5 日志文件不输出导致训练像卡死
现象:train.py 启动后终端长时间没有输出,看起来像是死循环,但 GPU 占用率是满的。
原因:LogUtils.py 配置的日志输出频率太低,或者标准输出缓冲没有 flush,训练在跑但你看不到进度。
解决:把日志打印间隔调小,比如每个 epoch 都打印一次 loss,并在 print 里加 flush=True。如果是按迭代打印,不要设成每 1000 步一次,改成每 50 步一次,能更快感知训练是否正常。
5. 推理与验证:test_deeplesion.py 和 EvaluationMetrics.py 的闭环用法
5.1 测试脚本的矩阵:训练完怎么跑推理
工程里 test_deeplesion.py 负责在 DeepLesion 数据集上做推理,test_module.py 和 test_something.py 像是调试用的临时脚本,test_clinic.py 面向临床数据。跑推理前一定要改对配置,用 DeeplesionInDuDoNetInference.yaml 而不是训练用的 YAML,因为推理阶段需要加载的是已保存的 checkpoint,而不是从零开始训练。
python test_deeplesion.py --config DeeplesionInDuDoNetInference.yaml推理脚本内部一般分四步:加载配置、加载模型权重、遍历测试集跑前向、把结果存入 inference_result_csv 目录。这里最容易忽略的是训练和推理时数据预处理要完全一致,特别是 FBP 重建这一步,如果推理时跳过了重建,输入网络的图像就和训练时对不上,结果直接崩掉。
5.2 评价指标:EvaluationMetrics.py 里有什么
CT 图像重建和分割的评价指标不只是准确率。EvaluationMetrics.py 里大概率实现了 PSNR(峰值信噪比)、SSIM(结构相似性)这两个重建任务核心指标,如果涉及分割,可能还有 Dice 系数和 IoU。PSNR 衡量图像灰度差距,SSIM 衡量结构相似度,这两个指标在金属伪影去除任务里经常出现“一个涨一个跌”的情况,因为去伪影过程往往伴随过平滑,SSIM 掉分但 PSNR 涨分。
# PSNR 和 SSIM 的计算入口(示意) from EvaluationMetrics import calculate_psnr, calculate_ssim psnr = calculate_psnr(output, target) ssim = calculate_ssim(output, target)跑完推理后,把这些指标输出到 CSV 文件,就是 inference_result_csv 目录里的内容。CSV 的好处是可以用 pandas 汇总多个样本的结果,算均值和方差,画曲线对比不同 checkpoint 的效果。很多复现者忽略这一步,只看可视化图像,这容易误判——视觉上“看着差不多”的图,PSNR 可能差了 2dB,在论文里这就是“显著提升”和“没有提升”的区别。
5.3 可视化链路:Visualizer、Html 与日志的关系
VisualizerUtil.py 和 Visualizer.py 承担了中间结果的保存功能,Html.py 则可以把训练过程中的图像整理成 HTML 页面方便对比。CT 任务里可视化尤其重要,因为金属伪影经常表现为局部的放射状条纹,光看数值无法定位问题。Visualizer 的作用就是把网络输出的图像、标签图像、重建误差图拼在一张图里,一眼看出伪影残余集中在哪个区域。
日志方面,development_log.txt 记录了开发过程中的版本变更,这对复现者来说比有些论文本身还有用——通常能看到作者改过哪些参数、修过哪些 bug,这些信息经常是论文摘要里不写的。看日志时重点找 FBP 相关和 loss 相关的改动记录,这两处改过的参数往往就是复现的关键开关。
6. 进阶迁移:把 InDuDoNet 接到自己的 CT 数据上
迁移第一步是改 build_gemotry.py 里的几何参数。以扇形束扫描为例,需要确认四组数:源到旋转中心的距离、探测器到中心的距离、探测器单元数、角度采样数。这四组数可以从 CT 设备的标定文件里拿到,也可以在原始投影数据里反推。反推几何参数的常用方法是用一组已知大小的高密度小球做扫描,从投影图像里测量小球直径随角度的变化,再对照 FBP 重建结果微调。跑 ProjectionTest.py 的意义不只是验证代码能跑,而是验证几何参数对不对。
改完 build_gemotry 后,把对应数字填进新 YAML 的 geometry 段。dataset_py_640geo.yaml 就是 640 角度几何的范例,新设备角度数变成 720 时,直接复制一份 YAML,只改角度数和探测器数即可。同步是关键:只改一处而 FBP 和 datasets_plus 拿到的还是旧值,网络前向维度虽然可能不报错,但重建图边缘会有条状伪影,这类问题从指标上很难看出来。
新数据集的样本格式要和现有预期对齐。从数据流看,每个样本至少需要四部分:sinogram、fb_image、metal_mask、target。如果只有重建好的图像而没有原始投影,build_gemotry 这段流程就得砍掉,直接从 FBP 重建图开始喂网络,但这样投影域分支就没有输入,Plus 版的增强能力会打折扣。我的建议是优先保留投影数据,因为双域网络的根基就是投影域和图像域的互校准。
训练策略上,固定优化器和学习率不动,先跑 5 个 epoch,只看 loss 曲线是否下降。如果 loss 前 5 个 epoch 波动幅度大但总体向下,说明链路通;如果 loss 完全不动,大概率是学习率太低或数据范围没归一化。CT 图像的像素值范围差异很大,有的设备输出 0 到 4095 的整数,有的已经归一化到 0 到 1,喂进网络前必须统一,否则梯度会剧烈震荡。
验证方法上,建议每次成功迁移后都把训练时最后一个 epoch 的 checkpoint 单独保存,推理脚本和评估脚本都基于这个 checkpoint 输出,保证可复现。我之前把一个牙科 CBCT 数据迁进这套框架时,漏改了探测器单元数,结果 FBP 重建图像分辨率不对,前向传播维度直接报错,排查花了大半天。从那以后我每次换数据集都强制走一遍「几何标定 → 单样本前向 → 单 epoch 验证」的三步流程,确认三步都通过才展开全量训练。这份源码的价值在于它把主干流程搭好了,你只需要改几何参数和 YAML,省去从零搭双域展开网络的功夫,希望帮到你。
本文还有配套的精品资源,点击获取