Entrena modelos con volúmenes, nubes de puntos y vídeo

Novedad de la v2.14.0

Este tutorial entrena tres modelos con datos 3D y en movimiento a partir de las anotaciones de un estudio de Potato. volume-unet entrena una U-Net 3D que segmenta volúmenes de TC y RM, external-3d ejecuta un detector de lidar que compilas y ejecutas tú, y video-track entrena un rastreador de objetos en vídeo. La página está dirigida a equipos que ya anotan volúmenes, nubes de puntos o vídeo en Potato y quieren un modelo que preetiquete el resto. Para cada entrenador, la página cubre el esquema que produce los datos de entrenamiento, una configuración que pasa potato validate, qué hace una ejecución, la puntuación sobre ítems reservados y cómo llegan las predicciones a los anotadores. El tutorial del primer modelo recorre la página Training con un clasificador de texto que no necesita GPU, y este tutorial parte de ese recorrido.

Resumen de los entrenadores

Los tres entrenadores ajustan un esquema por ejecución, puntúan el modelo sobre los ítems reservados val y test de la ejecución y escriben de vuelta las predicciones como preetiquetas que los anotadores editan. volume-unet y los detectores que hay detrás de external-3d necesitan una GPU en la práctica, así que las configuraciones de esos dos entrenadores envían las ejecuciones a un servidor con GPU por ssh, como configura el tutorial de entrenamiento remoto.

Todos los entrenadores leen los medios de archivos en disco. La ruta de un medio debe apuntar dentro de la carpeta de la tarea o de media_directory, y los ítems cuyo medio es una URL http(s) quedan fuera del entrenamiento y de la predicción.

Segmentar volúmenes de TC y RM con volume-unet

volume-unet aprende las máscaras que los anotadores pintan en un esquema volume_annotation y pinta máscaras en los volúmenes que nadie ha segmentado todavía. El esquema muestra un archivo NIfTI o una serie DICOM en tres vistas enlazadas (axial, coronal y sagital), y los anotadores pintan cada etiqueta con el pincel o cerrando un polígono. El campo del ítem que nombra source_field contiene un archivo .nii o .nii.gz, o un directorio con una serie DICOM, relativo a media_directory. Una caja dibujada sobre un corte no tiene máscara, así que una respuesta que solo contiene cajas no es un objetivo de entrenamiento. Potato es software de investigación sin autorización clínica ni regulatoria, y un modelo entrenado aquí propone máscaras para que los anotadores las corrijan.

La siguiente configuración segmenta el bazo y el hígado, entrena en un servidor con GPU y preetiqueta todos los volúmenes que quedan fuera de la partición train:

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

Cada línea de datos nombra un volumen:

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

El valor de patch fija el lado del parche cúbico con el que entrena la red, y debe ser múltiplo de 4 con los tres niveles predeterminados (2 elevado a levels menos 1). Un parche de 96 vóxeles da a la red más contexto alrededor de cada vóxel que el valor predeterminado de 64 y necesita más memoria de GPU. Con min_free_gpu_mb: 8000, Potato inicia una ejecución solo en una GPU con al menos 8000 MB libres. Potato toma la probabilidad media del modelo dentro de la máscara de cada etiqueta pintada y notifica la más baja de esas medias como la confianza del volumen. Con min_confidence: 0.6, un volumen en el que algún órgano se pintó con una probabilidad media inferior a 0,6 no recibe preetiqueta.

Pasos de entrenamiento de una ejecución de volume-unet

Una ejecución lee cada volumen anotado en el orden de vóxeles en que pinta el widget y convierte las máscaras en una etiqueta por vóxel. Donde dos etiquetas se solapan, gana la etiqueta posterior en el orden de etiquetas del paquete, igual que en la exportación NIfTI. Las intensidades se recortan a los percentiles 0,5 y 99,5 de cada volumen y se estandarizan. La red es una U-Net 3D en torch puro (tres niveles, normalización de instancia), entrenada con entropía cruzada más Dice suave sobre 100 parches aleatorios por época. La mitad de los parches se centran en un vóxel pintado, de modo que las estructuras pequeñas aparecen en el entrenamiento. Después de cada época, la ejecución puntúa el Dice medio sobre los volúmenes val y conserva la mejor época, y se detiene tras 10 épocas sin una puntuación mejor. La predicción desliza el parche por todo el volumen con un solapamiento de medio parche y promedia las probabilidades.

