Guía de clasificación de imágenes para Python

La tarea Clasificador de imágenes de MediaPipe te permite realizar la clasificación de imágenes. Puedes usar esta tarea para identificar lo que representa una imagen entre un conjunto de categorías definidas en el momento del entrenamiento. En estas instrucciones, se muestra cómo usar el clasificador de imágenes con Python.

Puedes ver esta tarea en acción en la demostración web. Para obtener más información sobre las capacidades, los modelos y las opciones de configuración de esta tarea, consulta la Descripción general.

Ejemplo de código

El código de ejemplo del clasificador de imágenes proporciona una implementación completa de esta tarea en Python para tu referencia. Este código te ayuda a probar esta tarea y comenzar a compilar tu propio clasificador de imágenes. Puedes ver, ejecutar y editar el código de ejemplo del clasificador de imágenes solo con tu navegador web.

Si quieres implementar el clasificador de imágenes para Raspberry Pi, consulta la app de ejemplo de Raspberry Pi.

Configuración

En esta sección, se describen los pasos clave para configurar tu entorno de desarrollo y los proyectos de código específicamente a fin de usar el clasificador de imágenes. Si deseas obtener información general sobre cómo configurar tu entorno de desarrollo para usar tareas de MediaPipe, incluidos los requisitos de versión de la plataforma, consulta la Guía de configuración de Python.

Paquetes

La tarea del clasificador de imágenes en el paquete pip de mediapipe Puedes instalar la dependencia con lo siguiente:

$ python -m pip install mediapipe

Importaciones

Importa las siguientes clases para acceder a las funciones de tareas del clasificador de imágenes:

import mediapipe as mp
from mediapipe.tasks import python
from mediapipe.tasks.python import vision

Modelo

La tarea del clasificador de imágenes de MediaPipe requiere un modelo entrenado que sea compatible con esta tarea. Si deseas obtener más información sobre los modelos entrenados disponibles para el clasificador de imágenes, consulta la descripción general de la tarea en la sección Modelos.

Selecciona y descarga un modelo, y almacénalo en un directorio local. Puedes usar el modelo recomendado EfficientNet-Lite0.

model_path = '/absolute/path/to/efficientnet_lite0_int8_2.tflite'

Especifica la ruta del modelo dentro del parámetro Nombre del modelo, como se muestra a continuación:

base_options = BaseOptions(model_asset_path=model_path)

Crea la tarea

Usa la función create_from_options para crear la tarea. La función create_from_options acepta opciones de configuración, incluidos el modo de ejecución, la configuración regional de los nombres visibles, la cantidad máxima de resultados, el umbral de confianza, la lista de categorías permitidas y la lista de bloqueo. Para obtener más información sobre las opciones de configuración, consulta Descripción general de la configuración.

La tarea del clasificador de imágenes admite 3 tipos de datos de entrada: imágenes fijas, archivos de video y transmisiones de video en vivo. Elige la pestaña correspondiente a tu tipo de datos de entrada para ver cómo crear la tarea y ejecutar la inferencia.

De imagen

import mediapipe as mp

BaseOptions = mp.tasks.BaseOptions
ImageClassifier = mp.tasks.vision.ImageClassifier
ImageClassifierOptions = mp.tasks.vision.ImageClassifierOptions
VisionRunningMode = mp.tasks.vision.RunningMode

options = ImageClassifierOptions(
    base_options=BaseOptions(model_asset_path='/path/to/model.tflite'),
    max_results=5,
    running_mode=VisionRunningMode.IMAGE)

with ImageClassifier.create_from_options(options) as classifier:
  # The classifier is initialized. Use it here.
  # ...
    

Video

import mediapipe as mp

BaseOptions = mp.tasks.BaseOptions
ImageClassifier = mp.tasks.vision.ImageClassifier
ImageClassifierOptions = mp.tasks.vision.ImageClassifierOptions
VisionRunningMode = mp.tasks.vision.RunningMode

