融合跨模态注意力与边界细化的心脏磁共振图像分割框架

A Cross-Modal Attention and Boundary Refinement Framework for Cardiac MRI Segmentation

---

后续的缺失模态稳健多序列 CMR 研究已独立为 MyoRouter-MD。两个项目拥有独立源码、实验结果与复现流程，本仓库继续保留原 coarse-to-fine 分割工程。

1. 项目背景与目标

本项目基于 MyoPS 2020 数据集，构建 双阶段 coarse-to-fine 心肌病灶分割框架。

• Stage 1（粗分割）：使用 PyMIC 2D U-Net 对多序列 CMR 进行心脏区域粗定位与初始分割，生成伪标签作为第 4 通道输入。
• Stage 2（精分割）：使用自定义的 CrossModalBoundaryNet，融合跨模态注意力与边界细化，实现高精度病灶分割。

分割目标（6 类）：背景(0) / 心肌 MYO(1) / 左心室 LV(2) / 右心室 RV(3) / 水肿 Edema(4) / 梗死 Scars(5)

核心贡献：

1. 多模态独立编码器保持模态特异性，避免 Early Fusion 信息混淆
2. Cross-Modal Attention Fusion 自适应融合 C0/LGE/T2/粗分割 mask
3. Boundary Refinement Module 利用边缘注意力增强小病灶边界精度
4. 综合损失 $L = L_{\text{dice}} + L_{\text{focal}} + \lambda L_{\text{boundary}}$

---

2. 项目结构

├── CodeForDataProcessing/           # Stage 1 粗分割数据处理
│   ├── PathConfig.py                # 全局路径配置
│   ├── CoarseStageCropping.py       # 粗定位裁剪
│   ├── CsvPreprocessing.py          # CSV 索引生成 & 5-fold 划分
│   ├── SegmentationAgentRunner.py   # PyMIC 训练/测试 Runner
│   ├── SegmentationAgentTester.py   # PyMIC 测试 Runner
│   ├── Postprocessing.py            # 连通域后处理
│   ├── FineStageCropping.py         # 精细裁剪 -> nnU-Net 格式
│   └── CreateDatasetJson.py         # 生成 dataset.json
│
├── FineStageModel/                  # Stage 2 精分割模型（核心）
│   ├── __init__.py
│   ├── MultimodalEncoder.py         # 多模态独立编码器
│   ├── CrossModalAttention.py       # 跨模态注意力融合模块
│   ├── BoundaryRefinement.py        # 边界细化模块
│   ├── SegmentationNetwork.py       # 完整网络 CrossModalBoundaryNet
│   ├── Losses.py                    # Dice + Focal + Boundary Loss
│   ├── Dataset.py                   # 多模态 CMR Dataset
│   ├── Train.py                     # 5-fold 训练脚本
│   ├── Predict.py                   # 推理脚本（5-fold 集成）
│   └── Evaluate.py                  # 评估脚本 (Dice/HD95/IoU)
│
├── CodeForFinalResult/
│   ├── GetFinalTest.py              # 最终结果导出
│   └── PredSuffix.py                # 文件重命名
│
├── Tools/
│   ├── ReportTrainingProgress.py    # 训练进度报告
│   └── VisualizeTrainingCurves.py   # 训练曲线可视化
│
├── mamba_ssm/                       # Mamba 兼容层（stub）
│   └── __init__.py
│
├── Config/
│   ├── train_val.cfg                # 粗分割训练配置
│   ├── test.cfg                     # 粗分割测试配置
│   ├── evaluation.cfg               # 粗分割评估配置
│   └── fine_stage.yaml              # 精分割超参数配置
│
├── Model/                           # 预训练模型
│   ├── unet2d/                      # 粗分割 PyMIC 5-fold checkpoints
│   └── fine_stage/                  # 精分割模型 checkpoints (训练后生成)
│
├── RunFullPipeline.sh               # 全流程自动化脚本
└── README.md

---

3. 环境依赖

pip install torch torchvision SimpleITK numpy scipy nibabel medpy tensorboard pyyaml batchgenerators pymic

• Python >= 3.8
• PyTorch >= 2.0
• GPU：NVIDIA RTX 4090 24GB（推荐）或同等显存

---

4. 数据流与完整流程

整体流程图

Raw Data -> [Stage 1] Coarse Crop -> Preprocessed Data -> PyMIC UNet2D Train/Test
                                                            |
                                                    Coarse Predictions (伪标签)
                                                            |
         [Stage 2] Fine Crop -> Task112_MyoPS (4 通道: C0/LGE/T2/coarse_seg)
                                       |
                              CrossModalBoundaryNet Train -> Predict -> Evaluate
                                       |
                                 Final Results