La partición train necesita al menos tres volúmenes anotados (min_instances). Un volumen cuya forma pintada difiere de la forma del archivo queda fuera con una advertencia en el registro de la ejecución, igual que un volumen cuyo archivo falta. Con varios lectores sobre un volumen, la ejecución entrena con las máscaras del lector al que resuelve aggregation, porque los agregadores comparan respuestas completas. Segmenta volúmenes de TC con varios lectores y STAPLE cubre el acuerdo entre lectores y una máscara de consenso construida fuera de Potato.

Puntuación sobre ítems reservados para volúmenes

Potato calcula el coeficiente de Dice de cada etiqueta en cada volumen reservado (dice_<label>) y notifica la media sobre todos los volúmenes y etiquetas como mean_dice. Una etiqueta que el modelo omitió en un volumen que la contiene puntúa 0. El entrenador elige la mejor época sobre val, así que la tabla de ejecuciones toma la puntuación principal de test (test_mean_dice). La lista completa de puntuaciones por etiqueta está en la página de la ejecución.

Preetiquetas de máscara en las vistas de cortes

Las máscaras predichas vuelven con la forma propia del widget, por etiqueta y corte axial, marcadas como prefilled_from: model, y aparecen en las vistas de cortes cuando un anotador abre el volumen. El anotador corrige las máscaras con el pincel y la goma, y las máscaras corregidas se guardan como la respuesta de ese anotador. Los ítems de la partición train de la ejecución nunca reciben preetiqueta. Sin include_unlabeled, una ejecución predice solo sobre los volúmenes val y test de la ejecución.

Detectar cuboides 3D con external-3d

external-3d entrega los escaneos anotados a un detector 3D que instalas en el host de entrenamiento, ejecuta tus comandos de entrenamiento y de predicción y lee las cajas predichas como cuboides. Potato no incluye un detector de nubes de puntos. Los dos detectores para los que se diseñó el entrenador, OpenPCDet y MMDetection3D, necesitan operadores CUDA compilados en el host con GPU. Por parte de Potato, el entrenador solo necesita numpy.

Los datos de entrenamiento proceden de un esquema spatial_annotation con la herramienta cuboid_3d. Los anotadores colocan cajas orientadas sobre escaneos PCD, PLY, LAS o KITTI .bin, almacenadas en metros en el sistema de referencia propio del escaneo (consulta Nubes de puntos 3D y cuboides). Los cuboides son los únicos objetivos, así que los puntos, las polilíneas y los segmentos de puntos del mismo esquema se ignoran. Cada ítem es un escaneo, y una secuencia de fotogramas con identificadores de pista no es un objetivo de entrenamiento.

Carpeta de datos escrita para el detector

Antes de que se ejecute el comando de entrenamiento, Potato escribe los escaneos en {dataset}, una carpeta llamada dataset dentro de la carpeta de la ejecución, con el formato de conjunto de datos personalizado de OpenPCDet:

• points/<id>.npy contiene un escaneo como filas float32 de x y z intensity, con intensidad 0 cuando el archivo no la tiene.
• labels/<id>.txt contiene una línea por cuboide de los escaneos train y val, como x y z dx dy dz heading label. El rumbo es la rotación alrededor de +z en radianes, y no se conservan el cabeceo ni el alabeo. Un espacio en el nombre de una etiqueta se convierte en un guion bajo.
• ImageSets/train.txt, val.txt y predict.txt listan los identificadores de los escaneos. predict.txt lista los escaneos val y test, más los escaneos sin etiquetar cuando writeback.include_unlabeled está activado.

