MICCAI STSR 2025 - Task 2 注册 Baseline Docker 镜像描述

镜像版本: v1.0

构建日期: 2025-04-16

联系方式: SemiTeethSegChallenge@aliyun.com

概述

• 赛事主页：https://www.codabench.org/competitions/6470/
• 挑战赛主页：https://songhen15.github.io/STSdevelop.github.io/miccai2025/index.html

计算机辅助诊断工具在现代牙科实践中越来越流行，特别是用于治疗计划或全面预后评估。随着口头数字化和三维扫描技术的进步，锥形束计算机断层扫描（CBCT）和口腔内扫描（IOS）已广泛应用于临床牙科处理中。INTRAORAL扫描（IOS）模型可提供高分辨率，可提供高分辨率，仅依靠IOS数据来依靠IOS数据，可带来不可能的根源信息的风险。相反，CBCT提供了有关牙冠和根部的全面信息，但其扫描分辨率较低，并且包含更多冗余信息。因此，结合两者可以更好地利用它们的互补优势，从而提高牙科治疗的精度和效率。

本 Docker 镜像旨在为 MICCAI STSR 2025 挑战赛 - Task 2（牙冠和牙根配准）提供一个开箱即用的深度学习基线（Baseline）解决方案。镜像内预装了所有必要的依赖库（PyTorch, PyTorch Lightning, MONAI, Trimesh, Hydra 等），并包含了两种不同的配准基线方法的完整代码工程，方便参赛者快速启动并进行开发和实验。

任务目标: 利用深度学习预测一个 4x4 刚体变换矩阵，将一个未对齐（模拟的）口内扫描 STL 文件（代表牙冠表面）精确地对齐到其在 CBCT 扫描坐标系中的正确位置（该位置由我们提供的 _gt.npy 真值矩阵定义）。

包含的 Baseline 方法

本镜像集成了两种不同思路的 Baseline 方法，存放于镜像内的 /workspace/ 目录下：

1. 多模态配准 Baseline (/workspace/multi_modal_registration_baseline/)

   • 方法: 该方法同时利用 CBCT 体数据 和 STL 点云 作为输入。
   • 网络结构:
     • 使用一个 3D 卷积神经网络（基于 MONAI 的 EfficientNetBN-B0 3D）作为 CBCT 编码器，提取体积特征。
     • 使用一个 PointNet 网络作为 STL 点云编码器，提取表面几何特征。
     • 将两种模态提取出的特征进行拼接（Concatenation）融合。
     • 通过一个多层感知机（MLP）回归头，从融合特征中预测 7 个参数（4 个四元数 + 3 个平移）来表示目标变换矩阵。
   • 优势: 理论上能利用 CBCT 提供的内部结构和密度信息，可能对包含部分牙根信息的配准更有利。
   • 劣势: 模型更复杂，训练所需 GPU 显存更大，训练时间可能更长。

2. 点云配准 Baseline (/workspace/point_cloud_registration_baseline/)

   • 方法: 该方法仅使用从 STL 文件中采样得到的点云作为输入。
   • 网络结构:
     • 使用一个 PointNet 网络作为编码器，提取输入点云的全局特征。
     • 通过一个多层感知机（MLP）回归头，直接从点云特征中预测 7 个参数（4 个四元数 + 3 个平移）来表示目标变换矩阵。
   • 优势: 模型相对简单，训练速度更快，所需 GPU 显存较少。易于理解和修改。
   • 劣势: 完全忽略了 CBCT 图像信息，可能无法充分利用 CBCT 提供的上下文信息，尤其是在仅靠 STL 表面难以区分精确位置的情况下。

重要说明 (两种 Baseline 共有):

• 无验证/测试标签: 这两个 Baseline 的训练流程都已适配挑战赛场景，即不需要验证集或测试集的真值标签 (_gt.npy)。模型将在训练集上进行训练和参数优化。
• 损失函数: 两个 Baseline 目前都使用参数化的 均方误差损失 (MSE Loss) 来比较预测的 7 个参数（四元数+平移）和真值参数。
• 检查点保存: 由于缺少验证集指标，模型的检查点（Checkpoint）默认配置为保存最后一个 Epoch 的状态 (last.ckpt)。参赛者可以修改此配置（见下文）。

用户可以修改和自定义的地方

我们鼓励参赛者在这些 Baseline 的基础上进行创新和优化。以下是一些常见的修改方向：

1. 配置文件 (configs/*.yaml):

   • 数据路径 (data.data_base_dir): 必须修改此路径，指向您存放解压后数据集（包含Train-Labeled, Validation-Public, Test-Labeled 文件夹）的宿主机目录。在 docker run 时需要将此宿主机目录挂载到容器内（通常是 /dataset 或 /data）。
   • 训练参数 (trainer.*): 调整 max_epochs, accelerator (gpu/cpu), devices (GPU ID), precision (16/32/bf16)。
   • 数据加载 (data.*): 调整 batch_size, num_workers, num_points (点云采样数量), cbct_img_size (CBCT 输入尺寸，仅多模态), augment_train (是否开启 CBCT 增强，仅多模态)。
   • 模型结构 (model.*): 调整编码器输出维度 (output_dim), MLP 隐藏层大小 (mlp_hidden_dims), fusion_dim (仅多模态), dropout_rate。
   • 优化器 (optimizer.*): 调整学习率 (lr), 权重衰减 (weight_decay)。
   • 损失函数 (loss.*): 调整不同损失分量的权重 (translation_weight, rotation_weight)。
   • 预测设置 (predict.*): 修改 checkpoint_path (指定训练好的模型进行预测), output_dir (保存预测结果的位置), predict_set (指定在哪个数据集子集上进行预测)。

2. 模型架构 (src_*/models/):

   • 替换编码器: 尝试使用更先进的 3D CNN (如 VNet, UNETR, SwinUNETR from MONAI) 或点云网络 (PointNet++, DGCNN, Point Transformer)。
   • 特征融合 (仅多模态): 实现更复杂的特征融合机制，例如 Attention 机制。
   • 回归头: 尝试不同的 MLP 结构或直接回归矩阵元素（通常不稳定）。

