Skip to main content

Entrenar y desplegar un modelo de clasificación personalizado con YOLOv8

Introducción

En esta guía explicaremos cómo entrenar y desplegar un modelo de clasificación personalizado con YOLOv8

Descripción general

Crearemos un entorno virtual donde instalaremos YOLOv8, descargaremos un modelo de clasificación de roboflow, lo entrenaremos y lo desplegaremos.

Clasificación de Imágenes

La clasificación de imágenes es la tarea más simple de visión por computadora e involucra clasificar una imagen en una de las clases predefinidas. Lo que obtenemos como salida es una sola etiqueta de clase y una puntuación de confianza.

La clasificación de imágenes es útil cuando no necesitamos conocer la ubicación del objeto en la imagen y solo necesitamos saber a qué clase pertenece la imagen.

Requisitos de Materiales

Configuración de Hardware

Para este tutorial, vamos a necesitar un Nvidia Jetson Orin NX 16GB.

Configuración de Software

  • JetPack 6.0 instalado en el reComputer
  • una cuenta de Roboflow para descargar el conjunto de datos

Preparando el reComputer

El reComputer J4012 de Seeed Studio es un Jetson Orin NX 16GB. Es una máquina potente, pero el Tegra Linux viene con muchas cosas y arranca en modo gráfico. Cambiemos eso.

note

Voy a ejecutar los ejemplos y programar remotamente usando VScode y una terminal SSH con reenvío X habilitado. El reenvío X es una opción con SSH que puede ejecutar algunas aplicaciones gráficas en nuestro lado de la conexión, en lugar de la computadora remota.

Si vas a conectarte a tu reComputer con monitor, teclado y ratón, omite el siguiente paso.

Cambiar modo de arranque

Todo está bien, pero no vamos a necesitar gráficos y, en modo inactivo, está consumiendo alrededor de 1.5GB de memoria.

Vamos a hacer que arranque a la línea de comandos en su lugar.

sudo systemctl set-default multi-user

A partir de ahora, nuestro reComputer al arrancar iniciará en la CLI. Si quieres puedes reiniciar ahora o simplemente podemos ir a la CLI con un comando.

sudo systemctl isolate multi-user

Ahora hemos pasado de usar 1.5GB de memoria a 700MB. Cada byte de memoria cuenta cuando se usa Aprendizaje Automático.

Cambiar modo de energía

Por defecto, nuestro reComputer debería estar ejecutándose en nivel 2 - 15W. Cuando entrenamos, o incluso inferimos un modelo de ML, si podemos ejecutar a máxima potencia, debería ser mejor.

Aprendamos cómo cambiarlo.

En el archivo /etc/nvpmodel.conf, tenemos los modos de energía disponibles.

< POWER_MODEL ID=0 NAME=MAXN >
< POWER_MODEL ID=1 NAME=10W >
< POWER_MODEL ID=2 NAME=15W >
< POWER_MODEL ID=3 NAME=25W >

Luego podemos usar sudo nvpmodel -m <número de modelo de energía> para cambiar el modo de energía. Y, según esta publicación del hilo, la configuración se mantiene incluso después de reinicios. Para ver en qué nivel de energía estamos ahora,

sudo nvpmodel -q

Seleccionemos el modo de máxima potencia para el entrenamiento de nuestro modelo

sudo nvpmodel -m 0

Después de un reinicio, podemos confirmar que estamos ejecutando a máxima potencia

Entrenando el modelo

Para el entrenamiento del modelo, vamos a usar YOLOv8. A continuación están los pasos necesarios para instalarlo con soporte CUDA. También vamos a necesitar una cuenta de roboflow.

El modelo

Voy a crear un modelo para clasificar aves. Esto es parte de un proyecto para un Comedero Inteligente de Aves que voy a colocar en mi jardín y quiero saber qué aves son las que se alimentan allí.

Debido a que esta es una tarea de clasificación, no necesitamos conocer la posición del ave en la foto.

Puedes usar otro conjunto de datos de tu elección, siempre que sea un conjunto de datos o modelo de Clasificación

He conseguido 12 clases de aves que sé que viven en mi área y son comunes cerca de mí y creé un conjunto de datos de clasificación en Roboflow.

Las clases de aves que voy a tratar de identificar son:

  • 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

Elige tu conjunto de datos y descárgalo desde roboflow. Una vez que hayas seleccionado tu conjunto de datos, selecciona "Download Dataset". - Necesitas una cuenta para eso.

A continuación, selecciona Folder Structure en el Formato y selecciona show download code.

A continuación, selecciona Jupyter si vas a usar un Jupyter Notebook o Terminal si planeas hacer esto en la terminal.

