AI

医学影像三维重建:从 nnUNet 分割到 VTK.js 实时渲染

基于 nnUNet 和 VTK.js 构建医学影像三维重建平台的开发

一、项目背景

医学影像三维重建是临床诊断和手术规划的重要工具。传统的三维重建软件如 3D Slicer 功能强大,但存在以下痛点:

  • 部署复杂:需要安装大量依赖,难以快速部署
  • 交互门槛高:学习曲线陡峭,医生需要专门培训
  • 无法 Web 化:无法集成到医院的 PACS 系统或远程会诊平台

本项目 Medical CT Viewer 的目标是构建一个轻量级的 Web 端医学影像三维重建平台:

  • 前端:React + VTK.js,支持浏览器内实时 3D 渲染
  • 后端:FastAPI + nnUNet + VTK Python,提供分割和转换服务
  • 核心功能:自动分割、三维重建、MPR 多平面重建

二、技术架构

┌─────────────────────────────────────────────────────────────────┐
│                        前端 (React + VTK.js)                     │
├─────────────────────────────────────────────────────────────────┤
│  Viewer3D                          MPRViewer                     │
│  ├── VolumeRenderer.ts             MPRRenderer.ts                │
│  │   ├── Marching Cubes            ├── ImageMapper              │
│  │   ├── Windowed Sinc             ├── ImageSlice               │
│  │   └── Gaussian Blur             └── 窗宽窗位交互              │
│  └── useViewerStore.ts                                            │
└─────────────────────────────────────────────────────────────────┘
                              ▲
                              │ VTI / JSON
                              ▼
┌─────────────────────────────────────────────────────────────────┐
│                     后端 (FastAPI + VTK Python)                   │
├─────────────────────────────────────────────────────────────────┤
│  app_nnunet.py          nii_to_vti.py          label_map.py     │
│  ├── 5 折集成预测        ├── NIfTI → VTI        ├── 标签融合     │
│  └── 异步执行            └── 窗宽窗位计算        └── 颜色映射     │
└─────────────────────────────────────────────────────────────────┘

三、核心实现

3.1 nnUNet 5 折集成预测

nnUNet 默认训练 5 折,推理时可以使用单折或 5 折集成。5 折集成能显著提升分割精度:

# backend/segmentation/app_nnunet.py
async def run_nnunet_predict(input_dir: str, output_dir: str, 
                             gpu_id: int = 0, folds: list[int] = None) -> str:
    if folds is None:
        folds = [0, 1, 2, 3, 4]  # 默认 5 折集成
    
    cmd = [
        "nnUNetv2_predict",
        "-i", input_dir,
        "-o", output_dir,
        "-d", str(DATASET_ID),
        "-c", CONFIGURATION,
        "-f", *[str(f) for f in folds],  # 多折预测
        "--disable_tta",  # 禁用 TTA 加速
    ]
    
    # 异步执行,避免阻塞
    process = await asyncio.create_subprocess_exec(*cmd, ...)

权衡:5 折集成精度高但推理慢(约 5 倍时间),实际使用时可根据场景选择:

场景 推荐配置 推理时间
批量处理 5 折 ~5 分钟
实时预览 单折 (fold=0) ~1 分钟

3.2 Marching Cubes + Windowed Sinc 表面重建

直接对分割标签进行 Marching Cubes 会产生明显的锯齿。我们采用了三阶段处理:

Binary Voxel → Gaussian Blur → Marching Cubes → Windowed Sinc → Smooth Mesh

Gaussian Blur 预处理

// frontend/src/vtk/VolumeRenderer.ts
private _gaussianBlur3D(data: Uint8Array, dims: number[], sigma: number): Float32Array {
  // sigma=1.0 平衡平滑度与细节保留
  const kernelRadius = Math.ceil(2.5 * sigma)
  // ... 三维高斯模糊实现
}

Windowed Sinc 平滑

// 参数映射:smoothness 0→1 → passBand 0.5→0.01
const passBand = 0.5 * Math.pow(0.02, smoothness)
const iterations = Math.round(10 + smoothness * 40)

const smoother = vtkWindowedSincPolyDataFilter.newInstance({
  numberOfIterations: iterations,
  passBand,  // 越小越平滑
  nonManifoldSmoothing: true,
  normalizeCoordinates: true,
})

这是 3D Slicer 使用的同款算法,能生成接近医学级质量的表面网格。

3.3 标签融合机制

模型输出的细小静脉(标签 6、8)单独重建效果较差,需要融合到下腔静脉(标签 1):