3. 损失函数 (src_*/losses/):

   • 实现几何损失: 实现 PointDistanceLoss（需要加载原始点云，并在训练/验证步骤中应用预测/真值矩阵进行计算），这通常比参数损失效果更好。可以将其与参数损失结合。
   • 添加旋转损失: 除了比较四元数，可以显式地计算预测矩阵和真值矩阵之间的角度差并加入损失。

4. 数据增强与预处理:

   • src_*/datasets/*.py: 在 Dataset 或 DataModule 中添加更丰富的数据增强。对点云可以加入随机旋转、平移、缩放、抖动等。对 CBCT (仅多模态) 可以使用 MONAI 提供的更多 Rand* 变换。
   • normalize_points: 尝试不同的点云归一化策略。
   • CBCT 预处理 (仅多模态): 尝试不同的强度归一化（如 NormalizeIntensityd）、裁剪 (CropForegroundd) 等。

5. 训练流程 (src_*/lightning/*.py, scripts/train_*.py):

   • 学习率调度器: 替换或添加更复杂的学习率调度策略 (e.g., Cosine Annealing)。
   • 检查点保存: 修改 ModelCheckpoint 的参数，例如设置 save_top_k（如果能找到合适的监控指标，如 train_loss，但不推荐）或 every_n_epochs 来定期保存。
   • 优化器: 尝试不同的优化器 (Adam, SGD with momentum)。

如何启动训练

1. 修改配置文件:

   • 使用 nano 或 vim 编辑 configs/multi_modal_config.yaml (或 point_cloud_config.yaml)。
   • 最重要: 将 data.data_base_dir 的值修改为您在容器内挂载的数据集路径，例如 /dataset (如果您的数据集根目录，即包含Train-Labeled的目录，被挂载到了 /dataset)。
   • 根据您的硬件和需求调整其他参数（如 batch_size, devices 等）。

2. 启动训练:

   # 如果使用多模态 Baseline
   python scripts/train_mm.py
   # 或者，如果使用点云 Baseline
   # python scripts/train_pc.py
   

   训练日志和模型检查点将默认保存在 /workspace/.../training_logs_* 目录下（具体路径见配置文件 trainer.default_root_dir）。

3. 监控训练 (可选): 如果您将容器的端口映射出来，可以在宿主机浏览器中通过 TensorBoard 查看日志。首先在容器内启动 TensorBoard:

   # 在容器内，训练运行时或训练结束后
   tensorboard --logdir ./training_logs_multimodal --bind_all
   # 或者
   # tensorboard --logdir ./training_logs_pointcloud --bind_all
   

   然后在宿主机的浏览器访问 http://<容器IP或localhost>:6006 (如果端口映射正确)。

如何进行预测 (生成提交文件)

1. 训练完成: 等待训练完成，或者选择一个您认为合适的检查点 (.ckpt 文件)。

2. 修改配置文件: 编辑相应的配置文件 (multi_modal_config.yaml 或 point_cloud_config.yaml)：

   • 设置 predict.checkpoint_path 为您选择的 .ckpt 文件的完整路径 (在容器内的路径)。
   • 设置 predict.output_dir 为您希望保存预测结果的目录 (例如 /workspace/.../submission_files)。
   • 设置 predict.predict_set 为您想要进行预测的数据集部分 (Validation 或 Test)。确保 data.data_base_dir 指向包含相应 Validation-Public/Images 或 Test-Labeled/Images 的目录。

3. 运行预测脚本:

   # 如果使用多模态 Baseline
   python scripts/predict_mm.py
   # 或者，如果使用点云 Baseline
   # python scripts/predict_pc.py
   

   脚本将在指定的 predict.output_dir 下生成符合提交规范的 .npy 文件。

其他说明

• 性能: 多模态 Baseline 对 GPU 显存的要求远高于点云 Baseline。请根据您的硬件资源选择合适的 Baseline 或调整模型大小/批次大小。
• 几何损失: 目前 Baseline 使用参数损失。切换到几何损失（如 PointDistanceLoss）通常能显著提高配准精度，但这需要额外的实现工作（在 Dataset 中传递原始点云，在 Lightning Module 中加载 GT 矩阵并应用变换）。
• 可复现性: 我们已设置了随机种子 (pl.seed_everything) 和 deterministic=True，但这可能无法完全保证跨不同硬件或库版本的绝对可复现性，并且可能略微降低性能。

希望这份详细的描述能帮助参赛者更好地理解和使用这个 Baseline Docker 镜像！祝您在比赛中取得好成绩！