Treine e implemente um modelo de classificação personalizado com YOLOv8
Introdução
Neste guia vamos explicar como treinar e implementar um modelo de classificação personalizado com YOLOv8
Visão geral
Vamos criar um ambiente virtual onde iremos instalar o YOLOv8, baixar um modelo de classificação do roboflow, treiná-lo e implantá-lo.
Classificação de Imagens
Classificação de imagens é a tarefa mais simples de visão computacional e envolve classificar uma imagem em uma das classes predefinidas. O que obtemos como saída é um único rótulo de classe e uma pontuação de confiança.
A classificação de imagens é útil quando não precisamos saber a localização do objeto na imagem e só precisamos saber a que classe a imagem pertence.
Requisitos de Materiais
Configuração de Hardware
Para este tutorial, vamos precisar de uma Nvidia Jetson Orin NX 16GB.

Configuração de Software
- JetPack 6.0 instalado no reComputer
- uma conta Roboflow para baixar o conjunto de dados
Preparando o reComputer
O reComputer J4012 da Seeed Studio é um Jetson Orin NX 16GB. É uma máquina poderosa, mas o Tegra Linux vem com muitas coisas e inicializa em modo gráfico. Vamos mudar isso.
Vou executar os exemplos e programar remotamente usando o VScode e um terminal SSH com X forwarding habilitado. X forwarding é uma opção do SSH que pode executar alguns aplicativos gráficos do nosso lado da conexão, em vez do computador remoto.
Se você vai se conectar ao seu reComputer com monitor, teclado e mouse, pule a próxima etapa.
Alterar modo de boot

Está tudo bem, mas não vamos precisar de gráficos e, em modo ocioso, está consumindo cerca de 1.5GB de memória.

Vamos fazer com que ele inicie na linha de comando.
sudo systemctl set-default multi-user
A partir de agora, nosso reComputer, ao inicializar, irá iniciar no CLI. Se quiser, você pode reiniciar agora ou podemos apenas ir para o CLI com um comando.
sudo systemctl isolate multi-user
Agora passamos de usar 1.5GB de memória para 700MB. Cada byte de memória conta quando usamos Machine Learning.
Alterar modo de energia
Por padrão, nosso reComputer deve estar rodando no nível 2 - 15W . Ao treinar, ou mesmo ao fazer inferência de um modelo de ML, se pudermos rodar em potência máxima, deve ser melhor.
Vamos aprender como mudar isso.
No arquivo /etc/nvpmodel.conf, temos os modos de energia disponíveis.
< POWER_MODEL ID=0 NAME=MAXN >
< POWER_MODEL ID=1 NAME=10W >
< POWER_MODEL ID=2 NAME=15W >
< POWER_MODEL ID=3 NAME=25W >
Podemos então usar sudo nvpmodel -m <power model number> para alterar o modo de energia. E, de acordo com este post do fórum, as configurações são mantidas mesmo após reinicializações.
Para ver em qual nível de energia estamos agora,
sudo nvpmodel -q

Vamos selecionar o modo de potência máxima para o treinamento do nosso modelo
sudo nvpmodel -m 0

Após uma reinicialização, podemos confirmar que estamos rodando em potência máxima

Treinando o modelo
Para o treinamento do modelo, vamos usar YOLOv8. Abaixo estão os passos necessários para instalá-lo com suporte a CUDA. Também vamos precisar de uma conta roboflow.
O modelo
Vou criar um modelo para classificar pássaros. Isto é parte de um projeto de um Comedouro Inteligente para Pássaros que vou colocar no meu jardim e quero saber quais são os pássaros que estão se alimentando ali.
Como esta é uma tarefa de classificação, não precisamos saber a posição do pássaro na foto.
Você pode usar outro conjunto de dados de sua escolha, desde que seja um conjunto de dados ou modelo de Classificação
Eu obtive 12 classes de pássaros que sei que vivem na minha região e são comuns perto de mim e criei um conjunto de dados de classificação no Roboflow.
As classes de pássaros que vou tentar identificar são:
- Andorinha-das-palhas (Barn Swallow)
- Estrelinha-de-cabeça-listada (Common Firecrest)
- Rouxinol-comum (Common Nightingale)
- Tentilhão-comum (Eurasian Chaffinch)
- Andorinhão-das-rochas (Eurasian Crag Martin)
- Pintassilgo-europeu (European Goldfinch)
- Verdilhão-europeu (European Greenfinch)
- Verdilhão-serim (European Serin)
- Pardal-doméstico (House Sparrow)
- Pardal-espanhol (Spanish Sparrow)
- Andorinha-das-chaminés-ocidental (Western House Martin)
- Lavadeira-branca (white Wagtail)
Escolha seu conjunto de dados e faça o download a partir do roboflow. Depois de selecionar seu conjunto de dados, selecione "Download Dataset". - Você precisa de uma conta para isso.

