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. 在工作流中使用
在前端工作流编辑器中,新模型会自动出现在模型节点列表中。