Treinar um modelo de detecção de objetos com TensorFlow

Agora que configuramos o TensorFlow, usaremos a arquitetura YOLO para treinar o modelo de detecção de objetos. YOLO é uma rede neural que prevê caixas delimitadoras e probabilidades de classe de uma imagem em uma única avaliação. Os modelos YOLO podem processar mais de 60 quadros por segundo, tornando-se uma ótima arquitetura para detectar objetos em vídeos. Você pode encontrar mais informações sobre como o YOLO funciona aqui.

Usando o YOLO

Primeiro, baixe este arquivo de exemplo YOLO, que contém scripts auxiliares para começar.

Ao usar o YOLO, temos três opções:

1. Use pesos de modelo pré-treinados para YOLO. O modelo foi pré-treinado em um grande conjunto de dados com 80 classes (categorias) para objetos cotidianos, como ônibus, pessoa, sanduíche etc. Se você quiser baixar um modelo YOLO pré-treinado no formato ONNX, poderá fazer isso aqui. Em seguida, você pode prosseguir para a fase final deste tutorial para saber como integrar esse modelo a um aplicativo.

2. Implemente o aprendizado de transferência com um conjunto de dados personalizado. O aprendizado de transferência é um método para usar um modelo treinado como ponto de partida para treinar um modelo resolvendo uma tarefa diferente, mas relacionada. Este tutorial usará os pesos YOLO pré-treinados com 80 classes para treinar um modelo com 20 classes com o conjunto de dados VOC. Se você quiser criar seu próprio conjunto de dados com classes personalizadas, confira as instruções aqui.

3. Treine YOLO do zero. Essa técnica não é recomendada, pois é muito difícil convergir. O artigo YOLO original treinou o Darknet na ImageNet (contendo centenas de milhares de fotos) antes de treinar toda a rede.

Implemente o aprendizado de transferência em pesos YOLO pré-treinados para o conjunto de dados VOC.

Vamos continuar com a segunda opção e implementar o aprendizado de transferência com as etapas a seguir.

1. Em uma janela miniconda, navegue até o diretório de exemplo yolo e execute o comando a seguir para instalar todos os pacotes pip necessários para YOLO.

pip install -r requirements.txt

2. Execute o script de instalação para baixar os dados e os pesos pré-treinados

python setup.py

3. Transforme o conjunto de dados. Consulte tools/voc2012.py para a implementação – esse formato é baseado na tensorflow object detection API. Muitos campos não são necessários, mas foram preenchidos aqui para compatibilidade com a API oficial.

python tools/voc2012.py \
  --data_dir './data/voc2012_raw/VOCdevkit/VOC2012' \
  --split train \
  --output_file ./data/voc2012_train.tfrecord

python tools/voc2012.py \
  --data_dir './data/voc2012_raw/VOCdevkit/VOC2012' \
  --split val \
  --output_file ./data/voc2012_val.tfrecord

4. Treinar o modelo. Execute os comandos a seguir:

python convert.py
python detect.py --image ./data/meme.jpg # Sanity check

python train.py \
	--dataset ./data/voc2012_train.tfrecord \
	--val_dataset ./data/voc2012_val.tfrecord \
	--classes ./data/voc2012.names \
	--num_classes 20 \
	--mode fit --transfer darknet \
	--batch_size 16 \
	--epochs 10 \
	--weights ./checkpoints/yolov3.tf \
	--weights_num_classes 80 

Agora você tem um modelo treinado novamente com 20 classes, pronto para uso.

Próximas etapas

Agora que criamos um modelo tensorFlow, precisamos convertê-lo no formato ONNX para uso com as APIs do Windows Machine Learning.