options = ImageClassifierOptions(
    base_options=BaseOptions(model_asset_path='/path/to/model.tflite'),
    max_results=5,
    running_mode=VisionRunningMode.VIDEO)

with ImageClassifier.create_from_options(options) as classifier:
  # The classifier is initialized. Use it here.
  # ...
    

Transmisión en vivo

import mediapipe as mp

BaseOptions = mp.tasks.BaseOptions
ImageClassifierResult = mp.tasks.vision.ImageClassifier.ImageClassifierResult
ImageClassifier = mp.tasks.vision.ImageClassifier
ImageClassifierOptions = mp.tasks.vision.ImageClassifierOptions
VisionRunningMode = mp.tasks.vision.RunningMode

def print_result(result: ImageClassifierResult, output_image: mp.Image, timestamp_ms: int):
    print('ImageClassifierResult result: {}'.format(result))

options = ImageClassifierOptions(
    base_options=BaseOptions(model_asset_path='/path/to/model.tflite'),
    running_mode=VisionRunningMode.LIVE_STREAM,
    max_results=5,
    result_callback=print_result)

with ImageClassifier.create_from_options(options) as classifier:
  # The classifier is initialized. Use it here.
  # ...
    

Si deseas ver un ejemplo completo de cómo crear un clasificador de imágenes para usar con una imagen, consulta el ejemplo de código.

Opciones de configuración

Esta tarea incluye las siguientes opciones de configuración para aplicaciones de Python:

Nombre de la opción Descripción Rango de valores Valor predeterminado
running_mode Establece el modo de ejecución para la tarea. Existen tres modos:

IMAGE: El modo para entradas de una sola imagen.

VIDEO: Es el modo para los fotogramas decodificados de un video.

LIVE_STREAM: Es el modo para una transmisión en vivo de datos de entrada, como los que provienen de una cámara. En este modo, se debe llamar a resultListener para configurar un objeto de escucha a fin de recibir resultados de forma asíncrona.
{IMAGE, VIDEO, LIVE_STREAM} IMAGE
display_names_locale Configura el idioma de las etiquetas que se usarán para los nombres visibles proporcionados en los metadatos del modelo de la tarea, si están disponibles. El valor predeterminado es en para inglés. Puedes agregar etiquetas localizadas a los metadatos de un modelo personalizado con la API de Metadata Writer de TensorFlow Lite Código de configuración regional en
max_results Configura la cantidad máxima opcional de resultados de clasificación con puntuación más alta que se mostrarán. Si es inferior a 0, se devolverán todos los resultados disponibles. Cualquier número positivo -1
score_threshold Establece el umbral de puntuación de predicción que anula el proporcionado en los metadatos del modelo (si corresponde). Se rechazarán los resultados inferiores a este valor. Cualquier número de punto flotante No establecida
category_allowlist Establece la lista opcional de nombres de categorías permitidas. Si no está vacío, se filtrarán los resultados de clasificación cuyo nombre de categoría no esté en este conjunto. Los nombres de categorías duplicados o desconocidos se ignoran. Esta opción es mutuamente excluyente con category_denylist, y el uso de ambos da como resultado un error. Cualquier cadena No establecida
category_denylist Establece la lista opcional de nombres de categorías que no están permitidas. Si no está vacío, se filtrarán los resultados de clasificación cuyo nombre de categoría se encuentre en este conjunto. Los nombres de categorías duplicados o desconocidos se ignoran. Esta opción es mutuamente exclusiva con category_allowlist, y el uso de ambas da como resultado un error. Cualquier cadena No establecida
result_callback Configura el objeto de escucha de resultados para recibir los resultados de clasificación de forma asíncrona cuando el clasificador de imágenes está en el modo de transmisión en vivo. Solo se puede usar cuando el modo de ejecución está configurado en LIVE_STREAM N/A No establecida

Preparar los datos

