Esta é a versão de desenvolvimento do código para o seguinte artigo:
Bugra Tekin, Sudipta N. Sinha e Pascal Fua, "Previsão de pose de objeto 6D contínua e contínua em tempo real", CVPR 2018.
O repositório original da base de código do artigo acima pode ser encontrado no link a seguir.
Propomos uma abordagem de disparo único para detectar simultaneamente um objeto em uma imagem RGB e prever sua pose 6D sem exigir vários estágios ou examinar múltiplas hipóteses. O componente principal do nosso método é uma nova arquitetura CNN inspirada no design da rede YOLO que prevê diretamente as localizações da imagem 2D dos vértices projetados da caixa delimitadora 3D do objeto. A pose 6D do objeto é então estimada usando um algoritmo PnP. Papel, arXiv
Se você usar este código, cite o seguinte
@inproceedings{tekin18,
TÍTULO = {{Predição de pose de objeto 6D de tiro único contínuo em tempo real}}, AUTOR = {Tekin, Bugra e Sinha, Sudipta N. e Fua, Pascal},
TÍTULO DO LIVRO = {CVPR},
ANO = {2018}
}
SingleShotPose é lançado sob a licença MIT (consulte o arquivo LICENSE para obter detalhes).
O código foi testado em Windows com CUDA v8 e cudNN v5.1. A implementação é baseada em PyTorch 0.4.1 e testada em Python3.6 . O código requer as seguintes dependências que podem ser instaladas com conda ou pip: numpy, scipy, PIL, opencv-python. Para uma versão anterior compatível com PyTorch 0.3.1 e testada em Python2.7, consulte a pasta py2 .
Dentro do diretório de código principal, execute o seguinte para baixar e extrair (1) o conjunto de dados LINEMOD pré-processado, (2) modelos treinados para o conjunto de dados LINEMOD, (3) o modelo treinado para o conjunto de dados OCCLUSION, (4) imagens de fundo do VOC2012 conjunto de dados 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, você pode acessar diretamente os links acima e baixar e extrair manualmente os arquivos nos diretórios correspondentes. Todo o processo de download pode demorar muito (cerca de 60 minutos). Esteja ciente também de que o acesso ao OneDrive em alguns países pode ser limitado.
Para treinar a execução do 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]
por exemplo
python train.py --datacfg cfg/ape.data --modelcfg cfg/yolo-pose.cfg --initweightfile cfg/darknet19_448.conv.23 --pretrain_num_epochs 15
se você gostaria de começar a partir dos pesos inicializados do ImageNet, ou
python train.py --datacfg cfg/ape.data --modelcfg cfg/yolo-pose.cfg --initweightfile backup/duck/init.weights
se você quiser começar com um modelo já pré-treinado no LINEMOD, para uma convergência mais rápida.
[datacfg] contém informações sobre as divisões de treinamento/teste, modelos de objetos 3D e parâmetros de câmera
[modelcfg] contém informações sobre a estrutura da rede
[initweightfile] contém pesos de inicialização. <
No início do treinamento você verá um 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
Isso define a estrutura da rede. Durante o treinamento, o melhor modelo de rede é salvo no arquivo “model.weights”. Para treinar redes para outros objetos, basta alterar o nome do objeto enquanto chama a função de trem, por exemplo, " python train.py --datacfg cfg/duck.data --modelcfg cfg/yolo-pose.cfg --initweightfile backup/duck/init.weights ". Se você encontrar erros de memória da GPU durante o treinamento, tente diminuir o tamanho do lote, por exemplo, 16 ou 8, para caber na memória. A versão de código aberto do código passou por uma forte refatoração e, além disso, alguns modelos tiveram que ser retreinados. Os modelos retreinados que fornecemos não mudam muito em relação aos resultados iniciais que fornecemos (às vezes um pouco piores e às vezes um pouco melhores).
Para testar a execução do modelo
python valid.py --datacfg [path_to_data_config_file] --modelcfg [path_to_model_config_file] --weightfile [path_to_trained_model_weights]
por exemplo
python valid.py --datacfg cfg/ape.data --modelcfg cfg/yolo-pose.cfg --weightfile backup/ape/model_backup.weights
Você também pode usar valid.ipynb para testar o modelo e visualizar os resultados.
Dentro da pasta multi_obj_pose_estimation/
Teste:
python valid_multi.py cfgfile weightfile
por exemplo
python valid_multi.py cfg/yolo-pose-multi.cfg backup_multi/model_backup.weights
Treinamento:
python train_multi.py datafile cfgfile weightfile
por exemplo,
python train_multi.py cfg/occlusion.data cfg/yolo-pose-multi.cfg backup_multi/init.weights
Nossos arquivos de rótulos consistem em 21 valores reais. Prevemos 9 pontos correspondentes ao centróide e aos cantos do modelo de objeto 3D. Além disso, prevemos a classe em cada célula. Isso faz com que 9x2+1 = 19 pontos. No treinamento multiobjeto, durante o treinamento, atribuímos qualquer caixa de âncora que tenha o tamanho mais semelhante ao objeto atual como a responsável por prever as coordenadas 2D desse objeto. Para codificar o tamanho dos objetos, temos 2 números adicionais para o intervalo nas dimensões x e y. Portanto, temos 9x2+1+2 = 21 números.
Respectivamente, 21 números correspondem ao seguinte: 1º número: rótulo de classe, 2º número: x0 (coordenada x do centróide), 3º número: y0 (coordenada y do centróide), 4º número: x1 (coordenada x de primeiro canto), 5º número: y1 (coordenada y do primeiro canto), ..., 18º número: x8 (coordenada x do oitavo canto), 19º número: y8 (coordenada y do oitavo canto), 20º número: intervalo x, 21º número: intervalo y.
As coordenadas são normalizadas pela largura e altura da imagem: x / image_width e y / image_height . Isso é útil para ter intervalos de saída semelhantes para tarefas de regressão de coordenadas e classificação de objetos.
Treinamos e testamos nossos modelos no conjunto de dados LINEMOD usando as mesmas divisões de treinamento/teste do método BB8 para validar nossa abordagem. Se desejar treinar um modelo em seu próprio conjunto de dados, você pode criar a mesma estrutura de pastas com o conjunto de dados LINEMOD fornecido e ajustar os caminhos em cfg/[OBJECT].data, [DATASET]/[OBJECT]/train.txt e Arquivos [DATASET]/[OBJECT]/test.txt. A pasta de cada objeto deve conter o seguinte:
(1) uma pasta contendo arquivos de imagem,
(2) uma pasta contendo arquivos de rótulos (consulte este link para obter uma explicação detalhada sobre como criar rótulos. Você também pode achar a caixa de ferramentas ObjectDatasetTools de terceiros útil para criar rótulos verdadeiros para estimativa de pose de objeto 6D),
(3) um arquivo de texto contendo os nomes dos arquivos das imagens de treinamento ( train.txt ),
(4) um arquivo de texto contendo os nomes dos arquivos das imagens de teste ( test.txt ),
(5) um arquivo .ply contendo o modelo do objeto 3D (a unidade do modelo do objeto é dada em metros),
(6) opcionalmente, uma pasta contendo máscaras de segmentação (Se você deseja alterar o fundo de suas imagens de treinamento para que sejam mais robustos a fundos diversos, isso seria essencial para uma melhor capacidade de generalização),
Certifique-se também de ajustar os seguintes valores nos arquivos de configuração de dados e modelo de acordo com suas necessidades:
diam " no arquivo de configuração de dados com o diâmetro do modelo do objeto em questão.yolo-pose.cfg ) e alguns parâmetros de aumento de dados (jitter, matiz, saturação, parâmetros de exposição em dataset.py ) também pode precisar ser ajustado para uma melhor convergência em seu conjunto de dados.yolo-pose-multi.cfg ). Certifique-se também de usar o número correto de classes e especificá-lo em yolo-pose-multi.cfg .fx , fy , u0 , v0 , width , height ) nos arquivos de configuração de dados por aqueles específicos do seu conjunto de dados. Ao criar um conjunto de dados de treinamento, amostrar um grande número de pontos de vista/distâncias e modelar uma grande variabilidade de configurações de iluminação/oclusão/fundo seria importante para aumentar a capacidade de generalização da abordagem em seu conjunto de dados. Se você quiser ajustar alguns parâmetros do modelo e de perda (por exemplo, fator de ponderação para diferentes termos de perda) para seus próprios propósitos, poderá fazê-lo no arquivo de configuração do modelo ( yolo-pose.cfg ).
O código foi escrito por Bugra Tekin e construído na implementação YOLOv2 do usuário do github @marvis
Para qualquer dúvida ou relato de bugs, entre em contato com Bugra Tekin