---

5. Stage 1：粗分割 (Coarse Stage)

5.1 数据预处理

# 粗定位裁剪
python CodeForDataProcessing/CoarseStageCropping.py

# 生成 CSV 索引 & 5-fold 划分
python CodeForDataProcessing/CsvPreprocessing.py

5.2 训练与推理

# 训练（fold 1~5）
python CodeForDataProcessing/SegmentationAgentRunner.py train Config/train_val.cfg 1

# 验证集推理 + 评估 + 后处理
python CodeForDataProcessing/SegmentationAgentRunner.py test Config/train_val.cfg 1
pymic_evaluate_seg Config/evaluation.cfg
python CodeForDataProcessing/Postprocessing.py result/unet2d result/unet2d_post

# 测试集推理 + 后处理
python CodeForDataProcessing/SegmentationAgentTester.py test Config/test.cfg
python CodeForDataProcessing/Postprocessing.py result/unet2d_test result/unet2d_test_post

注意：在执行前请确认 CodeForDataProcessing/PathConfig.py 中的路径配置正确，并安装 PyMIC 库。

---

6. Stage 1->2 过渡：精细裁剪

# 根据粗分割 mask 裁剪，生成 nnU-Net 格式 4 通道数据
python CodeForDataProcessing/FineStageCropping.py

# 生成 dataset.json
python CodeForDataProcessing/CreateDatasetJson.py

输出目录：Task112_MyoPS/，包含：

• imagesTr/CaseXXX_{0000,0001,0002,0003}.nii.gz（C0 / LGE / T2 / coarse_seg）
• labelsTr/CaseXXX.nii.gz（6 类 GT）
• imagesTs/（测试集）

---

7. Stage 2：精分割 (Fine Stage — CrossModalBoundaryNet)

7.1 网络架构

输入 (B, 4, H, W)
  |
MultimodalEncoder: 4 个独立编码器 x 4 尺度 [32, 64, 128, 256]
  |
MultiScaleFusion: 每尺度 CrossModalAttention (Multi-Head Attention + FFN)
  |
UNet Decoder: 3 层上采样 + skip connection
  |
BoundaryRefinementModule: Edge Attention + Edge-aware Conv + 辅助边缘监督
  |
Segmentation Head -> (B, 6, H, W)

7.2 训练

# 训练全部 5 折
python -m FineStageModel.Train --config Config/fine_stage.yaml

# 训练单个 fold
python -m FineStageModel.Train --config Config/fine_stage.yaml --fold 0

训练参数（Config/fine_stage.yaml）：

• Optimizer: AdamW (lr=1e-4, weight_decay=1e-4)
• Scheduler: CosineAnnealingLR (T_max=200, eta_min=1e-6)
• Loss: Dice + Focal + 0.5 x Boundary
• Batch size: 4, Patch size: 256x256
• 5-fold 交叉验证, 200 epochs

7.3 推理

# 5-fold 集成推理
python -m FineStageModel.Predict --config Config/fine_stage.yaml

7.4 评估

python -m FineStageModel.Evaluate \
    --pred_dir result/fine_stage/predictions \
    --gt_dir /path/to/Task112_MyoPS/labelsTr \
    --num_classes 6

评估指标：Dice / HD95 / IoU / Precision / Recall（逐类别 + 汇总）

---

8. 最终结果导出

# 导出到最终目录
python CodeForFinalResult/GetFinalTest.py result/fine_stage/predictions result/fine_stage/final

# 重命名为比赛格式 CaseXXX_pred.nii.gz
python CodeForFinalResult/PredSuffix.py

---

9. 消融实验

通过修改 Config/fine_stage.yaml 中的模型开关即可运行消融实验：

实验	use_shared_encoder	use_cross_attn	use_boundary
Baseline (Early Fusion UNet)	true	-	false
+ Multimodal Encoder	false	false	false
+ Cross-Modal Attention	false	true	false
+ Boundary Refinement (Full)	false	true	true

---

10. 预训练模型

模型	路径	说明
PyMIC 2D UNet (Stage 1)	Model/unet2d/fold_{1-5}/	粗分割 5-fold
CrossModalBoundaryNet	Model/fine_stage/fold_{0-4}/	精分割（训练后生成）

---

References:

[1] Zhai, S., Gu, R., Lei, W., & Wang, G. (2020). Myocardial edema and scar segmentation using a coarse-to-fine framework with weighted ensemble. In Myocardial Pathology Segmentation Combining Multi-Sequence CMR Challenge (pp. 49-59). Springer.