He seleccionado Jupyter, para usar esto en un notebook de Jupyter. Copia el código.

Creando el entorno

Vamos a crear un entorno virtual, instalar PyTorch e instalar YOLOv8. Según el consejo de la documentación de YOLOv8, es mejor instalar PyTorch primero y luego ultralytics.

Además, estoy instalando paquetes de jupyterlab para usar con VSCode. El notebook está adjunto a este tutorial.

Vamos a instalar algunas dependencias primero.

NOTA: Debido a que vamos a usar YOLOv8, necesitamos hacer algunos pasos que normalmente no necesitaríamos.

Solo instalar Torch siguiendo la documentación de aprendizaje profundo de NVIDIA es suficiente para tener Torch con soporte CUDA.

Si instalamos PyTorch normalmente con PIP, no tendría soporte CUDA.

Dependencias

sudo apt install libopenblas-dev cuda-toolkit libcudnn8 tensorrt python3-libnvinfer nvidia-l4t-dla-compiler

Crear el entorno virtual de Python

python -m venv birdClassificationModel

Si obtienes un error, es porque el paquete python3-venv no está instalado. Vamos a instalarlo y repetir el comando anterior.

sudo apt install python3-venv

Activar el entorno virtual

source birdClassificationModel/bin/activate

Puedes confirmar que está activo porque el nombre del mismo se coloca antes de tu prompt.

YOLOv8

Antes, y para seguir el consejo de la documentación, instalemos primero PyTorch.

Estoy usando JetPack 6.0, que viene con NVIDIA Jetson Linux 36.3 y CUDA 12.2. Actualicemos PIP primero

pip install -U pip

Instalando Torch para poder usarlo with YOLOv8, necesitamos seguir los pasos en los foros de NVIDIA.

Esto se hará con el entorno virtual activo, para que se instale en él. Descargar Torch versión 2.3 de 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

Después de esto, compilemos torchvision. Si lo instalamos desde las wheels, no tendrá soporte para CUDA.

La versión de la rama es para la versión de Torch instalada. Puedes ver más detalles en la página del foro.

Recuerda, necesitas tener el entorno virtual activo para que todo esto se instale en él.

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

Después de un tiempo, se compilará e instalará.

Después de instalar, veamos si Cuda está disponible.

Desde la línea de comandos

python -c "import torch;print (torch.cuda.is_available())"

Esto debería devolver True

Instalar YOLOv8

Ahora que tenemos PyTorch instalado con soporte CUDA, cuando instalemos YOLOv8, utilizará la versión instalada en lugar de intentar instalar un nuevo paquete (aunque sea la misma versión) sin el soporte CUDA.

pip install ultralytics

Instalemos roboflow y jupyterlab

pip install roboflow jupyterlab

Ahora, descarguemos el conjunto de datos. Si estás usando el notebook, simplemente reemplaza el código allí.

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")

Después de descargar el modelo, ahora tenemos un conjunto de tres directorios (test, train, valid), cada uno con un cierto número de imágenes de cada clase. Cada imagen de cada clase está en su propio directorio. Debido a que esto es para clasificación de imágenes, no necesitamos etiquetar las imágenes. YOLOv8 conocerá las clases, no solo del archivo de configuración que crearemos más tarde, sino también de los directorios.

Entrenar

Usualmente un conjunto de datos tiene las imágenes y las etiquetas (o anotaciones) con las coordenadas del objeto. Dado que esta es una tarea de clasificación, no necesitamos nada de eso. Solo que las imágenes estén en cada directorio que tenga el nombre de la clase.

Preparando el archivo de configuración

Aún necesitamos un archivo de configuración para que YOLOv8 reconozca las clases. Este archivo debe colocarse dentro del directorio del conjunto de datos, con extensión .yaml. El nombre no es importante.

cd <dataset_directory>
vi birdClassificationModel.yaml

Inserta el siguiente texto en el archivo

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 la clasificación, vamos a usar uno de los modelos pre-entrenados ya disponibles de Ultralytics.

Estos modelos han sido entrenados en ImageNet y están ajustados para clasificación. Vamos a usarlo y entrenarlo con nuestros datos.

Esto es lo que se conoce como aprendizaje por transferencia.

Voy a usar el modelo YOLOv8l-cls. Probablemente otros también funcionen bien, pero como no necesitamos tiempo real, es un equilibrio entre velocidad y precisión.

Entonces entrenemos el modelo, usando la interfaz CLI de YOLOv8

yolo task=classify mode=train model=yolov8l-cls.pt data=Bird-Classification-1 epochs=100
  • task=classify : Vamos a clasificar imágenes
  • mode=train : Estamos entrenando el modelo
  • model=yolov8l-cls.pt : Estamos usando un modelo preentrenado en clasificación
  • data=Bird-Classification-1 : el directorio donde se encuentra nuestro conjunto de datos
  • epochs=100 : por cuánto tiempo estamos entrenando el modelo.

