Skip to main content

AI 引擎扩展

AI 引擎 (packages/ai-engine) 是 RadStudio 的核心机器学习模块，基于 MONAI 和 PyRadiomics 构建。

架构​

packages/ai-engine/
├── workers/              # 训练/推理/特征提取 Worker
│   ├── train.py          # 训练 Worker
│   ├── inference.py      # 推理 Worker
│   └── features.py       # 特征提取 Worker
├── models/               # 模型定义
│   ├── segresnet.py      # MONAI SegResNet
│   ├── nnunet.py         # nnU-Net v2
│   └── custom.py         # 自定义模型
├── utils/                # 工具函数
├── requirements.txt      # 基础依赖
└── requirements-ml.txt   # ML 依赖（PyTorch 等）

添加自定义模型​

1. 定义模型​

在 packages/ai-engine/models/ 中创建模型文件：

# packages/ai-engine/models/my_model.py
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        # 定义网络结构

    def forward(self, x):
        # 前向传播
        return x

2. 注册模型​

在模型注册表中添加：

# packages/ai-engine/models/__init__.py
from .my_model import MyModel

MODEL_REGISTRY = {
    "segresnet": SegResNet,
    "nnunet": NNUNet,
    "my_model": MyModel,  # 新增
}

3. 在工作流中使用​

在前端工作流编辑器中，新模型会自动出现在模型节点列表中。

自定义特征提取​

PyRadiomics 特征扩展​

# packages/ai-engine/workers/features.py
import radiomics

# 配置特征提取器
settings = {
    'binWidth': 25,
    'resampledPixelSpacing': None,
    'interpolator': 'sitkBSpline',
    'enableCExtensions': True,
}

extractor = radiomics.featureextractor.RadiomicsFeatureExtractor(**settings)
extractor.enableFeatureClassByName('glcm', enabled=True)
extractor.enableFeatureClassByName('glrlm', enabled=True)

训练 Worker 开发​

# packages/ai-engine/workers/train.py
from monai.engines import SupervisedTrainer
from monai.handlers import StatsHandler, TensorBoardStatsHandler

def create_trainer(model, train_loader, optimizer, loss_function):
    trainer = SupervisedTrainer(
        device=device,
        max_epochs=100,
        train_data_loader=train_loader,
        network=model,
        optimizer=optimizer,
        loss_function=loss_function,
        inferer=SimpleInferer(),
        key_train_metric={"dice": DiceMetric()},
    )
    return trainer