Treine o YOLOv5 com dados personalizados#

📚 Este guia explica como treinar seu próprio conjunto de dados personalizado usando o modelo YOLOv5 🚀. Treinar modelos personalizados é uma etapa fundamental para adaptar soluções de visão computacional a aplicações específicas do mundo real, além da detecção de objetos genérica.

Antes de começar#

Primeiro, verifique se o ambiente necessário está configurado. Clone o repositório YOLOv5 e instale as dependências necessárias de requirements.txt. É essencial usar um ambiente com Python>=3.8.0 e PyTorch>=1.8. Os modelos e conjuntos de dados são baixados automaticamente da versão mais recente do YOLOv5 se não forem encontrados localmente.

git clone https://github.com/ultralytics/yolov5 # Clone the repository
cd yolov5
pip install -r requirements.txt # Install dependencies

Treinar com dados personalizados#

Desenvolver um modelo personalizado de detecção de objetos é um processo iterativo:

1. Coletar e organizar imagens: reúna imagens relevantes para sua tarefa específica. Dados de alta qualidade e diversificados são essenciais. Consulte nosso guia sobre coleta e anotação de dados.
2. Rotular objetos: anote com precisão os objetos de interesse nas imagens.
3. Treinar um modelo: use os dados rotulados para treinar seu modelo YOLOv5. Aproveite o aprendizado por transferência, começando com pesos pré-treinados.
4. Implantar e fazer previsões: use o modelo treinado para inferência em dados novos e não vistos.
5. Coletar casos extremos: identifique situações em que o modelo tem um desempenho ruim (casos extremos) e adicione dados semelhantes ao conjunto de dados para aumentar a robustez. Repita o ciclo.

A Plataforma Ultralytics oferece uma solução simplificada, sem código, para todo esse ciclo de operações de aprendizado de máquina (MLOps), incluindo gerenciamento de conjuntos de dados, treinamento de modelos e implantação.

Antes de iniciar o treinamento, é essencial preparar o conjunto de dados.

1. Criar um conjunto de dados#

Os modelos YOLOv5 precisam de dados rotulados para aprender as características visuais das classes de objetos. Organizar corretamente o conjunto de dados é fundamental.

1.1 Criar dataset.yaml#

O arquivo de configuração do conjunto de dados (por exemplo, coco128.yaml) descreve a estrutura do conjunto de dados, os nomes das classes e os caminhos para os diretórios de imagens. COCO128 é um pequeno conjunto de dados de exemplo, composto pelas primeiras 128 imagens do extenso conjunto de dados COCO. Ele é útil para testar rapidamente o pipeline de treinamento e diagnosticar possíveis problemas, como sobreajuste.

A estrutura do arquivo dataset.yaml inclui:

• path: o diretório raiz que contém o conjunto de dados.
• train, val, test: caminhos relativos de path para diretórios que contêm imagens ou arquivos de texto com listas de caminhos de imagens dos conjuntos de treinamento, validação e teste.
• names: um dicionário que associa índices de classe (começando em 0) aos nomes de classe correspondentes.

Você pode definir path como um diretório absoluto (por exemplo, /home/user/datasets/coco128) ou como um caminho relativo, como ../datasets/coco128, ao iniciar o treinamento a partir do diretório raiz do repositório YOLOv5.

Abaixo está a estrutura de coco128.yaml (veja no GitHub):

# Dataset root directory relative to the yolov5 directory
path: ../datasets/coco128

# Train/val/test sets: specify directories, *.txt files, or lists
train: images/train2017 # 128 images for training
val: images/train2017 # 128 images for validation
test: # Optional path to test images

# Classes (example using 80 COCO classes)
names:
    0: person
    1: bicycle
    2: car
    # ... (remaining COCO classes)
    77: teddy bear
    78: hair drier
    79: toothbrush

1.2 Aproveitar modelos para rotulagem automatizada#

A rotulagem manual é a abordagem mais comum, mas consome muito tempo. Modelos de base podem automatizar ou semiautomatizar a anotação e acelerar a criação de conjuntos de dados. Veja alguns exemplos de modelos que podem ajudar a gerar rótulos:

• Google Gemini: modelos multimodais de grande porte, como o Gemini, têm recursos avançados de compreensão de imagens. É possível solicitar que identifiquem e localizem objetos nas imagens, gerando caixas delimitadoras ou descrições que podem ser convertidas em rótulos no formato YOLO. Explore o potencial do modelo no notebook de tutorial fornecido.
• SAM2 (Modelo Segment Anything 2): modelos de base voltados à segmentação, como o SAM2, conseguem identificar e delinear objetos com alta precisão. Embora sejam usados principalmente para segmentação, as máscaras resultantes muitas vezes podem ser convertidas em anotações de caixas delimitadoras adequadas para tarefas de detecção de objetos.
• YOLOWorld: este modelo oferece recursos de detecção com vocabulário aberto. Você pode fornecer descrições textuais dos objetos de interesse, e o YOLOWorld pode localizá-los nas imagens sem ter sido treinado previamente nessas classes específicas. Isso pode servir como ponto de partida para gerar rótulos iniciais, que podem ser refinados depois.