Ahora que está ejecutándose, aquí hay algunas estadísticas usando jtop (tegra-stats)

Después de un par de horas, el entrenamiento está completo.

Ahora, veamos cómo se comporta el modelo. Vamos a probarlo.

yolo task=classify mode=predict model='./runs/classify/train6/weights/best.pt' source=Bird-Classification-1/test/**/*.jpg

Esto hará que yolo vaya a los directorios de prueba e intente predecir cada

Los resultados son todos correctos. Probemos con dos imágenes que nunca ha visto.

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

Voy a decir que estos resultados son excelentes

Exportando el modelo

Podemos usar el modelo tal como está para la inferencia, solo necesitamos abrirlo y usarlo. Para tiempos de inferencia más rápidos podemos exportarlo a TensorRT, ya que estamos en un NVIDIA Jetson Orin NX, o incluso ONNX, por ejemplo.

No es que necesitemos tiempos de inferencia más rápidos para este proyecto - no voy a usar esto en video en tiempo real - pero es bueno que podamos aprovechar la plataforma en la que estamos.

Desafortunadamente, debido al entorno virtual no pude exportarlo a TensorRT. Por alguna razón, no pude importar tensorrt en Python, pero fuera del entorno virtual, no tuve problemas con las librerías de tensorrt.

ONNX

Podemos exportar el modelo al formato ONNX de esta manera

yolo export model='./runs/classify/train6/weights/best.pt' format=onnx imgsz=640

Obtenemos un best.onnx que podemos usar para ejecutar inferencia.

Para ejecutar inferencia usando ONNX, necesitamos instalar el wheel onnxruntime_gpu.

Para instalar onnxruntime-gpu con JetPack 6.0, necesitamos descargarlo desde el Jetson Zoo.

Vamos a descargar onnxruntime_gpu 1.18.0

Descarga el wheel pip para nuestra versión de Python (Python-3.10)

wget https://nvidia.box.com/shared/static/48dtuob7meiw6ebgfsfqakc9vse62sg4.whl -O onnxruntime_gpu-1.18.0-cp310-cp310-linux_aarch64.whl

y luego, instálalo

pip install onnxruntime_gpu-1.18.0-cp310-cp310-linux_aarch64.whl

Inferencia

foto

He usado el siguiente código para ejecutar la inferencia con el modelo best.pt y ver los 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]}")

Lo que hace el código anterior es cargar el modelo, ejecutar inferencia en una imagen y guardar los resultados en la variable results.

Debido a que results es un objeto ultralytics.engine.results.Results de tipo lista con un elemento, que es una instancia de Results. El [0] en la variable results para contener el resultado de inferencia nos permitirá obtener los resultados que queremos.

results = bird_model("house_sparrow.jpg")[0]

A continuación, usamos los resultados para obtener los nombres de las clases. No es que no los conozcamos, sino que de esta manera permitiremos que este código funcione también en otros modelos.

class_names = results.names

Uno de los resultados es una variable top1 que contiene la clase TOP 1 con más probabilidades. Ese TOP1 es proporcionado por la lista probs.

top1 = results.probs.top1

A continuación, imprimimos la clase de mayor probabilidad que debería ser la especie de ave.

print (f" The detected bird is: {class_names[top1]}")
The detected bird is: House Sparrow

Cámara

Ahora, usemos una cámara para ejecutar inferencia.

Jetson puede usar una cámara USB o una cámara RPI. Voy a conectar una cámara USB.

El siguiente código verificará si se puede mostrar la transmisión de la cámara.

#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

Este soy yo, en mi computadora de escritorio. Solo usa ssh -X nombre_usuario@jetson_ip y la ventana X11 se reenviará a tu escritorio. Esto funciona porque también estoy usando Linux. Creo que WSL también podría funcionar.

Ahora, intentemos ejecutar inferencia en un feed de video, mostrando la clase con la mayor probabilidad

Aquí está el 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

Aquí hay un video que muestra la inferencia en una transmisión de video

✨ Proyecto de Colaborador

  • Este proyecto está respaldado por el Proyecto de Colaboradores de Seeed Studio.
  • Gracias a los esfuerzos de Bruno y tu trabajo será exhibido.

Soporte Técnico y Discusión de Productos

¡Gracias por elegir nuestros productos! Estamos aquí para brindarte diferentes tipos de soporte para asegurar que tu experiencia con nuestros productos sea lo más fluida posible. Ofrecemos varios canales de comunicación para satisfacer diferentes preferencias y necesidades.

Loading Comments...