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

La tarea Segmentador de imágenes de MediaPipe te permite dividir las imágenes en regiones según las categorías predefinidas para aplicar efectos visuales, como el desenfoque del fondo. En estas instrucciones, se muestra cómo usar el Segmentador de imágenes con el lenguaje Python. Para obtener más información sobre las funciones, 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 Segmentador 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 propia aplicación de segmentación de imágenes. Puedes ver, ejecutar y editar el código de ejemplo del Segmentador de imágenes solo con tu navegador web.

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 para usar el Segmentador 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. Puedes revisar el código fuente para este ejemplo en GitHub.

Paquetes

La tarea del Segmentador de imágenes de MediaPipe requiere el paquete mediapipe. Puedes instalar las dependencias necesarias con el siguiente comando:

$ python -m pip install mediapipe

Importaciones

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

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

Modelo

La tarea Segmentador 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 Segmentador de imágenes, consulta la sección de descripción general de la tarea Modelos.

Selecciona y descarga el modelo. Luego, almacénalo en el directorio de tu proyecto:

model_path = '/absolute/path/to/model.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

La tarea del Segmentador de imágenes de MediaPipe usa la función create_from_options para configurarla. La función create_from_options acepta valores para controlar las opciones de configuración. Para obtener más información sobre la configuración de tareas, consulta Opciones de configuración.

En estos ejemplos, también se incluyen las variaciones de la construcción de tareas para imágenes, archivos de video y transmisiones de video en vivo.

De imagen

BaseOptions = mp.tasks.BaseOptions
ImageSegmenter = mp.tasks.vision.ImageSegmenter
ImageSegmenterOptions = mp.tasks.vision.ImageSegmenterOptions
VisionRunningMode = mp.tasks.vision.RunningMode

# Create a image segmenter instance with the image mode:
options = ImageSegmenterOptions(
    base_options=BaseOptions(model_asset_path='/path/to/model.task'),
    running_mode=VisionRunningMode.IMAGE,
    output_category_mask=True)
with ImageSegmenter.create_from_options(options) as segmenter:
  

Video

BaseOptions = mp.tasks.BaseOptions
ImageSegmenter = mp.tasks.vision.ImageSegmenter
ImageSegmenterOptions = mp.tasks.vision.ImageSegmenterOptions
VisionRunningMode = mp.tasks.vision.RunningMode

# Create a image segmenter instance with the video mode:
options = ImageSegmenterOptions(
    base_options=BaseOptions(model_asset_path='/path/to/model.task'),
    running_mode=VisionRunningMode.VIDEO,
    output_category_mask=True)
with ImageSegmenter.create_from_options(options) as segmenter:
  

Transmisión en vivo

BaseOptions = mp.tasks.BaseOptions
ImageSegmenter = mp.tasks.vision.ImageSegmenter
ImageSegmenterOptions = mp.tasks.vision.ImageSegmenterOptions
VisionRunningMode = mp.tasks.vision.RunningMode

# Create a image segmenter instance with the live stream mode:
def print_result(result: List[Image], output_image: Image, timestamp_ms: int):
    print('segmented masks size: {}'.format(len(result)))

options = ImageSegmenterOptions(
    base_options=BaseOptions(model_asset_path='/path/to/model.task'),
    running_mode=VisionRunningMode.VIDEO,
    output_category_mask=True)
with ImageSegmenter.create_from_options(options) as segmenter:
    

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
output_category_mask Si se configura en True, el resultado incluye una máscara de segmentación como una imagen uint8, en la que cada valor de píxel indica el valor de la categoría ganadora. {True, False} False
output_confidence_masks Si se configura como True, el resultado incluye una máscara de segmentación como una imagen de valor flotante, en la que cada valor flotante representa el mapa de puntuación de confianza de la categoría. {True, False} True
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
result_callback Configura el objeto de escucha de resultados para recibir los resultados de la segmentación de forma asíncrona cuando el segmentador de imágenes está en modo de transmisión en vivo. Solo se puede usar cuando el modo de ejecución está configurado en LIVE_STREAM N/A N/A

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.

De imagen

# 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

# Use OpenCV’s VideoCapture to load the input video.
# Load the frame rate of the video using OpenCV’s CV_CAP_PROP_FPS
# You need the frame rate 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

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

Si deseas ver un ejemplo de código que muestra la preparación de datos para el Segmentador de imágenes, consulta el ejemplo de código.

Ejecuta la tarea

El segmentador de imágenes usa las funciones segment, segment_for_video y segment_async para activar las inferencias. Para la segmentación de imágenes, esto implica el procesamiento previo de los datos de entrada, la ejecución del modelo de segmentación y el procesamiento posterior de las salidas del modelo sin procesar a las máscaras segmentadas.

En los siguientes ejemplos de código, se muestra cómo ejecutar el procesamiento con el modelo de tareas.

De imagen

# Perform image segmentation on the provided single image.
# The image segmenter must be created with the image mode.
segmented_masks = segmenter.segment(mp_image)
    

Video

# Perform image segmentation on the provided single image.
# The image segmenter must be created with the video mode.
segmented_masks = segmenter.segment_for_video(mp_image, frame_timestamp_ms)
    

Transmisión en vivo

# Send live image data to perform image segmentation.
# The results are accessible via the `result_callback` provided in
# the `ImageSegmenterOptions` object.
# The image segmenter must be created with the live stream mode.
segmenter.segment_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 la marca de tiempo del fotograma de entrada a la tarea Segmentador de imágenes.
  • Cuando se ejecuta en la imagen o el modelo de video, la tarea Segmentador de imágenes bloqueará el subproceso actual hasta que termine de procesar la imagen o el marco de entrada.

Para ver un ejemplo más completo de cómo ejecutar inferencias del Segmentador de imágenes, consulta el ejemplo de código.

Cómo controlar y mostrar los resultados

El Segmentador de imágenes genera una lista de datos de Image. Si output_type es CATEGORY_MASK, el resultado es una lista que contiene una sola máscara segmentada como una imagen uint8. El píxel indica el índice de categoría reconocido de la imagen de entrada. Si output_type es CONFIDENCE_MASK, el resultado es un vector con el tamaño del número de categoría. Cada máscara segmentada es una imagen flotante dentro del rango [0,1], que representa la puntuación de confianza del píxel que pertenece a la categoría.

En las siguientes secciones, se muestran ejemplos de los datos de salida de esta tarea:

Nivel de confianza de la categoría

En las siguientes imágenes, se muestra una visualización del resultado de la tarea para una máscara de confianza de categoría. El resultado de la máscara de confianza contiene valores flotantes entre [0, 1].

Resultado original de la máscara de confianza de la imagen y la categoría. Imagen de origen del conjunto de datos Pascal VOC 2012.

Valor de la categoría

En las siguientes imágenes, se muestra una visualización del resultado de la tarea para una máscara de valor de categoría. El rango de máscara de categoría es [0, 255], y cada valor de píxel representa el índice de la categoría ganadora del resultado del modelo. El índice de la categoría ganadora tiene la puntuación más alta entre las categorías que el modelo puede reconocer.

Resultado original de la imagen y la máscara de categoría. Imagen de origen del conjunto de datos Pascal VOC 2012.