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

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

Usar YOLO

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

Ao usar YOLO, há três opções:

1. Usar pesos de modelo pré-treinados para YOLO. O modelo pré-treinado foi treinado em um grande conjunto de dados com 80 classes (categorias) para objetos diários, como ônibus, pessoa, sanduíche etc. Você pode baixar um modelo YOLO pré-treinado no formato ONNX aqui. Você pode prosseguir para a etapa final deste tutorial para saber como integrar esse modelo a um aplicativo.

2. Implementar a transferência de aprendizado com um conjunto de dados personalizado. A transferência de aprendizado é um método de uso de um modelo treinado como ponto de partida para treinar um modelo que resolve uma tarefa diferente, mas relacionada. Este tutorial vai 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 um conjunto de dados com classes personalizadas, confira as instruções aqui.

3. Treinar YOLO do zero. Essa técnica não é recomendada por ser muito difícil de convergir. O documento YOLO original treinou a darknet na imagenet (contendo centenas de milhares de fotos) antes de treinar toda a rede também.

Implementar a transferência de aprendizado em pesos YOLO pré-treinados para o conjunto de dados VOC:

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

1. Em uma janela do miniconda, navegue até o diretório de exemplo de 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. Confira tools/voc2012.py para implementação – esse formato é baseado na API de detecção de objetos tensorflow. Muitos campos não são necessários, mas foram preenchidos 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 e 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.