YOLO para principantes
En Eryx nos gusta compartir el conocimiento y eso es lo que hacemos en Hoy Aprendí. Descubrí tips e insights claves de nuestro equipo técnico para aplicarlos en tu día a día. En esta edición especial te contamos los básicos de YOLO.
Contexto
YOLO es un algoritmo muy conocido de detección y clasificación de objetos. Este utiliza una única red neuronal para detectar y/o clasificar uno o más objetos en una imagen o video. La detección es básicamente dibujar bounding boxes (cajitas) adonde se detecten los objetos. La clasificación consiste en calcular la probabilidad de que lo que se está viendo sea de cierta clase (por ejemplo, lo analizado parecer ser “87% una persona”). YOLO viene con algunas detecciones preentrenadas con el COCO dataset (traducido: Objetos Comunes en Contexto), por lo cual puede reconocer personas, valijas, manzanas, bananas, y muchos otros objetos comunes.
Lo que vamos a hacer ahora es un ejemplo básico de YOLO out-of-the-box: contar cuántas manzanas hay en esta imagen.
Cómo lo resolvimos
Usamos YOLOv8 con Python, por lo cual el primer paso es instalar la dependencia pip install ultralytics
Luego podemos correr este código:
from ultralytics import YOLO
# Cargar el modelo
model = YOLO("yolov8n.pt")
# Realizar predicción
image_name = "apples.webp"
image_path = f"assets/{image_name}"
results = model.predict(image_path, conf=0.4, save=True)Vayamos línea por línea rápidamente:
- Como pueden ver, primero se importa
YOLOy tenemos que cargar un modelo. El modelo que vamos a usar esyolov8n(o YOLOv8-N [Nano]). YOLOv8 es de los mejores hoy en día y el más utilizado, aunque la elección del mismo va a depender del problema a resolver. Apartado respecto al tamaño: un modelo como YOLOv8-N (Nano) tendría menos capas y menos complejas en comparación con YOLOv8-L (Grande), lo que lo hace más liviano y rápido, pero potencialmente menos preciso. - Al cargar el modelo con
YOLO("yolov8n.pt")el programa va a buscar ese archivo en nuestra computadora y si no lo encuentra lo va a bajar automáticamente. - Luego lo que hacemos es llamar a
model.predict(...)con la imagen de las manzanas que llaméapples.webpy guardé en la carpetaassets. El parámetroconfsirve para detectar objetos que superan la confianza 0.4, es decir, si YOLO cree que lo que está viendo es al menos 40% una manzana, entonces va a dibujar una bounding box con eso. El parámetrosaveva a guardar la imagen resultado en nuestro disco.
Si corremos este código, vamos a ver este resultado en la consola y esta es la imagen resultante:
image 1/1 assets/apples.webp: 384x640 8 apples, 195.9ms
Speed: 6.6ms preprocess, 195.9ms inference, 1389.0ms postprocess per image at shape (1, 3, 384, 640)
Results saved to runs/detect/predictAcá podemos ver qué imagen analizó, cuánto tardó (195.9ms), qué encontró (8 apples) y adónde dejó el resultado.
Y en runs/detect/predict hay una imagen apples.webp que se ve de esta forma:
¡Ya tenemos una imagen con las manzanas detectadas! Y vemos que detectó 8 manzanas, justo la cantidad de la imagen. Dicho esto, vemos que la manzana de atrás tiene una confianza de 0.44, es decir que si poniamos conf=0.5 en vez de conf=0.4 YOLO no la iba a detectar. Acá va la imagen con esta configuración para que vean la diferencia.
Ahora, yo tal vez quiera hacer algo con la cantidad de manzanas que reconoció YOLO. Para eso podemos agregar esto al código:
# Tomo el primer resultado ya que solo estamos analizando una imagen
result = results[0]
# Contar la cantidad de manzanas detectadas, es decir la cantidad de bounding boxes que dibujó YOLO
amount_of_apples_detected = len(result.boxes)
print(f"Número de manzanas detectadas: {amount_of_apples_detected}")model.predictda una lista de resultados por cada frame analizado, y en este caso le estamos pasando solo un frame, por lo cual tiene un solo resultado. Habría muchos más resultados si estuvieramos pasándole un video (YOLO soporta muchas fuentes para analizar, desde imágenes hasta videos o mismo streams, incluso hasta desde YouTube directamente).- Tomar los
boxesdel result es literalmente tomar los bounding boxes, las cajitas que dibujó YOLO y contarlas: eso nos dará la cantidad de manzanas detectadas. - Luego las puedo imprimir en pantalla, o hacer lo que querramos con ese dato.
¡Eso es todo por ahora! Esta fue una pequeña introducción al mundo de Computer Vision usando YOLO pero un gran paso para todos aquellos que aspiran ser los próximos ingenieros en IA.
Todavía queda mucho por explorar, por ejemplo entrenar a YOLO para que reconozca algo custom que no esté en el COCO dataset, usar detección de pose, segmentación, combinar modelos, etc. ¡Hasta la próxima!
Extras
Si quisieran dibujar la cantidad de manzanas que hay directamente en la imagen, podemos usar OpenCV (cv2 ) para dibujarlo de esta forma. Este sería el código completo:
from ultralytics import YOLO
import cv2
import os
# Cargar el modelo
model = YOLO("yolov8n.pt")
# Realizar predicción
image_name = "apples.webp"
image_path = f"assets/{image_name}"
results = model.predict(image_path, conf=0.5, save=True)
# Tomo el primer resultado ya que solo estamos analizando una imagen
result = results[0]
# Contar la cantidad de manzanas detectadas, es decir la cantidad de bounding boxes que dibujó YOLO
amount_of_apples_detected = len(result.boxes)
#### A partir de acá, uso OpenCV para dibujar sobre la imagen que devolvió YOLO ####
# Cargar la imagen guardada por YOLO con las detecciones
yolo_image_path = os.path.join(result.save_dir, image_name)
image_to_modify = cv2.imread(yolo_image_path)
# Agregar texto a la imagen
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = 0.5
position = (7, image_to_modify.shape[0] - 7) # Posición en la esquina inferior izquierda
cv2.putText(image_to_modify, f"Cantidad de manzanas: {amount_of_apples_detected}", position, font, font_scale, (255, 255, 255), 1, cv2.LINE_AA)
# Guardar la imagen con el texto añadido
cv2.imwrite(yolo_image_path, image_to_modify)Este es el resultado:
