Train Models on Volumes, Point Clouds, and Video

New in v2.14.0

This tutorial trains three models on 3D and moving data from a Potato study's own annotations. volume-unet trains a 3D U-Net that segments CT and MR volumes, external-3d runs a lidar detector that you build and run yourself, and video-track trains a video object tracker. The page is for teams who already annotate volumes, point clouds, or video in Potato and want a model to pre-label the rest. For each trainer the page covers the scheme that produces the training data, a config that passes potato validate, what a run does, the held-out score, and how the predictions reach annotators. The first-model tutorial walks through the Training page with a text classifier that needs no GPU, and this tutorial assumes that walkthrough.

Trainer summary

All three trainers fit one scheme per run, score the model on the run's held-out val and test items, and write predictions back as pre-labels that annotators edit. volume-unet and the detectors behind external-3d need a GPU in practice, so the configs for those two trainers send runs to a GPU server over ssh, as the remote training tutorial sets up.

Every trainer reads media from files on disk. A media path must point inside the task directory or media_directory, and items whose media is an http(s) URL are left out of training and prediction.

Segment CT and MR volumes with volume-unet

volume-unet learns the masks annotators paint in a volume_annotation scheme and paints masks on the volumes nobody has segmented yet. The scheme shows a NIfTI file or a DICOM series in three linked views (axial, coronal, and sagittal), and annotators paint each label with the brush or by closing a polygon. The item field that source_field names holds a .nii or .nii.gz file, or a directory holding one DICOM series, relative to media_directory. A box drawn on a slice has no mask, so an answer with boxes alone is not a training target. Potato is research software with no clinical or regulatory clearance, and a model trained here proposes masks for annotators to correct.

The config below segments the spleen and the liver, trains on a GPU server, and pre-labels every volume outside the train split:

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

Each data line names one volume:

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

The patch value sets the side of the cubic patch the network trains on, and must be a multiple of 4 at the default three levels (2 to the power of levels minus 1). A 96-voxel patch gives the network more context around each voxel than the default 64 and needs more GPU memory. With min_free_gpu_mb: 8000, Potato starts a run only on a GPU with at least 8,000 MB free. Potato takes the model's mean probability inside each painted label's mask and reports the lowest of those means as the volume's confidence. With min_confidence: 0.6, a volume where any one organ was painted at a mean probability under 0.6 gets no pre-label.

Training steps in a volume-unet run

A run reads each annotated volume in the voxel order the widget paints in and turns the masks into one label per voxel. Where two labels overlap, the later label in the bundle's label order wins, as the NIfTI export does. Intensities are clipped to each volume's 0.5 and 99.5 percentiles and standardized. The network is a 3D U-Net in plain torch (three levels, instance norm), trained with cross-entropy plus soft Dice on 100 random patches per epoch. Half the patches are centered on a painted voxel, so small structures appear in training. After each epoch the run scores mean Dice on the val volumes and keeps the best epoch, and the run stops after 10 epochs without a better score. Prediction slides the patch over the whole volume with half-patch overlap and averages the probabilities.

The train split needs at least three annotated volumes (min_instances). A volume whose painted shape differs from the file's shape is left out with a warning in the run log, and so is a volume whose file is missing. With several readers on one volume, the run trains on the one reader's masks that aggregation resolves to, because the aggregators compare whole answers. Segment CT Volumes With Several Readers and STAPLE covers reader agreement and a consensus mask built outside Potato.

Held-out score for volumes

Potato scores the Dice coefficient for each label on each held-out volume (dice_<label>) and reports the mean over every volume and label as mean_dice. A label the model left out of a volume that has the label scores 0. The trainer picks the best epoch on val, so the Runs table takes the headline from test (test_mean_dice). The full list of per-label scores is on the run page.

Mask pre-labels on the slice views

Predicted masks come back in the widget's own shape, per label and axial slice, marked prefilled_from: model, and appear on the slice views when an annotator opens the volume. The annotator corrects the masks with the brush and the eraser, and the corrected masks are saved as that annotator's answer. Items in the run's train split never get a pre-label. Without include_unlabeled, a run predicts only on the run's val and test volumes.

Detect 3D cuboids with external-3d

external-3d hands the annotated scans to a 3D detector that you install on the training host, runs your train and predict commands, and reads the predicted boxes back as cuboids. Potato does not bundle a point-cloud detector. The two detectors the trainer is built for, OpenPCDet and MMDetection3D, need CUDA operators compiled on the GPU host. On Potato's side the trainer needs only numpy.

The training data comes from a spatial_annotation scheme with the cuboid_3d tool. Annotators place oriented boxes on PCD, PLY, LAS, or KITTI .bin scans, stored in meters in the scan's own frame (see 3D Point Clouds and Cuboids). Cuboids are the only targets, so points, polylines, and point segments in the same scheme are ignored. Each item is one scan, and a sequence of frames with track ids is not a training target.

Dataset folder written for the detector

Before the train command runs, Potato writes the scans into {dataset}, a folder named dataset inside the run's folder, in OpenPCDet's custom dataset layout:

• points/<id>.npy holds one scan as float32 rows of x y z intensity, with intensity 0 when the file has none.
• labels/<id>.txt holds one line per cuboid for the train and val scans, as x y z dx dy dz heading label. The heading is the rotation about +z in radians, and pitch and roll are not kept. A space in a label name becomes an underscore.
• ImageSets/train.txt, val.txt, and predict.txt list the scan ids. predict.txt lists the val and test scans, plus the unlabeled scans when writeback.include_unlabeled is on.

Each <id> is the first 16 hex characters of the SHA-1 of the item id, because item ids may hold characters a file name cannot. The labels of test scans are never written into the folder. Potato writes a georeferenced scan relative to the scan's origin, because georeferenced coordinates are too large for float32, and shifts the predicted boxes back by the same offset.

Prediction files read back from the detector

The predict command writes one file per scan to {predictions}/<id>.txt, one box per line, as x y z dx dy dz heading label score. The score is optional. Potato leaves out a line that does not parse, holds nan or inf, names a label the scheme does not have, or describes a box with no size, and the run log counts the lines left out. The remaining boxes come back as cuboid_3d objects marked source: model, and each box's score becomes the box's confidence. Potato sets a scan's confidence to the scan's lowest box score, and min_confidence applies to that value.

Train and predict commands

