Clasificación de imágenes con MediaPipe | Python
¡Hola, hola Omesitos!.
Mediapipe nos ha traído nuevas tareas y es hora de explorarlas. Vamos con la segunda herramienta de visión por computador, la clasificación de imágenes. Al igual que en el post anterior veremos como funciona y como podemos aplicarla en imágenes, videos y videostreamings. Así que… ¡Vamos a por ello!.
CONTENIDO
- En visión artificial, ¿Qué es la clasificación de imágenes?
- Modelos de clasificación de imágenes con Mediapipe
- ¿Cómo usar la image classification en imágenes con Mediapipe – Python?
- ¿Cómo usar la clasificación de imágenes en video con Mediapipe – Python?
- ¿Cómo usar la image classification en video streaming con Mediapipe – Python?
En visión artificial, ¿Qué es la clasificación de imágenes?
La clasificación de imágenes es una de las tareas de visión por computador y machine learning en donde dada una imagen, se predecirá a qué clase o categoría pertenece. Para llevarlo a cabo necesitamos de un modelo. Este modelo es entrenado con muchísimas imágenes de cada una de las clases (o categorías), para que al terminar el entrenamiento, el modelo pueda distinguir a qué clase o categoría pertenece una nueva imagen que no haya pertenecido al conjunto de entrenamiento.

Y de hecho Mediapipe nos permite usar una de sus tareas de Visión artificial que es Image Classification. En ella nos proporciona algunos modelos ya entrenados que podemos usar y que podrían adaptarse a nuestros proyectos.

Mediapipe Vision Task
En el anterior blog post vimos como usar Object Detection, en esta ocasión veremos Image classification. La documentación oficial nos dice que con la Image Classification task podremos identificar qué representa una imagen dado el conjunto de categorías definidas durante el tiempo de entrenamiento. Así mismo, nos indica que podemos aplicarlo en imágenes o en flujo continuo como videos. Y que nos entregará una lista de potenciales categorías con su puntuación de probabilidad, de forma descendete.
En su documentación podemos encontrar las características de la tarea, el tipo de entradas que podemos usar como: imágenes, videos o video en vivo o videostreaming. Y también tenemos las salidas, como el score, category name, entre otros.
Modelos de clasificación de imágenes con Mediapipe
Pero para poder aplicar la clasificación de imágenes, así como lo hicimos para la detección de objetos, debemos tener un modelo que nos permita realizar esta tarea y en este caso MEDIAPIPE, nos deja escoger algunos que tendremos que descargar.