Prepara la entrada como un archivo de imagen o un array de NumPy y, luego, conviértela en un objeto mediapipe.Image. Si la entrada es un archivo de video o una transmisión en vivo desde una cámara web, puedes usar una biblioteca externa, como OpenCV, para cargar los fotogramas de entrada como arrays NumPy.

En los siguientes ejemplos, se explica y se muestra cómo preparar los datos para su procesamiento para cada uno de los tipos de datos disponibles.

De imagen

import mediapipe as mp

# Load the input image from an image file.
mp_image = mp.Image.create_from_file('/path/to/image')

# Load the input image from a numpy array.
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=numpy_image)
    

Video

import mediapipe as mp

# Use OpenCV’s VideoCapture to load the input video.

# Load the frame rate of the video using OpenCV’s CV_CAP_PROP_FPS
# You’ll need it to calculate the timestamp for each frame.

# Loop through each frame in the video using VideoCapture#read()

# Convert the frame received from OpenCV to a MediaPipe’s Image object.
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=numpy_frame_from_opencv)
    

Transmisión en vivo

import mediapipe as mp

# Use OpenCV’s VideoCapture to start capturing from the webcam.

# Create a loop to read the latest frame from the camera using VideoCapture#read()

# Convert the frame received from OpenCV to a MediaPipe’s Image object.
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=numpy_frame_from_opencv)
    

Ejecuta la tarea

Puedes llamar a la función de clasificación que corresponde a tu modo de ejecución para activar las inferencias. La API de Image Classifier mostrará las categorías posibles para el objeto dentro de la imagen o el marco de entrada.

De imagen

# Perform image classification on the provided single image.
classification_result = classifier.classify(mp_image)
    

Video

# Calculate the timestamp of the current frame
frame_timestamp_ms = 1000 * frame_index / video_file_fps

# Perform image classification on the video frame.
classification_result = classifier.classify_for_video(mp_image, frame_timestamp_ms)
    

Transmisión en vivo


# Send the latest frame to perform image classification.
# Results are sent to the `result_callback` provided in the `ImageClassifierOptions`.
classifier.classify_async(mp_image, frame_timestamp_ms)
    

Ten en cuenta lo siguiente:

  • Cuando se ejecuta en el modo de video o de transmisión en vivo, también debes proporcionar a la tarea Clasificador de imágenes la marca de tiempo del fotograma de entrada.
  • Cuando se ejecuta en la imagen o el modelo de video, la tarea Clasificador de imágenes bloqueará el subproceso actual hasta que termine de procesar la imagen o el marco de entrada.
  • Cuando se ejecuta en el modo de transmisión en vivo, la tarea Clasificador de imágenes no bloquea el subproceso actual, pero muestra el resultado de inmediato. Invocará su objeto de escucha de resultados con el resultado de la clasificación cada vez que termine de procesar un marco de entrada. Si se llama a la función classifyAsync cuando la tarea del clasificador de imágenes está ocupada procesando otro marco, la tarea ignora el nuevo marco de entrada.

Si deseas ver un ejemplo completo de cómo crear un clasificador de imágenes para usar con una imagen, consulta el ejemplo de código.

Cómo controlar y mostrar los resultados

Cuando se ejecuta la inferencia, la tarea del clasificador de imágenes muestra un objeto ImageClassifierResult que contiene la lista de categorías posibles para los objetos dentro de la imagen o el marco de entrada.

A continuación, se muestra un ejemplo de los datos de salida de esta tarea:

ImageClassifierResult:
 Classifications #0 (single classification head):
  head index: 0
  category #0:
   category name: "/m/01bwb9"
   display name: "Passer domesticus"
   score: 0.91406
   index: 671
  category #1:
   category name: "/m/01bwbt"
   display name: "Passer montanus"
   score: 0.00391
   index: 670

Este resultado se obtiene ejecutando el Clasificador de aves en:

En el código de ejemplo del clasificador de imágenes, se demuestra cómo mostrar los resultados de clasificación que muestra la tarea. Consulta el ejemplo de código para obtener más información.