Cada <id> son los primeros 16 caracteres hexadecimales del SHA-1 del identificador del ítem, porque los identificadores de ítem pueden contener caracteres que un nombre de archivo no admite. Las etiquetas de los escaneos test nunca se escriben en la carpeta. Potato escribe un escaneo georreferenciado relativo al origen del escaneo, porque las coordenadas georreferenciadas son demasiado grandes para float32, y desplaza las cajas predichas de vuelta con el mismo desplazamiento.

Archivos de predicción leídos del detector

El comando de predicción escribe un archivo por escaneo en {predictions}/<id>.txt, con una caja por línea, como x y z dx dy dz heading label score. La puntuación es opcional. Potato deja fuera una línea que no se puede analizar, que contiene nan o inf, que nombra una etiqueta que el esquema no tiene o que describe una caja sin tamaño, y el registro de la ejecución cuenta las líneas descartadas. Las cajas restantes vuelven como objetos cuboid_3d marcados con source: model, y la puntuación de cada caja pasa a ser la confianza de esa caja. Potato fija la confianza de un escaneo en la puntuación de caja más baja del escaneo, y min_confidence se aplica a ese valor.

Comandos de entrenamiento y de predicción

Ambos comandos se fijan como parámetros del entrenador. Potato divide cada comando en palabras como lo haría un shell, sustituye {dataset}, {workdir} (la carpeta de la ejecución) y {predictions} por las rutas correspondientes y ejecuta el comando en la carpeta de la ejecución sin shell. Otras llaves de un comando se dejan tal como están. Las tuberías, && y las $VARIABLES llegan al programa como palabras literales, de modo que el trabajo de varios pasos debe ir en un script. Añade este bloque a la configuración del estudio de nubes de puntos:

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

Las rutas están en el host de entrenamiento, donde también reside la carpeta de la ejecución. timeout_s limita cada comando, y las 12 horas de aquí sustituyen el valor predeterminado de 24. Un comando de entrenamiento que termina con un código distinto de cero hace fallar la ejecución, y el error de la ejecución muestra las últimas líneas de la salida del comando.

train.sh es tu propio script. Para OpenPCDet, el script apunta la configuración de conjunto de datos personalizado de OpenPCDet a {dataset}, construye los archivos de información con python -m pcdet.datasets.custom.custom_dataset create_custom_infos <dataset config>, ejecuta el script de entrenamiento de OpenPCDet y guarda el punto de control en {workdir}. Los nombres de clase de la configuración del detector deben coincidir con los nombres de etiqueta tal como están escritos en labels/. La guía de conjuntos de datos personalizados de MMDetection3D lee las nubes de puntos solo como archivos .bin, así que para MMDetection3D el script primero escribe cada points/<id>.npy como un archivo .bin. Deja train_command vacío para omitir el entrenamiento, y predict_command ejecuta entonces un modelo que entrenaste en otro lugar.

predict.py lee los identificadores de predict.txt, ejecuta el detector sobre cada escaneo y escribe las cajas. Las funciones load_model y detect son tu propio código alrededor del detector que entrenaste:

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

Sin train_command, el comando de predicción es toda la ejecución, y un comando de predicción que termina con un código distinto de cero hace fallar la ejecución. Un comando de predicción que termina con 0 pero no escribe ningún <id>.txt para ningún escaneo, o escribe solo líneas que Potato omite, cuenta también como fallido, porque una ejecución puntuada con esa salida se leería como un modelo que no encontró nada. Un archivo vacío es un escaneo sin objetos. Después de un comando de entrenamiento, un comando de predicción que falla no hace fallar la ejecución. La ejecución termina con una advertencia, «Prediction pass failed after 0 prediction(s)», y las últimas líneas del comando, y no tiene predicciones ni puntuación sobre ítems reservados. Consulta el registro de la ejecución en busca de esa advertencia cuando una ejecución terminada muestre «not scored».

Puntuación sobre ítems reservados para cuboides

