Nitrain：一款原生面向医学影像的人工智能框架

Nitrain（前身为 torchsample) 是一个框架无关的 Python 库，用于对医学影像进行采样和数据增强、在医学影像数据集上训练模型，以及在医学影像场景中可视化结果。

Nitrain 库的独特之处在于，它通过提供合理的默认设置和高层次的抽象，尽可能简化模型训练过程。此外，它还支持多种深度学习框架——PyTorch、TensorFlow 和 Keras——并且计划进一步扩展对更多框架的支持。

使用 Nitrain 训练医学影像 AI 模型的完整示例可以在 教程 页面中找到。如果您对医学影像 AI 更感兴趣，可以查看《使用 Python 的实用医学影像 AI 技术》（预计 2025 年初出版）。

快速入门

以下是一个使用 Nitrain 构建语义分割模型的示例，展示了其大部分核心功能。

import nitrain as nt
from nitrain.readers import ImageReader, ColumnReader

# 从图像文件夹和参与者文件创建数据集
dataset = nt.Dataset(inputs=ImageReader('sub-*/anat/*_T1w.nii.gz'),
                     outputs=ImageReader('sub-*/anat/*_aparc+aseg.nii.gz'),
                     transforms={
                         'inputs': tx.NormalizeIntensity(0,1),
                         ('inputs', 'outputs'): tx.Resize((64,64,64))
                     },
                     base_dir='~/desktop/ds004711/')

# 创建带有随机变换的数据加载器
loader = nt.Loader(dataset,
                   images_per_batch=4,
                   sampler=nt.SliceSampler(batch_size = 32, axis = 2)
                   transforms={
                           'inputs': tx.RandomNoise(sd=0.2)
                   })

# 根据架构创建模型
arch_fn = nt.fetch_architecture('unet', dim=2)
model = arch_fn(input_image_size=(64,64,1),
                mode='segmentation')

# 创建训练器并拟合模型
trainer = nt.Trainer(model, task='segmentation')
trainer.fit(loader, epochs=100)

安装

Nitrain 的最新版本可以从 PyPI 安装：

pip install nitrain

或者，您也可以直接从 GitHub 安装最新的开发版本：

python -m pip install git+github.com/nitrain/nitrain.git

依赖项

ants Python 包是关键依赖项，它使您能够高效地读取、处理和可视化医学影像。此外，您可以选择 Keras（tf.keras 或 keras3）、TensorFlow 或 PyTorch 作为后端来构建您的模型。

资源

以下链接可以帮助您更好地熟悉 Nitrain：

• 入门教程 [链接]
• 分割示例 [链接]
• 分类示例 [链接]
• 配准示例 [链接]

您还可以访问 ANTsPy 仓库，了解更多关于医学影像处理和传统医学影像分析的内容。

贡献

如果您有任何问题或 bug 报告，最好的求助方式是在 GitHub 页面上提交 issue。我非常欢迎任何新贡献者和项目建议。如果您想添加代码，最好的开始方式是提交一个 issue，或者直接联系我：nickcullen31@gmail.com。

您可以通过为本项目加星标或在 issues 标签页中提出功能请求来支持这项工作。这些行为有助于提升项目的影响力和社区覆盖面。

Nitrain 快速上手指南

Nitrain 是一个专为医学影像设计的 AI 框架（前身为 torchsample）。它支持 PyTorch、TensorFlow 和 Keras 等多种后端，旨在通过合理的默认设置和高层抽象，简化医学影像数据的采样、增强、模型训练及结果可视化流程。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux、macOS 或 Windows。
• Python 版本：建议 Python 3.8 及以上。
• 核心依赖：
  • ANTsPy：Nitrain 的关键依赖库，用于高效读取、操作和可视化医学影像（如 NIfTI 格式）。
  • 深度学习后端：根据需求安装其中之一：
    • PyTorch
    • TensorFlow / Keras (tf.keras 或 keras3)

提示：由于 ANTsPy 包含 C++ 扩展，在某些环境下安装可能需要编译工具（如 gcc 或 Visual Studio Build Tools）。如遇困难，建议优先使用预编译的 wheel 包或 conda 环境。

安装步骤

您可以选择从 PyPI 安装稳定版，或直接从 GitHub 安装最新开发版。

方式一：安装稳定版（推荐）

pip install nitrain

国内用户加速建议：如果下载速度较慢，可使用国内镜像源：

pip install nitrain -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：安装开发版

如需体验最新功能，可直接从 GitHub 安装：

python -m pip install git+https://github.com/nitrain/nitrain.git

注意：安装前请确保已单独安装好 antspy 和您选择的深度学习框架（如 torch 或 tensorflow），因为 nitrain 本身主要提供逻辑封装。

基本使用

以下示例展示了如何使用 Nitrain 构建一个语义分割模型。该流程涵盖了数据集加载、数据增强、模型构建及训练的核心步骤。

import nitrain as nt
from nitrain.readers import ImageReader, ColumnReader
# 假设 tx 为 transforms 模块的别名，实际使用中需根据文档导入具体变换类
# 此处沿用原文逻辑示意
import nitrain.transforms as tx 

# 1. 创建数据集
# 从文件夹读取图像和对应的标签文件，并定义预处理变换
dataset = nt.Dataset(inputs=ImageReader('sub-*/anat/*_T1w.nii.gz'),
                     outputs=ImageReader('sub-*/anat/*_aparc+aseg.nii.gz'),
                     transforms={
                         'inputs': tx.NormalizeIntensity(0,1),
                         ('inputs', 'outputs'): tx.Resize((64,64,64))
                     },
                     base_dir='~/desktop/ds004711/')

# 2. 创建数据加载器
# 配置批量大小、切片采样策略以及随机数据增强
loader = nt.Loader(dataset,
                   images_per_batch=4,
                   sampler=nt.SliceSampler(batch_size = 32, axis = 2),
                   transforms={
                           'inputs': tx.RandomNoise(sd=0.2)
                   })

# 3. 构建模型
# 获取预定义的 UNet 架构并初始化
arch_fn = nt.fetch_architecture('unet', dim=2)
model = arch_fn(input_image_size=(64,64,1),
                mode='segmentation')

# 4. 训练模型
# 创建训练器并开始拟合
trainer = nt.Trainer(model, task='segmentation')
trainer.fit(loader, epochs=100)

代码说明

• Dataset: 自动处理医学影像文件的路径匹配（支持通配符），并应用强度归一化和尺寸调整。
• Loader: 结合 SliceSampler 实现 3D 数据的切片采样，并在训练时动态添加随机噪声增强。
• Architecture: fetch_architecture 允许快速调用内置的经典网络结构（如 UNet）。
• Trainer: 封装了训练循环，只需指定任务类型即可开始训练。

更多详细案例（包括分类、配准等任务）请参考官方 Tutorials。