Em seguida, selecione Folder Structure em Format e selecione show download code.

Depois, selecione Jupyter se você for usar um Jupyter Notebook ou Terminal se estiver planejando fazer isso no terminal.
Eu selecionei Jupyter, para usar isso em um notebook Jupyter. Copie o código.

Criando o ambiente
Vamos criar um ambiente virtual, instalar o PyTorch e instalar o YOLOv8. Per YOLOv8 documentation tip, é melhor instalar o PyTorch primeiro e depois o ultralytics.
Além disso, estou instalando pacotes jupyterlab para usar com o VSCode. O notebook está anexado a este tutorial.
Vamos instalar algumas dependências primeiro.
NOTA: Como vamos usar YOLOv8, precisamos fazer alguns passos que normalmente não precisaríamos.
Apenas instalar o Torch seguindo a documentação de deep learning da NVIDIA é suficiente para ter o Torch com suporte a CUDA.
Se instalarmos o PyTorch normalmente com PIP, ele não terá suporte a CUDA.
Dependências
sudo apt install libopenblas-dev cuda-toolkit libcudnn8 tensorrt python3-libnvinfer nvidia-l4t-dla-compiler
Crie o ambiente virtual Python
python -m venv birdClassificationModel
Se você obtiver um erro, é porque o pacote python3-venv não está instalado. Vamos instalá-lo e repetir o comando acima.
sudo apt install python3-venv
Ative o ambiente virtual
source birdClassificationModel/bin/activate
Você pode confirmar que ele está ativo porque o nome dele é colocado antes do seu prompt.
YOLOv8
Antes, e para seguir a dica da documentação, vamos primeiro instalar o PyTorch.
Estou usando JetPack 6.0, que vem com NVIDIA Jetson Linux 36.3 e CUDA 12.2. Vamos atualizar o PIP primeiro
pip install -U pip
Para instalar o Torch para poder usá-lo com YOLOv8, precisamos seguir os passos nos fóruns da NVIDIA.
Isso será feito com o ambiente virtual ativo, para que seja instalado nele. Baixe a versão 2.3 do Torch da NVIDIA
wget https://nvidia.box.com/shared/static/mp164asf3sceb570wvjsrezk1p4ftj8t.whl -O torch-2.3.0-cp310-cp310-linux_aarch64.whl
sudo apt-get install python3-pip libopenblas-base libopenmpi-dev libomp-dev
pip3 install 'Cython<3'
pip install numpy torch-2.3.0-cp310-cp310-linux_aarch64.whl
Depois disso, vamos compilar o torchvision. Se o instalarmos pelos wheels, ele não terá suporte a CUDA.
A versão do branch é para a versão do Torch instalada. Você pode ver mais detalhes na página do fórum.
Lembre-se, você precisa ter o ambiente virtual ativo para que tudo isso seja instalado nele.
sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libopenblas-dev libavcodec-dev libavformat-dev libswscale-dev
git clone --branch v0.18.0 https://github.com/pytorch/vision torchvision
cd torchvision/
export BUILD_VERSION=0.18.0
python setup.py install
Depois de um tempo, ele será compilado e instalado.

Depois de instalar, vamos ver se o Cuda está disponível.

Pela linha de comando
python -c "import torch;print (torch.cuda.is_available())"
Isso deve retornar True
Instalar YOLOv8
Agora que temos o PyTorch instalado com suporte a CUDA, quando instalarmos o YOLOv8, ele usará a versão instalada em vez de tentar instalar um novo pacote (embora seja a mesma versão) sem o suporte a CUDA.
pip install ultralytics
Vamos instalar o roboflow e o jupyterlab
pip install roboflow jupyterlab
Agora, vamos baixar o dataset. Se você estiver usando o notebook, apenas substitua o código lá.
rf = Roboflow(api_key="<your_api_key>")
project = rf.workspace("bruno-santos-omqsq").project("bird-classification-19z7c")
version = project.version(1)
dataset = version.download("folder")
Depois de baixar o modelo, agora temos um conjunto de três diretórios (test, train, valid), cada um com um certo número de imagens de cada classe. Cada imagem de cada classe está no seu próprio diretório. Como isto é para classificação de imagens, não precisamos rotular as imagens. O YOLOv8 saberá as classes, não apenas a partir do arquivo de configuração que vamos criar depois, mas também a partir dos diretórios.