Potato puntúa la precisión media 3D por etiqueta con un IoU volumétrico de 0,5 (ap_3d_50_<label>). El IoU es exacto para cuboides orientados con cualquier rotación. Potato notifica también la media sobre las etiquetas con IoU 0,25 y 0,5, y la precisión, la exhaustividad y la F1 con 0,5. La puntuación principal es map_3d_50. Con un train_command, la tabla de ejecuciones toma la puntuación principal primero de test, porque al comando de entrenamiento se le entregan las etiquetas val en {dataset} y puede elegir un punto de control con esas etiquetas. Sin train_command, el modelo nunca vio los ítems val del estudio, y la puntuación principal procede de val.

Preetiquetas de cuboide en el visor de nubes de puntos

Los cuboides predichos aparecen en el visor de nubes de puntos como cajas editables, y una caja editada se guarda como la respuesta del anotador. Potato comprueba cada objeto predicho antes de que el visor lo dibuje, y deja fuera un objeto cuya geometría no se puede analizar, un cuboide sin tamaño y una etiqueta que el esquema no tiene. Cómo anotar nubes de puntos 3D cubre el lado de la anotación y el acuerdo sobre cajas 3D.

Rastrear objetos en vídeo con video-track

video-track aprende de las pistas que los anotadores dibujan en un esquema video_annotation en modo tracking o combined, y predice pistas en los vídeos que nadie ha rastreado. Una pista es una etiqueta y una caja en cada fotograma clave, y el widget interpola la caja entre fotogramas clave. Un esquema video_annotation en otro modo almacena segmentos en lugar de pistas, y video-track no se ofrece para ese esquema.

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

Cada línea de datos nombra un clip bajo media_directory:

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

Fija video_fps en la tasa de fotogramas real de los clips. Los fotogramas clave se numeran por fotograma, y el entrenador numera los fotogramas clave predichos con el video_fps del esquema (30 si no se fija). El entrenador lee la caja de un fotograma muestreado según el ajuste interpolation de cada pista (lineal, Catmull-Rom cúbica o constante). Por tanto, las cajas con las que aprende el modelo siguen la trayectoria que vieron los anotadores. Como en el widget, una pista con un solo fotograma clave se mantiene en todos los fotogramas, y una pista a la que el widget dio un startFrame y un endFrame se extiende entre esos dos fotogramas. Un fotograma clave fuera de esos dos fotogramas se sigue mostrando en su propio fotograma. predict_fps: 5 detecta sobre cinco fotogramas por segundo en lugar de dos. Cinco fotogramas por segundo dan a un objeto que se mueve rápido más fotogramas clave y cuestan más pasadas del detector. Una pista termina tras más de max_age fotogramas muestreados sin una detección coincidente, así que max_age: 4 con 5 fotogramas por segundo permite que una pista sobreviva 0,8 segundos de detecciones perdidas. Un peatón que pasa detrás de un poste conserva una sola pista cuando el detector lo vuelve a encontrar en menos de 0,8 segundos. La configuración entrena en la máquina que ejecuta Potato (el backend local), y el backend gpu-server de la configuración de volúmenes también funciona con video-track.

Entrenamiento y enlace de pistas en una ejecución de video-track

Una ejecución muestrea cada vídeo anotado a frame_fps fotogramas por segundo, desde el inicio hasta el último fotograma clave, y lee cada fotograma muestreado con OpenCV. Los píxeles del lienzo dependen de la pantalla del anotador, así que el entrenador lee la caja de cada pista en ese fotograma como una fracción de la imagen. Un fotograma muestreado sin ninguna pista se mantiene en el entrenamiento como ejemplo de una escena vacía. La ejecución ajusta un Faster R-CNN de torchvision con esos fotogramas, con los parámetros y valores predeterminados de torchvision-detect. El detector parte de los pesos COCO de torchvision. Después de cada época, la ejecución puntúa el mAP de cajas con IoU 0,5 sobre los fotogramas val (frame_map_50) y conserva la mejor época.

