Esta es la versión de desarrollo del código para el siguiente artículo:
Bugra Tekin, Sudipta N. Sinha y Pascal Fua, "Predicción de pose de objeto 6D de disparo único sin interrupciones en tiempo real", CVPR 2018.
El repositorio original del código base del artículo anterior se puede encontrar en el siguiente enlace.
Proponemos un enfoque de disparo único para detectar simultáneamente un objeto en una imagen RGB y predecir su pose 6D sin requerir múltiples etapas ni tener que examinar múltiples hipótesis. El componente clave de nuestro método es una nueva arquitectura CNN inspirada en el diseño de la red YOLO que predice directamente las ubicaciones de las imágenes 2D de los vértices proyectados del cuadro delimitador 3D del objeto. Luego se estima la pose 6D del objeto utilizando un algoritmo PnP. Papel, arXiv
Si utiliza este código, cite lo siguiente
@inprocedimientos{tekin18,
TÍTULO = {{Predicción de pose de objeto 6D de disparo único en tiempo real}}, AUTOR = {Tekin, Bugra y Sinha, Sudipta N. y Fua, Pascal},
TÍTULO DEL LIBRO = {CVPR},
AÑO = {2018}
}
SingleShotPose se publica bajo la licencia MIT (consulte el archivo de LICENCIA para obtener más detalles).
El código se prueba en Windows con CUDA v8 y cudNN v5.1. La implementación se basa en PyTorch 0.4.1 y se prueba en Python3.6 . El código requiere las siguientes dependencias que podrían instalarse con conda o pip: numpy, scipy, PIL, opencv-python. Para obtener una versión anterior que sea compatible con PyTorch 0.3.1 y probada en Python2.7, consulte la carpeta py2
.
Dentro del directorio de código principal, ejecute lo siguiente para descargar y extraer (1) el conjunto de datos LINEMOD preprocesado, (2) modelos entrenados para el conjunto de datos LINEMOD, (3) el modelo entrenado para el conjunto de datos OCCLUSION, (4) imágenes de fondo de VOC2012 conjunto de datos respectivamente.
wget -O LINEMOD.tar --no-check-certificate "https://onedrive.live.com/download?cid=05750EBEE1537631&resid=5750EBEE1537631%21135&authkey=AJRHFmZbcjXxTmI"
wget -O backup.tar --no-check-certificate "https://onedrive.live.com/download?cid=0C78B7DE6C569D7B&resid=C78B7DE6C569D7B%21191&authkey=AP183o4PlczZR78"
wget -O multi_obj_pose_estimation/backup_multi.tar --no-check-certificate "https://onedrive.live.com/download?cid=05750EBEE1537631&resid=5750EBEE1537631%21136&authkey=AFQv01OSbvhGnoM"
wget https://pjreddie.com/media/files/VOCtrainval_11-May-2012.tar
wget https://pjreddie.com/media/files/darknet19_448.conv.23 -P cfg/
tar xf LINEMOD.tar
tar xf backup.tar
tar xf multi_obj_pose_estimation/backup_multi.tar -C multi_obj_pose_estimation/
tar xf VOCtrainval_11-May-2012.tar
Alternativamente, puede ir directamente a los enlaces anteriores y descargar y extraer manualmente los archivos en los directorios correspondientes. Todo el proceso de descarga puede tardar bastante (~60 minutos). También tenga en cuenta que el acceso a OneDrive en algunos países puede ser limitado.
Para entrenar la ejecución del modelo,
python train.py --datacfg [path_to_data_config_file] --modelcfg [path_to_model_config_file] --initweightfile [path_to_initialization_weights] --pretrain_num_epochs [number_of_epochs to pretrain]
p.ej
python train.py --datacfg cfg/ape.data --modelcfg cfg/yolo-pose.cfg --initweightfile cfg/darknet19_448.conv.23 --pretrain_num_epochs 15
si desea comenzar desde los pesos inicializados de ImageNet, o
python train.py --datacfg cfg/ape.data --modelcfg cfg/yolo-pose.cfg --initweightfile backup/duck/init.weights
si desea comenzar con un modelo ya entrenado previamente en LINEMOD, para una convergencia más rápida.
[datacfg] contiene información sobre las divisiones de entrenamiento/prueba, modelos de objetos 3D y parámetros de la cámara.
[modelcfg] contiene información sobre la estructura de la red
[initweightfile] contiene pesos de inicialización. <
Al comienzo de la capacitación verá un resultado como este:
layer filters size input output
0 conv 32 3 x 3 / 1 416 x 416 x 3 -> 416 x 416 x 32
1 max 2 x 2 / 2 416 x 416 x 32 -> 208 x 208 x 32
2 conv 64 3 x 3 / 1 208 x 208 x 32 -> 208 x 208 x 64
3 max 2 x 2 / 2 208 x 208 x 64 -> 104 x 104 x 64
...
30 conv 20 1 x 1 / 1 13 x 13 x1024 -> 13 x 13 x 20
31 detection
Esto define la estructura de la red. Durante el entrenamiento, el mejor modelo de red se guarda en el archivo "model.weights". Para entrenar redes para otros objetos, simplemente cambie el nombre del objeto mientras llama a la función de entrenamiento, por ejemplo, " python train.py --datacfg cfg/duck.data --modelcfg cfg/yolo-pose.cfg --initweightfile backup/duck/init.weights
". Si encuentra errores de memoria de la GPU durante el entrenamiento, puede intentar reducir el tamaño del lote, por ejemplo a 16 u 8, para que quepa en la memoria. La versión de código abierto del código ha sido objeto de una fuerte refactorización y, además, algunos modelos tuvieron que volver a entrenarse. Los modelos reentrenados que proporcionamos no cambian mucho con respecto a los resultados iniciales que brindamos (a veces un poco peores y otras veces un poco mejores).
Para probar la ejecución del modelo
python valid.py --datacfg [path_to_data_config_file] --modelcfg [path_to_model_config_file] --weightfile [path_to_trained_model_weights]
p.ej
python valid.py --datacfg cfg/ape.data --modelcfg cfg/yolo-pose.cfg --weightfile backup/ape/model_backup.weights
También puedes usar valid.ipynb para probar el modelo y visualizar los resultados.
Dentro de la carpeta multi_obj_pose_estimation/
Pruebas:
python valid_multi.py cfgfile weightfile
p.ej
python valid_multi.py cfg/yolo-pose-multi.cfg backup_multi/model_backup.weights
Capacitación:
python train_multi.py datafile cfgfile weightfile
p.ej,
python train_multi.py cfg/occlusion.data cfg/yolo-pose-multi.cfg backup_multi/init.weights
Nuestros archivos de etiquetas constan de 21 valores de verdad sobre el terreno. Predecimos 9 puntos correspondientes al centroide y las esquinas del modelo de objetos 3D. Además predecimos la clase en cada celda. Eso hace 9x2+1 = 19 puntos. En el entrenamiento de objetos múltiples, durante el entrenamiento, asignamos el cuadro de anclaje que tenga el tamaño más similar al objeto actual como el responsable de predecir las coordenadas 2D de ese objeto. Para codificar el tamaño de los objetos, tenemos 2 números adicionales para el rango en las dimensiones x e y. Por lo tanto, tenemos 9x2+1+2 = 21 números.
Respectivamente, 21 números corresponden a lo siguiente: 1er número: etiqueta de clase, 2do número: x0 (coordenada x del centroide), 3er número: y0 (coordenada y del centroide), 4to número: x1 (coordenada x de la primera esquina), quinto número: y1 (coordenada y de la primera esquina), ..., número 18: x8 (coordenada x de la octava esquina), número 19: y8 (coordenada y de la octava esquina) , Número 20: rango x, Número 21: rango y.
Las coordenadas están normalizadas por el ancho y alto de la imagen: x / image_width
e y / image_height
. Esto es útil para tener rangos de salida similares para las tareas de regresión de coordenadas y clasificación de objetos.
Entrenamos y probamos nuestros modelos en el conjunto de datos LINEMOD utilizando las mismas divisiones de entrenamiento/prueba con el método BB8 para validar nuestro enfoque. Si desea entrenar un modelo en su propio conjunto de datos, puede crear la misma estructura de carpetas con el conjunto de datos LINEMOD proporcionado y ajustar las rutas en cfg/[OBJECT].data, [DATASET]/[OBJECT]/train.txt y Archivos [CONJUNTO DE DATOS]/[OBJETO]/test.txt. La carpeta de cada objeto debe contener lo siguiente:
(1) una carpeta que contiene archivos de imágenes,
(2) una carpeta que contiene archivos de etiquetas (consulte este enlace para obtener una explicación detallada sobre cómo crear etiquetas. También puede encontrar útil la caja de herramientas ObjectDatasetTools de terceros para crear etiquetas reales para la estimación de la pose de objetos 6D),
(3) un archivo de texto que contiene los nombres de los archivos de las imágenes de entrenamiento ( train.txt
),
(4) un archivo de texto que contiene los nombres de los archivos de las imágenes de prueba ( test.txt
),
(5) un archivo .ply que contiene el modelo de objeto 3D (la unidad del modelo de objeto se proporciona en metros),
(6) opcionalmente, una carpeta que contiene máscaras de segmentación (si desea cambiar el fondo de sus imágenes de entrenamiento para que sean más resistentes a diversos fondos, esto sería esencial para una mejor capacidad de generalización),
Asegúrese también de ajustar los siguientes valores en los archivos de configuración de datos y modelos según sus necesidades:
diam
" en el archivo de configuración de datos con el diámetro del modelo de objeto en cuestión.yolo-pose.cfg
) y algunos parámetros de aumento de datos (jitter, tono, saturación, parámetros de exposición en dataset.py
). Es posible que también sea necesario ajustarlo para lograr una mejor convergencia en su conjunto de datos.yolo-pose-multi.cfg
). Asegúrese también de utilizar el número correcto de clases y especifíquelo en yolo-pose-multi.cfg
.fx
, fy
, u0
, v0
, width
, height
) en los archivos de configuración de datos por los específicos de su conjunto de datos. Al crear un conjunto de datos de entrenamiento, tomar muestras de una gran cantidad de puntos de vista/distancias y modelar una gran variabilidad de configuraciones de iluminación/oclusión/fondo sería importante para aumentar la capacidad de generalización del enfoque en su conjunto de datos. Si desea ajustar algunos parámetros del modelo y de pérdida (por ejemplo, el factor de ponderación para diferentes términos de pérdida) para sus propios fines, puede hacerlo en el archivo de configuración del modelo ( yolo-pose.cfg
).
El código está escrito por Bugra Tekin y se basa en la implementación YOLOv2 del usuario de github @marvis.
Para cualquier pregunta o informe de errores, comuníquese con Bugra Tekin