跳转到内容

在体数据、点云和视频上训练模型

在绘制的 CT 掩码上拟合 3D U-Net,在激光雷达长方体上运行自己的 OpenPCDet 检测器,并训练视频跟踪器,每种模型都在留出条目上评分。

v2.14.0 新增

本教程根据 Potato 研究自己的标注,在 3D 数据和运动数据上训练三个模型。volume-unet 训练分割 CT 和 MR 体数据的 3D U-Net,external-3d 运行由你自己构建和运行的激光雷达检测器,video-track 训练视频目标跟踪器。本页面向已经在 Potato 中标注体数据、点云或视频,并希望用模型为其余条目预标注的团队。对每个训练器,本页介绍产生训练数据的方案、能通过 potato validate 的配置、一次运行做了什么、留出得分,以及预测如何送到标注者面前。训练你的第一个模型教程用一个不需要 GPU 的文本分类器演示了训练页面,本教程以该教程为前提。

训练器概览

三个训练器每次运行都只拟合一个方案,在该次运行留出的 val 和 test 条目上给模型评分,并把预测作为供标注者编辑的预标签回写。volume-unet 和 external-3d 背后的检测器在实践中需要 GPU,所以这两个训练器的配置通过 ssh 把运行发送到 GPU 服务器,设置方法见远程训练教程。

训练器方案条目在训练机器上安装主得分
volume-unetvolume_annotation(绘制的掩码)NIfTI 体数据、DICOM 序列pip install "potato-annotation[train-vision,volumes]";实践中需要 GPUmean_dice
external-3dspatial_annotation(长方体)点云你自己的检测器(OpenPCDet、MMDetection3D)map_3d_50
video-tracktracking 或 combined 模式的 video_annotation视频pip install "potato-annotation[train-vision]" opencv-pythonidf1

每个训练器都从磁盘上的文件读取媒体。媒体路径必须指向任务目录或 media_directory 之内,媒体为 http(s) URL 的条目不参与训练和预测。

用 volume-unet 分割 CT 和 MR 体数据

volume-unet 学习标注者在 volume_annotation 方案中绘制的掩码,并在还没有人分割的体数据上绘制掩码。该方案在三个联动的视图(轴状面、冠状面和矢状面)中显示 NIfTI 文件或 DICOM 序列,标注者用画笔绘制每个标签,或者闭合一个多边形。source_field 指定的条目字段保存一个 .nii 或 .nii.gz 文件,或者保存包含一个 DICOM 序列的目录,路径相对于 media_directory。在切片上画的框没有掩码,所以只含框的答案不是训练目标。Potato 是研究软件,没有临床或监管许可,在这里训练的模型提出的掩码需要由标注者修正。

下面的配置分割脾脏和肝脏,在 GPU 服务器上训练,并为 train 划分之外的每个体数据预标注:

yaml
annotation_task_name: "Spleen segmentation"
task_dir: .
media_directory: media
data_files:
  - data/scans.jsonl
item_properties:
  id_key: id
  text_key: volume
output_annotation_dir: annotation_output
user_config:
  allow_all_users: true
 
annotation_schemes:
  - annotation_type: volume_annotation
    name: organs
    description: "Paint the spleen and the liver on every slice where each organ appears."
    source_field: volume
    labels:
      - {name: Spleen, color: "#e6553a", key_value: "1"}
      - {name: Liver, color: "#3a8ee6", key_value: "2"}
    tools: [brush, eraser, polygon]
    default_window: Soft tissue
 
model_training:
  enabled: true
  backend: gpu-server
  backends:
    gpu-server:
      type: ssh
      host: gpu.example.edu
      python: /opt/conda/bin/python3
      gpus: auto
      min_free_gpu_mb: 8000
  trainer_params:
    volume-unet:
      patch: 96
  writeback:
    schemas: [organs]
    min_confidence: 0.6
    include_unlabeled: true

每行数据对应一个体数据:

json
{"id": "case_001", "volume": "case_001.nii.gz"}
{"id": "case_002", "volume": "case_002_ct"}

patch 的值设定网络训练所用立方体块的边长,在默认的三个层级下必须是 4 的倍数(2 的 levels 减 1 次方)。96 体素的块让网络在每个体素周围看到比默认 64 更多的上下文,也需要更多的 GPU 内存。设置 min_free_gpu_mb: 8000 后,Potato 只在空闲内存至少 8,000 MB 的 GPU 上启动运行。Potato 取模型在每个已绘制标签的掩码内的平均概率,并把这些平均值中最低的一个作为该体数据的置信度。设置 min_confidence: 0.6 后,只要有任何一个器官的绘制平均概率低于 0.6,该体数据就不会得到预标签。

volume-unet 运行中的训练步骤

一次运行按控件绘制时的体素顺序读取每个已标注的体数据,并把掩码变成每个体素一个标签。两个标签重叠的地方,以包中标签顺序靠后的标签为准,与 NIfTI 导出一致。强度被裁剪到每个体数据的 0.5 和 99.5 百分位并做标准化。网络是纯 torch 实现的 3D U-Net(三个层级,instance norm),用交叉熵加 soft Dice 训练,每个 epoch 取 100 个随机块。其中一半的块以已绘制的体素为中心,所以小结构会出现在训练中。每个 epoch 之后,运行在 val 体数据上计算平均 Dice 并保留最佳 epoch,连续 10 个 epoch 没有更好的得分时运行停止。预测时让块以半块重叠的方式滑过整个体数据,并对概率取平均。

train 划分至少需要三个已标注的体数据(min_instances)。已绘制形状与文件形状不同的体数据会被排除,并在运行日志中给出警告,文件缺失的体数据也一样。同一个体数据有多位读片者时,运行用 aggregation 解析出的那一位读片者的掩码训练,因为聚合器比较的是整个答案。Segment CT Volumes With Several Readers and STAPLE 介绍读片者之间的一致性,以及在 Potato 之外构建的共识掩码。

体数据的留出得分

Potato 对每个留出体数据上的每个标签计算 Dice 系数(dice_<label>),并把所有体数据和标签上的平均值报告为 mean_dice。模型在含有某标签的体数据中漏掉该标签时,得分为 0。训练器在 val 上选出最佳 epoch,所以 Runs 表的主得分取自 test(test_mean_dice)。各标签得分的完整列表在运行页面上。

切片视图上的掩码预标签

预测的掩码以控件自己的形状(每个标签、每个轴状面切片)返回,标记为 prefilled_from: model,标注者打开该体数据时会显示在切片视图上。标注者用画笔和橡皮擦修正掩码,修正后的掩码保存为该标注者的答案。属于该次运行 train 划分的条目永远不会得到预标签。没有 include_unlabeled 时,一次运行只在该次运行的 val 和 test 体数据上预测。

用 external-3d 检测 3D 长方体

external-3d 把已标注的扫描交给你安装在训练主机上的 3D 检测器,运行你的训练和预测命令,再把预测的框作为长方体读回。Potato 不捆绑点云检测器。该训练器针对的两个检测器 OpenPCDet 和 MMDetection3D 需要在 GPU 主机上编译 CUDA 算子。在 Potato 一侧,该训练器只需要 numpy。

训练数据来自带 cuboid_3d 工具的 spatial_annotation 方案。标注者在 PCD、PLY、LAS 或 KITTI .bin 扫描上放置有朝向的框,以米为单位保存在扫描自己的坐标系中(见3D 点云和长方体)。长方体是唯一的目标,所以同一方案中的点、折线和点段会被忽略。每个条目是一次扫描,带轨迹 id 的帧序列不是训练目标。

为检测器写出的数据集文件夹

训练命令运行之前,Potato 把扫描写入 {dataset},这是运行文件夹内名为 dataset 的文件夹,采用 OpenPCDet 的自定义数据集布局:

  • points/<id>.npy 把一次扫描保存为 x y z intensity 的 float32 行,文件没有强度时强度为 0。
  • labels/<id>.txt 为 train 和 val 扫描的每个长方体写一行,格式为 x y z dx dy dz heading label。heading 是绕 +z 的旋转,单位为弧度,俯仰和横滚不保留。标签名中的空格变成下划线。
  • ImageSets/train.txt、val.txt 和 predict.txt 列出扫描 id。predict.txt 列出 val 和 test 扫描,在 writeback.include_unlabeled 开启时还包括未标注的扫描。

每个 <id> 是条目 id 的 SHA-1 的前 16 个十六进制字符,因为条目 id 可能含有文件名不能使用的字符。test 扫描的标签永远不会写入该文件夹。Potato 以扫描的原点为基准写出带地理参考的扫描,因为地理参考坐标对 float32 来说太大,并把预测的框按同样的偏移量移回。

从检测器读回的预测文件

预测命令为每次扫描向 {predictions}/<id>.txt 写一个文件,每行一个框,格式为 x y z dx dy dz heading label score。score 是可选的。Potato 会排除无法解析、含有 nan 或 inf、使用了方案中没有的标签,或者描述的框没有大小的行,运行日志会统计被排除的行数。其余的框作为标记为 source: model 的 cuboid_3d 对象返回,每个框的 score 成为该框的置信度。Potato 把一次扫描的置信度设为该扫描中最低的框得分,min_confidence 作用于这个值。

训练和预测命令

两个命令都设为训练器参数。Potato 按 shell 的方式把每个命令拆成词,把 {dataset}、{workdir}(运行的文件夹)和 {predictions} 替换为对应路径,并在运行的文件夹中不经 shell 运行该命令。命令中的其他花括号保持原样。管道、&& 和 $VARIABLES 会作为字面的词传给程序,所以多步工作应放进脚本。把下面的块添加到点云研究的配置中:

yaml
model_training:
  enabled: true
  backend: gpu-server
  backends:
    gpu-server:
      type: ssh
      host: gpu.example.edu
      python: /opt/conda/envs/pcdet/bin/python
      gpus: auto
  trainer_params:
    external-3d:
      train_command: "/opt/pcdet-jobs/train.sh {dataset} {workdir}"
      predict_command: "/opt/conda/envs/pcdet/bin/python /opt/pcdet-jobs/predict.py {dataset} {workdir} {predictions}"
      timeout_s: 43200
  writeback:
    schemas: [objects]
    min_confidence: 0.3
    include_unlabeled: true

这些路径位于训练主机上,运行的文件夹也在那里。timeout_s 限制每个命令的时长,这里的 12 小时取代了默认的 24 小时。训练命令以非零状态退出会使运行失败,运行的错误信息会显示该命令输出的最后几行。

train.sh 是你自己的脚本。对于 OpenPCDet,脚本把 OpenPCDet 的自定义数据集配置指向 {dataset},用 python -m pcdet.datasets.custom.custom_dataset create_custom_infos <dataset config> 构建信息文件,运行 OpenPCDet 的训练脚本,并把检查点保存在 {workdir} 下。检测器配置中的类别名必须与 labels/ 中写出的标签名一致。MMDetection3D 的自定义数据集指南只读取 .bin 文件形式的点云,所以对 MMDetection3D,脚本要先把每个 points/<id>.npy 写成 .bin 文件。train_command 留空则跳过训练,此时 predict_command 运行你在别处训练的模型。

predict.py 读取 predict.txt 中的 id,对每次扫描运行检测器,并写出框。load_model 和 detect 函数是你围绕所训练的检测器写的代码:

python
"""predict.py DATASET WORKDIR PREDICTIONS: one line per box, one file per scan."""
import os
import sys
 
import numpy as np
 
from my_detector import load_model, detect  # your code around OpenPCDet or MMDetection3D
 
dataset, workdir, out = sys.argv[1:4]
model = load_model(os.path.join(workdir, "output"))  # where train.sh saved the checkpoint
os.makedirs(out, exist_ok=True)
 
with open(os.path.join(dataset, "ImageSets", "predict.txt")) as fh:
    scan_ids = fh.read().split()
 
for scan_id in scan_ids:
    points = np.load(os.path.join(dataset, "points", scan_id + ".npy"))  # float32, (N, 4)
    with open(os.path.join(out, scan_id + ".txt"), "w") as fh:
        for x, y, z, dx, dy, dz, heading, label, score in detect(model, points):
            fh.write(f"{x} {y} {z} {dx} {dy} {dz} {heading} {label} {score}\n")

没有 train_command 时,预测命令就是整个运行,预测命令以非零状态退出会使运行失败。预测命令即使以 0 退出,但没有为任何扫描写出 <id>.txt,或只写出 Potato 会舍弃的行,也算失败,因为按这种输出评分的运行会像是模型什么也没找到。空文件表示该扫描没有物体。在训练命令之后,失败的预测命令不会使运行失败。运行会带着警告“Prediction pass failed after 0 prediction(s)”和该命令的最后几行结束,此时运行没有预测,也没有留出得分。已结束的运行显示“not scored”时,请在运行日志中查找这条警告。

长方体的留出得分

Potato 在体积 IoU 为 0.5 时为每个标签计算 3D 平均精度(ap_3d_50_<label>)。对任意旋转的有朝向长方体,该 IoU 是精确的。Potato 还报告各标签在 IoU 0.25 和 0.5 时的平均值,以及 0.5 时的精确率、召回率和 F1。主得分是 map_3d_50。有 train_command 时,Runs 表的主得分优先取自 test,因为训练命令在 {dataset} 中拿到了 val 标签,可能据此选择检查点。没有 train_command 时,模型从未见过研究的 val 条目,主得分取自 val。

点云查看器中的长方体预标签

预测的长方体在点云查看器中显示为可编辑的框,编辑后的框保存为标注者的答案。Potato 在查看器绘制每个预测对象之前先检查它,并排除几何无法解析的对象、没有大小的长方体,以及方案中没有的标签。How to Annotate 3D Point Clouds 介绍标注一侧,以及 3D 框上的一致性。

用 video-track 跟踪视频中的目标

video-track 从标注者在 tracking 或 combined 模式的 video_annotation 方案中画的轨迹学习,并在还没有人跟踪的视频上预测轨迹。一条轨迹由一个标签和每个关键帧上的一个框组成,控件在关键帧之间对框做插值。其他模式的 video_annotation 方案保存的是片段而不是轨迹,该方案不提供 video-track。

yaml
annotation_task_name: "Pedestrian tracking"
task_dir: .
media_directory: media
data_files:
  - data/clips.jsonl
item_properties:
  id_key: id
  text_key: video
output_annotation_dir: annotation_output
user_config:
  allow_all_users: true
 
annotation_schemes:
  - annotation_type: video_annotation
    name: tracks
    description: "Track every pedestrian and cyclist with a box at each keyframe."
    source_field: video
    mode: tracking
    video_fps: 25
    labels:
      - {name: pedestrian, color: "#FF6B6B"}
      - {name: cyclist, color: "#45B7D1"}
    tracking_options:
      interpolation: linear
      auto_advance_frames: 5
    frame_stepping: true
 
model_training:
  enabled: true
  trainer_params:
    video-track:
      frame_fps: 2
      predict_fps: 5
      max_age: 4
  writeback:
    schemas: [tracks]
    include_unlabeled: true

每行数据对应 media_directory 下的一个视频片段:

json
{"id": "clip_001", "video": "/media/clip_001.mp4"}
{"id": "clip_002", "video": "/media/clip_002.mp4"}

请把 video_fps 设为片段的真实帧率。关键帧按帧编号,训练器按方案的 video_fps(未设置时为 30)为预测的关键帧编号。训练器按每条轨迹自己的 interpolation 设置(线性、三次 Catmull-Rom 或常数)读取采样帧上的框,所以模型学习的框与标注者看到的路径一致。与控件中一样,只有一个关键帧的轨迹在每一帧都保持不动,控件为其指定了 startFrame 和 endFrame 的轨迹则在这两帧之间运行。位于这两帧之外的关键帧仍会显示在它自己的那一帧上。predict_fps: 5 表示每秒在五帧上检测,而不是两帧。每秒五帧让快速运动的目标得到更多关键帧,代价是更多的检测器遍数。一条轨迹在连续超过 max_age 个采样帧没有匹配的检测之后结束,所以在每秒 5 帧下 max_age: 4 让轨迹能挺过 0.8 秒的漏检。行人从杆子后面经过时,只要检测器在 0.8 秒内重新发现该行人,就仍是同一条轨迹。该配置在运行 Potato 的机器上训练(local 后端),体数据配置中的 gpu-server 后端同样适用于 video-track。

video-track 运行中的训练和轨迹连接