La predicción muestrea cada vídeo a predict_fps, detecta sobre cada fotograma muestreado y enlaza las detecciones en pistas. Las detecciones de cada fotograma se emparejan una a una con las pistas vivas de la misma etiqueta, según el solapamiento con la última caja de cada pista, y un emparejamiento necesita un IoU de al menos match_iou (0,3). Una detección sin emparejar con una puntuación de al menos new_track_score (0,5) inicia una pista nueva. El entrenador descarta una pista predicha con menos de min_track_keyframes detecciones (2 por defecto) antes de escribirla y puntuarla. De lo contrario, el widget mantendría en todos los fotogramas una pista de una sola detección. Cada detección se convierte en un fotograma clave, de modo que el widget interpola entre los fotogramas clave del modelo como lo hace entre los de un anotador.

La partición train necesita al menos dos vídeos anotados (min_instances). Un vídeo que OpenCV no puede decodificar queda fuera con una advertencia que nombra el ítem, y la ejecución continúa con el resto.

Puntuación sobre ítems reservados para pistas

Potato compara las pistas predichas con las pistas anotadas en fotogramas muestreados a 2 por segundo a lo largo de cada vídeo reservado, más todos los fotogramas clave anotados. Las puntuaciones son MOTA, IDF1, cambios de identidad, omisiones y falsos positivos, y una caja se cuenta como encontrada con IoU 0,5. MOTA imputa las omisiones, los falsos positivos y los cambios de identidad al número de cajas anotadas. IDF1 es la puntuación F1 de las cajas emparejadas con la identidad correcta a lo largo de todo el vídeo, y IDF1 es la puntuación principal (idf1). El entrenador elige la mejor época sobre val, así que la tabla de ejecuciones toma la puntuación principal de test.

Preetiquetas de pista en el widget de vídeo

Las pistas predichas llenan el widget de vídeo como pistas con fotogramas clave marcados como salida de IA y con la puntuación del detector. El widget convierte las cajas a los píxeles de su propio lienzo al dibujar las pistas, y el anotador corrige las pistas como cualquier otra pista. La confianza de un vídeo es la puntuación media de los fotogramas clave del modelo. En modo combined, una preetiqueta contiene solo pistas, y el anotador añade los segmentos. Cómo anotar el seguimiento de objetos en vídeo cubre el lado de la anotación, la oclusión y el acuerdo temporal.

Iniciar una ejecución y comprobar el resultado

Reinicia el servidor después de añadir model_training, abre /admin/training, elige el esquema y el entrenador y pulsa Train now. La página ofrece solo los entrenadores que encajan con cada esquema, y lista un entrenador que no está instalado como «not installed», con el comando de instalación. Para comprobar un servidor con GPU antes de la primera ejecución, ejecuta:

potato training check gpu-server config.yaml

La comprobación notifica la versión de Python remota, cada GPU con su memoria libre, la GPU que recibiría una ejecución y los paquetes de entrenamiento que faltan en el intérprete remoto. Una ejecución también puede iniciarse desde la API, con parámetros por ejecución que sustituyen a 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}}'

Cuando la ejecución termina, la tabla de ejecuciones lista la puntuación sobre ítems reservados, con la partición y la métrica en una información emergente. Las preetiquetas condicionan a los anotadores, así que activa include_unlabeled en un estudio nuevo solo cuando la puntuación sobre ítems reservados justifique mostrar las respuestas del modelo. El tutorial del modelo en el bucle cubre cómo servir primero las preetiquetas de menor confianza y reentrenar con las correcciones, y Configuración y rutas del entrenamiento de modelos lista todas las claves, los backends y las comprobaciones de la escritura de vuelta.

Más información

• Entrenamiento de modelos con tus propias anotaciones
• Entrenadores de modelos en Potato, agrupados por tarea
• El tutorial de modelos de visión, para detección y segmentación 2D
• Seguimiento de vídeo y propagación de máscaras
• Anotación de vídeo
• Anotación de imágenes médicas
• Consenso de máscaras con STAPLE
• Referencia upstream de volume-unet, external-3d y video-track

Referencias

• 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