Treinar
Normalmente um dataset tem as imagens e os rótulos (ou anotações) com as coordenadas do objeto. Como esta é uma tarefa de classificação, não precisamos de nada disso. Apenas que as imagens estejam em cada diretório que é o nome da classe.
Preparando o arquivo de configuração
Ainda precisamos de um arquivo de configuração para o YOLOv8 reconhecer as classes.
Este arquivo deve ser colocado dentro do diretório do dataset, com extensão .yaml. O nome não é importante.
cd <dataset_directory>
vi birdClassificationModel.yaml
Insira o seguinte texto no arquivo
train: train/
valid: valid/
test: test/
# number of classes
nc: 12
# class names
names: ["Barn Swallow","Common Firecrest","Common Nightingale","Eurasian Chaffinch","Eurasian Crag Martin","European Goldfinch","European Greenfinch","European Serin","House Sparrow","Spanish Sparrow","Western House Martin","white Wagtail"]
Para a classificação, vamos usar um dos modelos pré‑treinados já disponíveis da Ultralytics.

Esses modelos foram treinados no ImageNet e são ajustados para classificação. Vamos usá‑lo e treiná‑lo com nossos dados.
Isto é o que é conhecido como transfer learning.
Vou usar o modelo YOLOv8l-cls. Provavelmente outros também funcionarão bem, mas como não precisamos de tempo real, é um compromisso entre velocidade e precisão.
Então vamos treinar o modelo, usando a interface de linha de comando (CLI) do YOLOv8
yolo task=classify mode=train model=yolov8l-cls.pt data=Bird-Classification-1 epochs=100
- task=classify : Vamos classificar imagens
- mode=train : Estamos treinando o modelo
- model=yolov8l-cls.pt : Estamos usando um modelo pré‑treinado em classificação
- data=Bird-Classification-1 : o diretório onde nosso dataset está localizado
- epochs=100 : por quanto tempo estamos treinando o modelo.
Agora que está executando, aqui estão algumas estatísticas usando jtop (tegra-stats)



Depois de algumas horas, o treinamento é concluído.

Agora, vamos ver como o modelo se comporta. Vamos testá‑lo.
yolo task=classify mode=predict model='./runs/classify/train6/weights/best.pt' source=Bird-Classification-1/test/**/*.jpg
Isto fará o yolo entrar nos diretórios de teste e tentar prever cada




Os resultados estão todos corretos. Vamos tentar com duas imagens que ele nunca viu.


yolo task=classify mode=predict model='./runs/classify/train6/weights/best.pt' source=house_sparrow.jpg

yolo task=classify mode=predict model='./runs/classify/train6/weights/best.pt' source=white_wagtail.jpg