Tenemos el modelo recomendado EfficientNet-Lite0, que es el modelo que nos recomienda Mediapipe. Este ha sido entrenado con la arquitectura EfficientNet y usó como dataset ImageNet, por lo que puede reconocer entre 1000 clases. De hecho si damos clic en full list, podremos ver todas las categorías que podría reconocer, entre árboles, animales, comida, vehículos, entre otros.
Este modelo está disponible en dos versiones: int8 y float32, y para ambos casos el tamaño de la imagen de entrada sería de 224×224, pero de redimensionar se ocupa el propio Mediapipe. Aunque hay que tener en cuenta que no se debería tener como entrada una imagen con tamaño menor a 224×224 debido a que podría entorpecer los resultados. Además se debería respectar la relación de aspecto.
Int8 y float32 hacen referencia a la precisión numérica utilizada para representar los valores de los parámetros de un modelo. Por ejemplo, un modelo representado con 8 bits tiene menos precisión, pero mucho más rápido que uno de 32 bits.
Para descargar cada modelo simplemente damos clic sobre su nombre. Entonces lo que tendremos que hacer es descargar aquellos que deseemos probar. En lo personal recomendaría probarlos todos y ver cual se adapta mejor al problema a resolver.
Y por cierto, si aún no has instalado Mediapipe puedes hacerlo a través de:
pip install mediapipe
Una vez que tenemos el o los modelos, es hora de pasar a la programación…
¿Cómo usar la image classification en imágenes con Mediapipe – Python?
Vamos a crear un nuevo archivo de python. Y vamos a empezar por la importación de las librerías.
import mediapipe as mp from mediapipe.tasks.python import vision from mediapipe.tasks.python import BaseOptions import cv2
Lo primero que haremos es importar mediapipe, luego de mediapipe.task.python importaremos vision, ya que estaremos empleado una tarea de visión por computador, la cual nos permitirá manejar imágenes o videos. Luego importaremos BaseOptions que permitirá básicamente leer el modelo que se usará y finalmente importamos OpenCV.
# Especificar la configuración del clasificador options = vision.ImageClassifierOptions( base_options = BaseOptions(model_asset_path="/path/efficientnet_lite_int8.tflite"), max_results=5, score_threshold=0.01, running_mode=vision.RunningMode.IMAGE) classifier = vision.ImageClassifier.create_from_options(options)
Y precisamente vamos a empezar con las opciones de configuración a través de vision.ImageClassifierOptions. En primer lugar vamos a establecer el path del modelo que estaremos usando, en este caso efficientnet_lite_int8 (el modelo recomendado por Mediapipe). Luego tenemos a max_results en donde estableceremos 5, esto para que se muestre el top 5 de resultados, aunque podríamos ubicar otro número, por ejemplo el 1, 2, 3, etc. Para que se muestre ese top. Si especificamos -1, se nos devolverán todos los resultados (que estén por encima del score_threshold).
Vamos a establecer un umbral de puntuación para la probabilidad, en este caso he ubicado 0.01 que corresponderá al 1%. Y finalmente en running mode vamos a ubicar vision.RunningMode.Image, debido a que vamos a estar usando una imagen de entrada.
Una vez que tenemos las opciones de configuración, vamos a pasarlas al clasificador, para ello usamos vision.ImageClassifier.create_from_options(options).
Nota: Para más información sobre las opciones de configuración, no olvides echarle un vistazo a la sección Configuration Options de Mediapipe.
# Leer la imagen de entrada image = cv2.imread("/path/image.jpg") image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image_rgb = mp.Image(image_format=mp.ImageFormat.SRGB, data=image_rgb)
Ahora vamos a leer la imagen de entrada, para ello usaremos cv2.imread. Tendremos que especificar su ubicación o simplemente su nombre junto con su extensión si se encuentra en el mismo directorio que el archivo de Python en el que estamos trabajando.
Vamos a proceder a cambiar el orden de los canales de la imagen de entrada, transformaremos de BGR a RGB. Y además tendremos que convertir la imagen a un objeto mediapipe.Image. Para ello especificamos mp.ImageFormat.SRGB y data=image_rgb.
# Obtener los resultados del clasificador classifier_result = classifier.classify(image_rgb)
Para empezar con la clasificación tendremos que usar classifier.classify, y especificar la imagen que deseamos analizar.
y_aux = 0 for classification in classifier_result.classifications: for category in classification.categories: print(category) cv2.rectangle(image, (90, 0), (560,100), (255,255,255), -1) cv2.putText(image, f"{category.category_name}: {category.score * 100:.2f}%", (95, 80 + y_aux), cv2.FONT_HERSHEY_SIMPLEX, 2.1, (214, 88, 107), 5) y_aux += 40
Vamos a recorrer cada uno de los resultados obtenidos, estos valores se encontrarán en classifier_result.classifications. entonces tendremos que añadir un nuevo for para ir por cada una de las categorías, a través de clasification.categories.
Cada resultado está ordenado de forma descendente, y cuenta con:
- Index: el índice de la categoría en los resultados del modelo.
- Score: el nivel de confianza para la categoría, generalmente una probabilidad en el rango de [0,1].
- Category_name: nombre de visualización para la categoría según lo especificado en los Metadatos del Modelo TFLite
Recordemos que estos resultados los obtenemos debido a la configuración que establecimos previamente.
cv2.imshow("Image", image) cv2.waitKey(0) cv2.destroyAllWindows()
Finalmente visualizamos la imagen. Esperamos a que se presione una tecla para finalmente cerrar la ventana de visualización.
Una vez terminado el programa, probemos el modelo con un par de imágenes:


En la imagen de la izquierda tenemos que se ha clasificado a la imagen como un helado con un 23.44%, un mortero y una taza con 7.81%, una cuchara de madera con 7.03% y salsa de chocolate con 6.64%.
En la imagen de la derecha tenemos tiger cat con un 46.48%, gato egipcio con 29.30% y atigrado con 20.31%. Ojo que habíamos establecido 5 en max_results que estén por encima de 1% en su score. Tenemos solo 3 resultados porque los demás son inferiores a esa puntuación.
A continuación tenemos el programa completo sobre la clasificación de imágenes sobre una imagen:
import cv2 import mediapipe as mp from mediapipe.tasks.python import vision from mediapipe.tasks.python import BaseOptions # Especificar la configuración del clasificador options = vision.ImageClassifierOptions( base_options = BaseOptions(model_asset_path="/path/efficientnet_lite_int8.tflite"), max_results=5, score_threshold=0.01, running_mode=vision.RunningMode.IMAGE) classifier = vision.ImageClassifier.create_from_options(options) # Leer la imagen de entrada image = cv2.imread("/path/image.jpg") image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image_rgb = mp.Image(image_format=mp.ImageFormat.SRGB, data=image_rgb) # Obtener los resultados del clasificador classifier_result = classifier.classify(image_rgb) #print(classifier_result.classifications) y_aux = 0 for classification in classifier_result.classifications: for category in classification.categories: print(category) cv2.rectangle(image, (90, 0), (560,100), (255,255,255), -1) cv2.putText(image, f"{category.category_name}: {category.score * 100:.2f}%", (95, 80 + y_aux), cv2.FONT_HERSHEY_SIMPLEX, 2.1, (214, 88, 107), 5) y_aux += 40 cv2.imshow("Image", image) cv2.waitKey(0) cv2.destroyAllWindows()
¿Cómo usar la clasificación de imágenes en video con Mediapipe – Python?
Debido a que el programa para la clasificación en video presenta similitudes con la clasificación de imágenes sobre una imagen, procederemos a explicar las diferencias entre ambos, así que vamos a por ello.
import cv2 import mediapipe as mp from mediapipe.tasks.python import vision from mediapipe.tasks.python import BaseOptions # Especificar la configuración del clasificador options = vision.ImageClassifierOptions( base_options=BaseOptions(model_asset_path="/path/efficientnet_lite_float32.tflite"), max_results=5, score_threshold=0.01, running_mode=vision.RunningMode.VIDEO) classifier = vision.ImageClassifier.create_from_options(options)
La principal diferencia con el código anterior está presente en la línea 11, ya que en running_mode debemos especificar vision.RunningMode.VIDEO, puesto a que el análisis se llevará a cabo en un video de entrada. Además ahora estamos usando el modelo efficientnet_lite_float32.tflite (recuerda que puedes usar los otros modelos disponibles en Mediapipe).
# Leer el video de entrada cap = cv2.VideoCapture("path/video.mp4") frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT) fps = cap.get(cv2.CAP_PROP_FPS)
Con cap.get(cv2.CAP_PROP_FRAME_COUNT), obtendremos la cantidad total de fotogramas que posee el video de entrada. Mientras que con cap.get(cv2.CAP_PROP_FPS) obtendremos los fps del video. Estos datos los usaremos más adelante.
for frame_index in range(int(frame_count)): ret, frame = cap.read() if ret == False: break frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame_rgb = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame_rgb) # Calcular la marca temporal del frame actual (en milisegundos) frame_timestamp_ms = int(1000 * frame_index / fps) # Obtener los resultados del clasificador classifier_result = classifier.classify_for_video(frame_rgb, frame_timestamp_ms) print(f"classifier_result: {classifier_result}") y_aux = 0 for classification in classifier_result.classifications: for category in classification.categories: print(f"Clase: {category.category_name}, Puntaje: {category.score * 100}") cv2.putText(frame, f"{category.category_name}: {category.score * 100:.2f}%", (10, 20 + y_aux), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2) y_aux += 20 cv2.imshow('Video', frame) if cv2.waitKey(1) & 0xFF == 27: break cap.release() cv2.destroyAllWindows()
Vamos a recorrer cada uno de los fotogramas del video, y como lo podemos ver en la línea 24 y 25 también tendremos que transformar cada frame de BGR a RGB y lo convertimos en un objeto mediapipe.Image.
Lo nuevo viene en la línea 28, con el cálculo de frame_timestamp_ms, que es una marca temporal en milisegundos. Esta irá incrementando conforme pase cada fotograma, esto le ayudará a mediapipe a manejar correctamente el flujo de video, de tal forma que los datos estén correctamente ordenados. (Puedes echarle un vistazo a este link para más información).
Entonces en la línea 28, tendemos que pasar al detector no solo la imagen, sino también la marca de tiempo calculada.
Una vez terminado el programa, probemos el modelo con un video:


En la imagen de la izquierda tenemos que se ha clasificado la imagen como cheetah, mientras que en la imagen de la derecha tenemos un mayor porcentaje para plate.
A continuación tenemos el programa completo sobre la Image Clasification en un video:
import cv2 import mediapipe as mp from mediapipe.tasks.python import vision from mediapipe.tasks.python import BaseOptions # Especificar la configuración del clasificador options = vision.ImageClassifierOptions( base_options=BaseOptions(model_asset_path="/path/efficientnet_lite_float32.tflite"), max_results=5, score_threshold=0.01, running_mode=vision.RunningMode.VIDEO) classifier = vision.ImageClassifier.create_from_options(options) # Leer el video de entrada cap = cv2.VideoCapture("path/video.mp4") frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT) fps = cap.get(cv2.CAP_PROP_FPS) for frame_index in range(int(frame_count)): ret, frame = cap.read() if ret == False: break frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame_rgb = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame_rgb) # Calcular la marca temporal del frame actual (en milisegundos) frame_timestamp_ms = int(1000 * frame_index / fps) # Obtener los resultados del clasificador classifier_result = classifier.classify_for_video(frame_rgb, frame_timestamp_ms) print(f"classifier_result: {classifier_result}") y_aux = 0 for classification in classifier_result.classifications: for category in classification.categories: print(f"Clase: {category.category_name}, Puntaje: {category.score * 100}") cv2.putText(frame, f"{category.category_name}: {category.score * 100:.2f}%", (10, 20 + y_aux), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2) y_aux += 20 cv2.imshow('Video', frame) if cv2.waitKey(1) & 0xFF == 27: break cap.release() cv2.destroyAllWindows()
¿Cómo usar la image classification en video streaming con Mediapipe – Python?
Finalmente veremos como podemos aplicar la clasificación de imágenes sobre un video streaming, o video en directo. Para ello tendremos que hacerle unas cuantas modificaciones al programa anterior, veamos:
import cv2 import mediapipe as mp from mediapipe.tasks.python import vision from mediapipe.tasks.python import BaseOptions import time classification_result_list = [] # Función callback para procesar los resultados de la clasificación def classification_callback(result, output_image, timestamp_ms): classification_result_list.append(result)
En la línea 7 definiremos una lista vacía que nos ayudará más adelante a almacenar los resultados de la clasificación. Mientras que en la línea 10, tenemos la definicion de una función callback, la cual nos ayudará a procesar los mismos. Esta se llama automáticamente cada vez que el clasificador produce un nuevo resultado.
# Especificar la configuración del clasificador options = vision.ImageClassifierOptions( base_options=BaseOptions(model_asset_path="/path/efficientnet_lite2_float32.tflite"), max_results=5, score_threshold=0.01, running_mode=vision.RunningMode.LIVE_STREAM, result_callback=classification_callback) classifier = vision.ImageClassifier.create_from_options(options)
En el apartado de opciones de configuración, tendremos que definir vision.RunningMode.LIVE_STREAM para el running_mode, y en classification_callback especificamos la función que habíamos creado recientemente, esta nos permitirá manejar los resultados de las predicciones.
# Leer el video de entrada cap = cv2.VideoCapture(0) while True: ret, frame = cap.read() if ret == False: break frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame_rgb = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame_rgb) # Obtener los resultados del clasificador classifier_result = classifier.classify_async(frame_rgb, time.time_ns() // 1_000_000) y_aux = 0 if classification_result_list: for classification in classification_result_list[0].classifications: for category in classification.categories: print(f"Clase: {category.category_name}, Puntaje: {category.score * 100}") cv2.putText(frame, f"{category.category_name}: {category.score * 100:.2f}%", (10, 20 + y_aux), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (253, 81, 139), 2) y_aux += 20 cv2.imshow('Video', frame) classification_result_list.clear() if cv2.waitKey(1) & 0xFF == 27: break cap.release() cv2.destroyAllWindows()
En la línea 34 vamos a llamar al modelo para realizar las predicciones. Para ello usaremos classifier.classify_async, al cual tendremos que darle el frame que va a analizar junto con la marca de tiempo en milisegundos.
IMPORTANTE: La descripción en cuanto a la clasificación de imágenes en un videostreming nos dice lo siguiente: «El método classify_async
está diseñado para procesar datos de transmisión en vivo, como la entrada de una cámara. Para reducir la latencia general, el clasificador de imágenes puede descartar las imágenes de entrada si es necesario. En otras palabras, no se garantiza que haya una salida por cada imagen de entrada». (Fuente)
import cv2 import mediapipe as mp from mediapipe.tasks.python import vision from mediapipe.tasks.python import BaseOptions import time classification_result_list = [] # Función callback para procesar los resultados de la clasificación def classification_callback(result, output_image, timestamp_ms): classification_result_list.append(result) # Especificar la configuración del clasificador options = vision.ImageClassifierOptions( base_options=BaseOptions(model_asset_path="./Modelos/efficientnet_lite2_float32.tflite"), max_results=5, score_threshold=0.01, running_mode=vision.RunningMode.LIVE_STREAM, result_callback=classification_callback) classifier = vision.ImageClassifier.create_from_options(options) # Leer el video de entrada cap = cv2.VideoCapture(0) while True: ret, frame = cap.read() if ret == False: break frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame_rgb = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame_rgb) # Obtener los resultados del clasificador classifier_result = classifier.classify_async(frame_rgb, time.time_ns() // 1_000_000) y_aux = 0 if classification_result_list: for classification in classification_result_list[0].classifications: for category in classification.categories: print(f"Clase: {category.category_name}, Puntaje: {category.score * 100}") cv2.putText(frame, f"{category.category_name}: {category.score * 100:.2f}%", (10, 20 + y_aux), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (253, 81, 139), 2) y_aux += 20 cv2.imshow('Video', frame) classification_result_list.clear() if cv2.waitKey(1) & 0xFF == 27: break cap.release() cv2.destroyAllWindows()
Una vez concluido el programa, podremos hacer pruebas sobre este:


En la imagen de la izquierda podemos ver que con un 32.30% se ha predicho un avión de pasajeros, mientras que en la imagen de la derecha tenemos a una abeja con el 19.23%.
Y bien, hemos llegado al final. ¡Espero que les haya gustado!. Nos vemos en un siguiente post. ????
Referencias
- https://ai.google.dev/edge/mediapipe/solutions/vision/image_classifier
- https://github.com/google-ai-edge/mediapipe/blob/master/mediapipe/tasks/python/vision/image_classifier.py
- https://ai.google.dev/edge/mediapipe/solutions/vision/image_classifier/python
- https://colab.research.google.com/github/googlesamples/mediapipe/blob/main/examples/image_classification/python/image_classifier.ipynb#scrollTo=Iy4r2_ePylIa
- https://github.com/google-ai-edge/mediapipe-samples/blob/main/examples/image_classification/raspberry_pi/classify.py