一次运行以每秒 frame_fps 帧的速率对每段已标注的视频采样,从开头到最后一个关键帧,并用 OpenCV 读取每个采样帧。画布像素取决于标注者的屏幕,所以训练器把每条轨迹在该帧上的框读取为画面的一个比例。采样帧上没有轨迹时,该帧仍留在训练中,作为空场景的示例。运行在这些帧上微调 torchvision 的 Faster R-CNN,参数和默认值与 torchvision-detect 相同。检测器从 torchvision 的 COCO 权重开始。每个 epoch 之后,运行在 val 帧上计算 IoU 为 0.5 时的框 mAP(frame_map_50)并保留最佳 epoch。

预测以 predict_fps 对每段视频采样,在每个采样帧上检测,并把检测连接成轨迹。每一帧的检测与同一标签的活动轨迹一对一匹配,依据是与每条轨迹最后一个框的重叠,匹配要求 IoU 至少为 match_iou(0.3)。得分至少为 new_track_score(0.5)的未匹配检测开始一条新轨迹。训练器在写回和评分之前,会丢弃检测数少于 min_track_keyframes(默认 2)的预测轨迹。否则控件会让只有一个检测的轨迹在每一帧都保持不动。每个检测都成为一个关键帧,所以控件在模型的关键帧之间做插值,与在标注者的关键帧之间一样。

train 划分至少需要两段已标注的视频(min_instances)。OpenCV 无法解码的视频会被排除,并给出点名该条目的警告,运行用其余的继续。

轨迹的留出得分

Potato 在每段留出视频上每秒采样 2 帧,再加上每个已标注的关键帧,把预测的轨迹与已标注的轨迹比较。得分包括 MOTA、IDF1、ID 切换、漏检和误检,IoU 达到 0.5 的框算作找到。MOTA 把漏检、误检和身份切换按已标注框的数量计入代价。IDF1 是在整段视频中以正确身份匹配的框的 F1 得分,IDF1 是主得分(idf1)。训练器在 val 上选出最佳 epoch,所以 Runs 表的主得分取自 test。

视频控件中的轨迹预标签

预测的轨迹以轨迹的形式填入视频控件,关键帧标记为 AI 输出并带有检测器的得分。控件在绘制轨迹时把框转换为控件自己的画布像素,标注者像修正其他轨迹一样修正这些轨迹。一段视频的置信度是模型关键帧得分的平均值。在 combined 模式下,预标签只含轨迹,片段由标注者添加。How to Annotate Video Object Tracking 介绍标注一侧、遮挡和时间上的一致性。

开始一次运行并检查结果

添加 model_training 之后重启服务器,打开 /admin/training,选择方案和训练器,然后点击 Train now。该页面只提供适合每个方案的训练器,并把未安装的训练器列为“not installed”,附上安装命令。要在第一次运行之前检查 GPU 服务器,请运行:

bash
potato training check gpu-server config.yaml

该检查会报告远程 Python 版本、每块 GPU 及其空闲内存、一次运行将分到的 GPU,以及远程解释器缺少的训练包。也可以通过 API 启动运行,每次运行的参数会覆盖 trainer_params:

bash
curl -X POST -H "X-API-Key: $KEY" -H "Content-Type: application/json" \
  http://localhost:8000/admin/training/api/runs \
  -d '{"trainer": "video-track", "schemas": ["tracks"], "params": {"predict_fps": 10}}'

运行结束后,Runs 表列出留出得分,划分和指标显示在提示中。预标签会锚定标注者,所以请等留出得分足以支持展示模型的答案之后,再在新研究中开启 include_unlabeled。模型在环教程介绍优先提供置信度最低的预标签,以及在修正上重新训练,模型训练配置和路由列出每个键、后端和回写检查。

延伸阅读

参考文献

  • Bernardin, K., and Stiefelhagen, R. (2008). Evaluating Multiple Object Tracking Performance: The CLEAR MOT Metrics. EURASIP Journal on Image and Video Processing. https://doi.org/10.1155/2008/246309
  • Çiçek, Ö., Abdulkadir, A., Lienkamp, S. S., Brox, T., and Ronneberger, O. (2016). 3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation. MICCAI 2016. https://arxiv.org/abs/1606.06650
  • Ristani, E., Solera, F., Zou, R. S., Cucchiara, R., and Tomasi, C. (2016). Performance Measures and a Data Set for Multi-Target, Multi-Camera Tracking. ECCV 2016 Workshop on Benchmarking Multi-Target Tracking. https://arxiv.org/abs/1609.01775