Both commands are set as trainer parameters. Potato splits each command into words the way a shell would, replaces {dataset}, {workdir} (the run's folder), and {predictions} with the matching paths, and runs the command in the run's folder without a shell. Other braces in a command are left as written. Pipes, &&, and $VARIABLES reach the program as literal words, so multi-step work belongs in a script. Add this block to the point-cloud study's config:

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

The paths are on the training host, where the run's folder also lives. timeout_s caps each command, and 12 hours here replaces the default of 24. A train command that exits non-zero fails the run, and the run's error shows the last lines of the command's output.

train.sh is your own script. For OpenPCDet, the script points OpenPCDet's custom dataset config at {dataset}, builds the info files with python -m pcdet.datasets.custom.custom_dataset create_custom_infos <dataset config>, runs OpenPCDet's training script, and saves the checkpoint under {workdir}. The class names in the detector's config must match the label names as written in labels/. MMDetection3D's custom dataset guide reads point clouds only as .bin files, so for MMDetection3D the script first writes each points/<id>.npy out as a .bin file. Leave train_command empty to skip training, and predict_command then runs a model you trained elsewhere.

predict.py reads the ids in predict.txt, runs the detector on each scan, and writes the boxes. The load_model and detect functions are your own code around the detector you trained:

"""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")

Without a train_command, the predict command is the whole run, and a predict command that exits non-zero fails the run. A predict command that exits 0 but writes no <id>.txt for any scan, or writes only lines Potato leaves out, counts as failing too, because a run scored on that output would read as a model that found nothing. An empty file is a scan with no objects. After a train command, a failing predict command does not fail the run. The run finishes with a warning, "Prediction pass failed after 0 prediction(s)", and the command's last lines, and the run has no predictions and no held-out score. Check the run log for that warning when a finished run shows "not scored".

Held-out score for cuboids

Potato scores 3D average precision per label at a volumetric IoU of 0.5 (ap_3d_50_<label>). The IoU is exact for oriented cuboids at any rotation. Potato also reports the mean over labels at IoU 0.25 and 0.5, and precision, recall, and F1 at 0.5. The headline is map_3d_50. With a train_command, the Runs table takes the headline from test first, because the train command is handed the val labels in {dataset} and may choose a checkpoint on those labels. Without a train_command, the model never saw the study's val items, and the headline comes from val.

Cuboid pre-labels in the point-cloud viewer

Predicted cuboids appear in the point-cloud viewer as editable boxes, and an edited box is saved as the annotator's answer. Potato checks each predicted object before the viewer draws the object, and leaves out an object whose geometry does not parse, a cuboid with no size, and a label the scheme does not have. How to Annotate 3D Point Clouds covers the annotation side and agreement on 3D boxes.

Track objects in video with video-track

video-track learns from the tracks annotators draw in a video_annotation scheme in tracking or combined mode, and predicts tracks on the videos nobody has tracked. A track is a label and a box at each keyframe, and the widget interpolates the box between keyframes. A video_annotation scheme in another mode stores segments rather than tracks, and video-track is not offered for that scheme.

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

Each data line names a clip under media_directory:

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

Set video_fps to the clips' real frame rate. Keyframes are numbered by frame, and the trainer numbers the predicted keyframes at the scheme's video_fps (30 when unset). The trainer reads the box on a sampled frame by each track's own interpolation setting (linear, cubic Catmull-Rom, or constant). The boxes the model learns from therefore follow the path the annotators saw. As in the widget, a track with one keyframe is held at every frame, and a track the widget gave a startFrame and an endFrame runs between those two frames. A keyframe outside those two frames still shows on the keyframe's own frame. predict_fps: 5 detects on five frames per second instead of two. Five frames per second give a fast-moving object more keyframes and cost more detector passes. A track ends after more than max_age sampled frames without a matching detection, so max_age: 4 at 5 frames per second lets a track survive 0.8 seconds of missed detections. A pedestrian who passes behind a pole keeps one track when the detector finds the pedestrian again within 0.8 seconds. The config trains on the machine running Potato (the local backend), and the gpu-server backend from the volume config works for video-track too.

Training and track linking in a video-track run

A run samples each annotated video at frame_fps frames per second, from the start to the last keyframe, and reads each sampled frame with OpenCV. Canvas pixels depend on the annotator's screen, so the trainer reads each track's box on that frame as a fraction of the picture. A sampled frame with no track on the frame stays in training as an example of an empty scene. The run fine-tunes a torchvision Faster R-CNN on those frames, with the parameters and defaults of torchvision-detect. The detector starts from torchvision's COCO weights. After each epoch the run scores box mAP at IoU 0.5 on the val frames (frame_map_50) and keeps the best epoch.

Prediction samples each video at predict_fps, detects on every sampled frame, and links the detections into tracks. Each frame's detections are matched one to one with the live tracks of the same label, by the overlap with each track's last box, and a match needs an IoU of at least match_iou (0.3). An unmatched detection with a score of at least new_track_score (0.5) starts a new track. The trainer drops a predicted track with fewer than min_track_keyframes detections (2 by default) before write-back and scoring. The widget would otherwise hold a one-detection track on every frame. Every detection becomes a keyframe, so the widget interpolates between the model's keyframes as the widget does between an annotator's.

The train split needs at least two annotated videos (min_instances). A video that OpenCV cannot decode is left out with a warning naming the item, and the run continues with the rest.

Held-out score for tracks

Potato compares the predicted tracks with the annotated tracks on frames sampled at 2 per second across each held-out video, plus every annotated keyframe. The scores are MOTA, IDF1, ID switches, misses, and false positives, with a box counted as found at IoU 0.5. MOTA charges misses, false positives, and identity switches against the number of annotated boxes. IDF1 is the F1 score of boxes matched to the correct identity across the whole video, and IDF1 is the headline (idf1). The trainer picks the best epoch on val, so the Runs table takes the headline from test.

Track pre-labels in the video widget

The predicted tracks fill the video widget as tracks with keyframes marked as AI output and carrying the detector's score. The widget converts the boxes to the widget's own canvas pixels when drawing the tracks, and the annotator corrects the tracks like any other track. A video's confidence is the mean score of the model's keyframes. In combined mode a pre-label holds tracks only, and the annotator adds the segments. How to Annotate Video Object Tracking covers the annotation side, occlusion, and temporal agreement.

Start a run and check the result

Restart the server after adding model_training, open /admin/training, pick the scheme and the trainer, and click Train now. The page offers only the trainers that fit each scheme, and lists a trainer that is not installed as "not installed", with the install command. To check a GPU server before the first run, run:

potato training check gpu-server config.yaml

The check reports the remote Python version, each GPU with the GPU's free memory, the GPU a run would get, and the training packages the remote interpreter lacks. A run can also start from the API, with per-run parameters that override trainer_params:

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}}'

When the run finishes, the Runs table lists the held-out score, with the split and the metric in a tooltip. Pre-labels anchor annotators, so turn on include_unlabeled in a new study once the held-out score supports showing the model's answers. The model-in-the-loop tutorial covers serving the least-confident pre-labels first and retraining on corrections, and Model Training Configuration and Routes lists every key, the backends, and the write-back checks.

Further reading

• Model Training on Your Own Annotations
• Model Trainers in Potato, Grouped by Task
• The vision-models tutorial, for 2D detection and segmentation
• Video Tracking and Mask Propagation
• Video annotation
• Medical imaging annotation
• Mask consensus with STAPLE
• Upstream reference for volume-unet, external-3d, and video-track

References

• 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