O uso desses modelos pode servir como uma etapa de "pré-rotulagem", reduzindo o trabalho manual necessário. No entanto, é fundamental revisar e refinar os rótulos gerados automaticamente para garantir a precisão e a consistência, pois a qualidade deles afeta diretamente o desempenho do modelo YOLOv5 treinado. Depois de gerar (e possivelmente refinar) os rótulos, verifique se eles seguem o formato YOLO: um arquivo *.txt por imagem, com cada linha representando um objeto como class_index x_center y_center width height (coordenadas normalizadas, classe indexada a partir de zero). Se uma imagem não tiver objetos de interesse, não é necessário criar um arquivo *.txt correspondente.

As especificações dos arquivos *.txt no formato YOLO são precisas:

• Uma linha por caixa delimitadora.
• Cada linha deve conter: class_index x_center y_center width height.
• As coordenadas devem ser normalizadas para um intervalo entre 0 e 1. Para isso, divida os valores em pixels de x_center e width pela largura total da imagem e divida y_center e height pela altura total da imagem.
• Os índices das classes começam em zero (ou seja, a primeira classe é representada por 0, a segunda por 1 e assim por diante).

O arquivo de rótulos correspondente à imagem acima, com dois objetos 'person' (índice de classe 0) e um objeto 'tie' (índice de classe 27), ficaria assim:

1.3 Organizar diretórios#

Estruture o diretório dos seus conjuntos de dados conforme ilustrado abaixo. Por padrão, o YOLOv5 espera que o diretório do conjunto de dados (por exemplo, /coco128) fique dentro de uma pasta /datasets, localizada ao lado do diretório do repositório /yolov5.

O YOLOv5 localiza automaticamente os rótulos de cada imagem substituindo a última ocorrência de /images/ no caminho da imagem por /labels/. Por exemplo:

../datasets/coco128/images/im0.jpg # Path to the image file
../datasets/coco128/labels/im0.txt # Path to the corresponding label file

A estrutura de diretórios recomendada é:

/datasets/
└── coco128/  # Dataset root
    ├── images/
    │   ├── train2017/  # Training images
    │   │   ├── 000000000009.jpg
    │   │   └── ...
    │   └── val2017/    # Validation images (optional if using same set for train/val)
    │       └── ...
    └── labels/
        ├── train2017/  # Training labels
        │   ├── 000000000009.txt
        │   └── ...
        └── val2017/    # Validation labels (optional if using same set for train/val)
            └── ...

2. Selecionar um modelo#

Escolha um modelo pré-treinado para iniciar o processo de treinamento. Começar com pesos pré-treinados acelera significativamente o aprendizado e melhora o desempenho em comparação com o treinamento do zero. O YOLOv5 oferece modelos de vários tamanhos, cada um equilibrando velocidade e precisão de forma diferente. Por exemplo, o YOLOv5s é o segundo menor e o segundo mais rápido, sendo adequado para ambientes com recursos limitados. Consulte a tabela do README para ver uma comparação detalhada de todos os modelos disponíveis.

3. Treinar#

Inicie o treinamento do modelo com o script train.py. Os argumentos essenciais incluem:

• --img: define o tamanho da imagem de entrada (por exemplo, --img 640). Tamanhos maiores geralmente proporcionam mais precisão, mas exigem mais memória da GPU.
• --batch: determina o tamanho do lote (por exemplo, --batch 16). Escolha o maior tamanho que sua GPU consegue suportar.
• --epochs: especifica o número total de épocas de treinamento (por exemplo, --epochs 100). Uma época corresponde a uma passagem completa por todo o conjunto de dados de treinamento.
• --data: caminho para seu arquivo dataset.yaml (por exemplo, --data coco128.yaml).
• --weights: caminho para o arquivo de pesos iniciais. É altamente recomendável usar pesos pré-treinados (por exemplo, --weights yolov5s.pt) para obter convergência mais rápida e resultados melhores. Para treinar do zero (não recomendado, a menos que você tenha um conjunto de dados muito grande e necessidades específicas), use --weights '' --cfg yolov5s.yaml.

Os pesos pré-treinados são baixados automaticamente da versão mais recente do YOLOv5 se não forem encontrados localmente.

# Example: Train YOLOv5s on the COCO128 dataset for 3 epochs
python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --weights yolov5s.pt

Todos os resultados do treinamento, incluindo pesos e logs, são salvos no diretório runs/train/. Cada sessão de treinamento cria um novo subdiretório (por exemplo, runs/train/exp, runs/train/exp2 etc.). Para uma experiência interativa e prática, explore a seção de treinamento nos nossos notebooks de tutorial oficiais:

4. Visualizar#

O YOLOv5 se integra perfeitamente a várias ferramentas para visualizar o progresso do treinamento, avaliar resultados e monitorar o desempenho em tempo real.

Registro e visualização com Comet#

A integração com o Comet está completa e permite acompanhar experimentos de forma abrangente. Visualize métricas em tempo real, salve hiperparâmetros, gerencie conjuntos de dados e checkpoints de modelos e analise previsões dos modelos usando os painéis personalizados do Comet interativos.

É fácil começar:

pip install comet_ml                                                          # 1. Install Comet library
export COMET_API_KEY=YOUR_API_KEY_HERE                                        # 2. Set your Comet API key (create a free account at Comet.ml)
python train.py --img 640 --epochs 3 --data coco128.yaml --weights yolov5s.pt # 3. Train your model - Comet automatically logs everything!

Saiba mais sobre os recursos compatíveis no nosso guia de integração com o Comet. Conheça melhor os recursos do Comet na documentação oficial. Experimente o notebook do Comet no Colab para ver uma demonstração ao vivo:

Registro e automação com ClearML#

A integração com o ClearML permite acompanhar experimentos detalhadamente, gerenciar versões de conjuntos de dados e até executar treinamentos remotamente. Ative o ClearML seguindo estas etapas simples:

• Instale o pacote: pip install clearml
• Inicialize o ClearML: execute clearml-init uma vez para se conectar ao servidor ClearML (auto-hospedado ou no plano gratuito).

O ClearML captura automaticamente detalhes dos experimentos, envios de modelos, comparações, alterações de código não confirmadas e pacotes instalados, garantindo a reprodutibilidade completa. Você pode agendar facilmente tarefas de treinamento em agentes remotos e gerenciar versões de conjuntos de dados usando o ClearML Data. Consulte o guia de integração com o ClearML para ver todos os detalhes.

Registro local#

Os resultados do treinamento são registrados automaticamente usando o TensorBoard e salvos como arquivos CSV no diretório do experimento correspondente (por exemplo, runs/train/exp). Os dados registrados incluem:

• Perda e métricas de desempenho do treinamento e da validação.
• Imagens de amostra mostrando as transformações aplicadas (como mosaicos).
• Rótulos de referência junto com as previsões do modelo para inspeção visual.
• Métricas de avaliação importantes, como curvas de precisão-revocação (PR).
• Matrizes de confusão para uma análise detalhada do desempenho por classe.

O arquivo results.csv é atualizado após cada época e, ao fim do treinamento, é plotado como results.png. Você também pode plotar manualmente qualquer arquivo results.csv usando a função utilitária fornecida:

from utils.plots import plot_results

# Plotar resultados de um diretório específico de execução do treinamento
plot_results("runs/train/exp/results.csv")  # Isso vai gerar 'results.png' no mesmo diretório

5. Próximas etapas#

Após a conclusão bem-sucedida do treinamento, o checkpoint do modelo com melhor desempenho (best.pt) é salvo e fica pronto para implantação ou refinamento adicional. Algumas próximas etapas possíveis:

• Execute a inferência em novas imagens ou vídeos usando o modelo treinado por meio da CLI ou do Python.
• Execute val.py para avaliar a precisão e a capacidade de generalização do modelo em diferentes divisões dos dados (por exemplo, um conjunto de teste separado).
• Exporte o modelo para vários formatos de implantação, como ONNX, TensorFlow SavedModel ou TensorRT, para obter inferência otimizada em diversas plataformas.
• Use a evolução de hiperparâmetros para tentar obter ganhos adicionais de desempenho.
• Continue aprimorando seu modelo seguindo nossas dicas para obter os melhores resultados de treinamento e adicionando iterativamente dados mais diversificados e desafiadores com base na análise de desempenho.

Ambientes compatíveis#

A Ultralytics oferece ambientes prontos para uso, equipados com dependências essenciais como CUDA, cuDNN, Python e PyTorch, para facilitar o início.

• Notebooks gratuitos com GPU:
  • 
  • 
  • 
• Plataformas na nuvem:
  • Google Cloud: Guia de início rápido do GCP
  • Amazon AWS: Guia de início rápido do AWS
  • Microsoft Azure: Guia de início rápido do AzureML
• Configuração local:
  • Docker: Guia de início rápido do Docker

Estado do projeto#

Este selo indica que todos os testes de integração contínua (CI) do GitHub Actions para YOLOv5 estão sendo aprovados. Esses rigorosos testes de CI abrangem as principais funcionalidades, incluindo treinamento, validação, inferência, exportação e benchmarks, nos sistemas operacionais macOS, Windows e Ubuntu. Os testes são executados automaticamente a cada 24 horas e a cada confirmação de código, garantindo estabilidade consistente e desempenho ideal.

Perguntas frequentes#