# backend/vtk_viewer/core/label_map.py
LABEL_MERGE_MAP = {
    6: 1,  # Right_TinyVein → Postcava
    8: 1,  # Left_TinyVein → Postcava
}

融合实现(在 label map 构建时):

for label_id in unique_labels:
    target_label_id = LABEL_MERGE_MAP.get(label_id, label_id)
    seg_dicts.append({
        "path": str(seg_path),
        "label_id": target_label_id,      # 写入目标标签值
        "is_multilabel": True,
        "extract_label": label_id          # 从原始标签提取
    })

关键点:使用 extract_label 指定从哪个原始标签提取,label_id 指定写入哪个目标标签。

3.4 MPR 多平面重建

三视图(Axial/Coronal/Sagittal)共享同一份 CT 数据和 Label 数据:

// frontend/src/components/ViewerMPR/MPRViewer.tsx
const sharedCTData: { data: any | null } = { data: null }
const sharedLabelData: { data: any | null } = { data: null }

async function loadSharedData(caseId: string): Promise<void> {
  // 只加载一次,所有视图共享
  const ctData = await loadCT(caseId)
  sharedCTData.data = ctData
  
  const labelData = await loadLabelMap(caseId)
  sharedLabelData.data = labelData
}

四、踩坑记录

4.1 VTK.js 窗宽窗位影响分割颜色

问题:拖动 MPR 视图调整窗宽窗位时,分割的颜色会跟着变化。

原因vtkInteractorStyleImage 默认 currentImageNumber=-1,会操作最后一个添加的 actor。在 MPR 视图中,label overlay 是后添加的,所以拖动时修改的是 label 的颜色而非 CT。

解决

// frontend/src/vtk/MPRRenderer.ts
const style = vtkInteractorStyleImage.newInstance()
style.setInteractionMode('IMAGE_SLICING')
// 强制操作第一个 actor(CT)
style.setCurrentImageNumber(0)

4.2 标签 6、8 未重建

问题:融合后的标签 6、8 没有出现在 3D 视图中。

原因:原始逻辑检测到标签在 LABEL_MERGE_MAP 中就 continue 跳过了:

# 错误的实现
for label_id in unique_labels:
    if label_id in LABEL_MERGE_MAP:
        continue  # 直接跳过,导致缺失

解决:使用 extract_label 提取原始值,写入目标标签:

# 正确的实现
for label_id in unique_labels:
    target_label_id = LABEL_MERGE_MAP.get(label_id, label_id)
    # 提取原始标签,写入目标标签
    seg_dicts.append({
        "extract_label": label_id,
        "label_id": target_label_id,
    })

4.3 Node.js 版本兼容

问题:前端启动报错 SyntaxError: Unexpected reserved word

原因:Vite 5 需要 Node.js 18+,系统默认版本 v12.22.9 太旧。

解决

# 使用 nvm 切换版本
source ~/.nvm/nvm.sh
nvm use 18
npm run dev

五、性能优化

5.1 Gaussian Blur 参数调优

sigma 值 效果 适用场景
0.5 轻微平滑,细节保留好 高精度重建
1.0 平衡平滑与细节 推荐默认值
1.5 强平滑,细节丢失 快速预览

5.2 Marching Cubes 异步执行

// 让出主线程,保持 UI 响应
for (const label of labels) {
  await new Promise<void>(resolve => setTimeout(resolve, 0))
  runMarchingCubes(label)
}

5.3 VTI 压缩

# backend/vtk_viewer/core/nii_to_vti.py
writer = vtk.vtkXMLImageDataWriter()
writer.SetCompressorTypeToZLib()  # ZLib 压缩
writer.Write()

六、经验总结

  1. Web 端医学影像的核心不是算法创新,而是工程落地。Marching Cubes 和 Windowed Sinc 都是成熟算法,但把它们组合成一个流畅的 Web 应用需要大量的工程优化。

  2. 标签融合的设计思路:让模型自由输出所有标签,后处理阶段再决定是否合并。这样模型架构不需要改动,灵活性更高。

  3. VTK.js 交互的坑:默认行为往往不符合预期。比如 currentImageNumber=-1 操作最后一个 actor,而我们期望操作第一个。遇到这类问题,优先查阅 VTK.js 源码。

  4. 异步处理的重要性:无论是后端的 nnUNet 推理还是前端的 Marching Cubes,都需要异步执行,否则会阻塞主线程导致 UI 卡死。

  5. 调试医学影像的三板斧

    • 打印数据形状和数值范围
    • 用 Matplotlib 可视化中间结果
    • 对照 3D Slicer 的输出验证正确性