SOLQ: segmentación de objetos mediante consultas de aprendizaje

Este repositorio es una implementación oficial del documento SOLQ de NeurIPS 2021: segmentación de objetos mediante consultas de aprendizaje.

Introducción

TL; DR. SOLQ es un marco de segmentación de instancias de extremo a extremo con Transformer. Produce directamente las máscaras de instancia sin ninguna dependencia de caja.

Resumen. En este documento, proponemos un marco de extremo a extremo para la segmentación de instancias. Basado en el DETR recientemente introducido, nuestro método, denominado SOLQ, segmenta objetos mediante el aprendizaje de consultas unificadas. En SOLQ, cada consulta representa un objeto y tiene múltiples representaciones: clase, ubicación y máscara. Las consultas de objetos aprendidas realizan clasificación, regresión de cuadro y codificación de máscara simultáneamente en una forma de vector unificado. Durante la fase de entrenamiento, los vectores de máscara codificados son supervisados ​​por la codificación de compresión de máscaras espaciales sin procesar. En tiempo de inferencia, los vectores de máscara producidos pueden transformarse directamente en máscaras espaciales mediante el proceso inverso de codificación de compresión. Los resultados experimentales muestran que SOLQ puede lograr un rendimiento de vanguardia, superando la mayoría de los enfoques existentes. Además, el aprendizaje conjunto de la representación de consultas unificadas puede mejorar en gran medida el rendimiento de detección del DETR original. Esperamos que nuestro SOLQ pueda servir como base de referencia sólida para la segmentación de instancias basadas en Transformer.

Actualizaciones

• (10/12/2021) Lanzamiento de D-DETR+SQR log.txt en SQR.
• (29/09/2021) Nuestro SOLQ ha sido aceptado por NeurIPS 2021.
• (14/07/2021) Se reporta un mayor rendimiento (Box AP=56.5, Mask AP=46.7) al entrenar con el lado largo 1536 en la columna vertebral Swin-L, en lugar del lado largo 1333.

Resultados principales

Instalación

El código base está construido sobre Deformable DETR.

Requisitos

• Linux, CUDA>=9.2, GCC>=5.4

• Pitón>=3.7

  Le recomendamos que utilice Anaconda para crear un entorno conda:

  conda create -n deformable_detr python=3.7 pip

  Luego, activa el entorno:

  conda activate deformable_detr

• PyTorch>=1.5.1, torchvision>=0.6.1 (siguiendo las instrucciones aquí)

  Por ejemplo, si su versión de CUDA es 9.2, podría instalar pytorch y torchvision de la siguiente manera:

  conda install pytorch=1.5.1 torchvision=0.6.1 cudatoolkit=9.2 -c pytorch

• Otros requerimientos

  pip install -r requirements.txt

• Construir MultiEscalaDeformableAtención

  cd ./models/ops
  sh ./make.sh

Uso

Preparación del conjunto de datos

Descargue COCO y organícelos de la siguiente manera:

mkdir data && cd data
ln -s /path/to/coco coco

Capacitación y Evaluación

Entrenamiento en un solo nodo

Entrenando SOLQ en 8 GPU de la siguiente manera:

sh configs/r50_solq_train.sh

Evaluación

Puede descargar el modelo previamente entrenado de SOLQ (el enlace está en la sesión “Resultados principales”), luego ejecute el siguiente comando para evaluarlo en el conjunto de datos de valor de COCO 2017:

sh configs/r50_solq_eval.sh

Evaluación del conjunto de datos de prueba y desarrollo de COCO 2017

Puede descargar el modelo previamente entrenado de SOLQ (el enlace está en la sesión “Resultados principales”), luego ejecutar el siguiente comando para evaluarlo en el conjunto de datos de desarrollo de prueba COCO 2017 (enviar al servidor):

sh configs/r50_solq_submit.sh

Visualización en el conjunto de datos val COCO 2017

Puede visualizar en la imagen de la siguiente manera:

EXP_DIR=/path/to/checkpoint
python visual.py 
       --meta_arch solq 
       --backbone resnet50 
       --with_vector 
       --with_box_refine 
       --masks 
       --batch_size 2 
       --vector_hidden_dim 1024 
       --vector_loss_coef 3 
       --output_dir ${EXP_DIR} 
       --resume ${EXP_DIR}/solq_r50_final.pth 
       --eval    

Citando SOLQ

Si encuentra SOLQ útil en su investigación, considere citar:

@article{dong2021solq,
  title={SOLQ: Segmenting Objects by Learning Queries},
  author={Dong, Bin and Zeng, Fangao and Wang, Tiancai and Zhang, Xiangyu and Wei, Yichen},
  journal={NeurIPS},
  year={2021}
}