Vou dizer que esses resultados estão ótimos
Exportando o modelo
Podemos usar o modelo como está para a inferência, só precisamos abri‑lo e usá‑lo. Para tempos de inferência mais rápidos podemos exportá‑lo para TensorRT, já que estamos em um NVIDIA Jetson Orin NX, ou mesmo ONNX, por exemplo.
Não é que precisemos de tempos de inferência mais rápidos para este projeto - não vou usar isso em vídeo em tempo real - mas é legal podermos tirar proveito da plataforma em que estamos.
Infelizmente, devido ao ambiente virtual, não consegui exportá‑lo para TensorRT. Por alguma razão, não consegui importar o tensorrt em Python, mas fora do ambiente virtual, não tive problema com as bibliotecas tensorrt.
ONNX
Podemos exportar o modelo para o formato ONNX assim
yolo export model='./runs/classify/train6/weights/best.pt' format=onnx imgsz=640
Obtemos um best.onnx que podemos usar para executar inferência.
Para executar inferência usando ONNX, precisamos instalar o wheel onnxruntime_gpu.
Para instalar o onnxruntime-gpu com o JetPack 6.0, precisamos baixá‑lo do Jetson Zoo.
Vamos baixar o onnxruntime_gpu 1.18.0
Baixe o wheel do pip para nossa versão do Python (Python-3.10)
wget https://nvidia.box.com/shared/static/48dtuob7meiw6ebgfsfqakc9vse62sg4.whl -O onnxruntime_gpu-1.18.0-cp310-cp310-linux_aarch64.whl
e então, instale‑o
pip install onnxruntime_gpu-1.18.0-cp310-cp310-linux_aarch64.whl
Inferência
foto
Usei o seguinte código para executar inferência com o modelo best.pt e ver os resultados
# running inference
from ultralytics import YOLO
# load the model
bird_model = YOLO("./runs/classify/train6/weights/best.pt")
#run inference
results = bird_model("house_sparrow.jpg")[0]
# get class names
class_names = results.names
# get top class with more probability
top1 = results.probs.top1
# print the class name with the highest probability
print (f" The detected bird is: {class_names[top1]}")
O que o código acima faz é carregar o modelo, executar inferência em uma imagem e salvar os resultados na variável results.
Como results é um objeto ultralytics.engine.results.Results do tipo lista com um item, que é uma instância de Results. O [0] na variável results para manter o resultado da inferência nos permitirá chegar aos resultados que queremos.
results = bird_model("house_sparrow.jpg")[0]
Em seguida, usamos os resultados para obter os nomes das classes. Não que não os conheçamos, mas assim permitirá que este código funcione em outros modelos também.
class_names = results.names
Um dos resultados é uma variável top1 que guarda a classe TOP 1 com mais probabilidades. Esse TOP1 é dado pela lista probs.
top1 = results.probs.top1
Em seguida, imprimimos a classe com a probabilidade mais alta, que deve ser a espécie do pássaro.
print (f" The detected bird is: {class_names[top1]}")
The detected bird is: House Sparrow
Câmera
Agora, vamos usar uma câmera para executar inferência.
O Jetson pode usar uma câmera USB ou uma câmera RPI. Vou conectar uma câmera USB.
O código a seguir verificará se é possível exibir um feed da câmera.
#Lets test if we can use a USB camera
import cv2
cap = cv2.VideoCapture(0)
while True:
ret, img = cap.read()
cv2.imshow('Camera', img)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows
Sou eu, no meu computador desktop. Basta usar ssh -X username@jetson_ip e a janela X11 será encaminhada para o seu desktop. Isso funciona porque também estou usando Linux. Acho que o WSL também poderia funcionar.

Agora, vamos tentar executar inferência em um feed de vídeo, exibindo a classe com a maior probabilidade
Aqui está o código
# again, save this code in a file a run it from the Jetson
import cv2
from ultralytics import YOLO
import time
#define confidence level
#only equal or above this level we say it's a class of bird
confidence = 0.95
# time when processed last frame
prev_frame = 0
# time processed current frame
cur_time = 0
# load the model
bird_model = YOLO("./runs/classify/train6/weights/best.pt")
# cv2 font
font = cv2.FONT_HERSHEY_SIMPLEX
# open camera
cap = cv2.VideoCapture(0)
while True:
ret, img = cap.read()
# to display fps
cur_frame = time.time()
fps = 1 / (cur_frame - prev_frame)
prev_frame = cur_frame
fps = int(fps)
fps = str(fps)
cv2.putText (img, fps, (550,50), font, 1, (124,10,120), 2, cv2.LINE_AA)
# inference current frame
results = bird_model(img, verbose=False)[0]
# get class names
class_names = results.names
# get top class with more probability
top1 = results.probs.top1
top1conf = results.probs.top1conf.tolist()
# we will only show the class name if the confidence is higher than defined level
# print the class name with the highest probability
if (top1conf >= confidence):
bird_class = class_names[top1]
print (f" The detected bird is: {class_names[top1]}")
# color is in BGR
confid = round(top1conf,2)
img = cv2.putText(img, bird_class, (50,50), font, 0.9, (0, 0, 255), 2, cv2.LINE_AA)
img = cv2.putText(img, "Conf: " + str(confid), (50,80), font, 0.6, (255, 0, 255), 1, cv2.LINE_AA)
cv2.imshow('Camera', img)
else:
img = cv2.imshow('Camera', img)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows
Aqui está um vídeo mostrando a inferência em um fluxo de vídeo
✨ Projeto de Colaborador
- Este projeto é apoiado pelo Seeed Studio Contributor Project.
- Obrigado pelos esforços do Bruno e seu trabalho será exibido.
Suporte Técnico e Discussão de Produtos
Obrigado por escolher nossos produtos! Estamos aqui para fornecer diferentes tipos de suporte para garantir que sua experiência com nossos produtos seja a mais tranquila possível. Oferecemos vários canais de comunicação para atender a diferentes preferências e necessidades.