Segmentación múltiple con MediaPipe en Python
¡Hola, hola Omesitos!.
Seguimos analizando los modelos de Mediapipe Image Segmentation. En el anterior post vimos como usar el Hair segmentation model, y en esta ocasión veremos como segmentar múltiples clases con: Multi-class selfie segmentation model. Así que… ¡Vamos a por ello!.
CONTENIDO
- ¿Qué es la segmentación selfie multiclase que ofrece MediaPipe?
- ¿Cómo usar el Multi-class Selfie Segmentation Model de MediaPipe?
- ¿Cómo usar Multi-class Selfie Segmentation Model de MediaPipe en una imagen con Python?
- ¿Cómo usar Multi-class Selfie Segmentation Model de MediaPipe en videos con Python?
- ¿Cómo usar Multi-class Selfie Segmentation Model de MediaPipe en videostream con Python?
¿Qué es la segmentación selfie multiclase que ofrece MediaPipe?
El Multi-class Selfie Segmentation Model es un modelo de machine learning para visión artificial diseñado para segmentar múltiples partes de una imagen o video que contienen a una persona, identificando diversas regiones específicas. Este modelo tiene la capacidad de distinguir entre varias categorías dentro de una misma imagen, generando una máscara donde cada píxel es etiquetado según su categoría.
Las categorías que el modelo puede segmentar son las siguientes:
- 0: Fondo
- 1: Cabello
- 2: Piel del cuerpo
- 3: Piel del rostro
- 4: Ropa
- 5: Otros (accesorios)
Como podemos observar en la lista, cada categoría tiene un número asociado, el cual estará presente en la máscara resultante generada por el modelo. Ya lo veremos más adelante. Por ahora veamos un ejemplo de como el modelo puede segmentar a una imagen:
A la izquierda tenemos la imagen de entrada, mientras que a la derecha, la imagen segmentada. Podemos ver que el fondo se ha coloreado de azul, el cabello de verde, la piel de la cara de celeste, mientras que la piel del cuerpo de rojo, la ropa de rosa y los accesorios, en este caso los lentes de amarillo. Y esto es precisamente lo que haremos en este tutorial, vamos a identificar las diferentes regiones y colorearlas.
Pero antes de ir a por más, debemos tomar en cuenta que el modelo necesita como entrada una imagen de 256×256 pixeles. Es decir que necesitaremos una imagen cuadrada, con una relación de aspecto de 1. Podríamos probar con otras relaciones de aspecto, sin embargo es importante respetar los requerimientos que nos especifica la documentación para obtener los mejores resultados.
Para descargar el modelo, simplemente vamos a tener que dar clic sobre su nombre.
En la documentación también podremos encontrar el apartado de Model Card, donde podremos obtener más información del modelo como: uso, limitaciones, métricas, etc.
¿Cómo usar el Multi-class Selfie Segmentation Model de MediaPipe?
Para empezar a usar las soluciones y modelos de Mediapipe, debemos instalar este framework, así que para ello nos ayudaremos de pip:
pip install mediapipe
Puedes revisar los modelos disponibles para la segmentación de imagen en la documentación de Mediapipe:
Al igual que en los anteriores blog posts, debemos tener en cuenta que los modelos de segmentación nos van a devolver predicciones sobre cada pixel de la imagen de entrada. En sí debemos tener en claro las siguientes salidas:
CATEGORY_MASK: Lista que contiene una máscara segmentada en formato uint8. Cada valor de píxel indica que es parte de una categoría de segmento específica.
CONFIDENCE_MASK: Lista de canales que contiene una máscara segmentada con valores de píxel en formato float32. Cada valor de píxel indica el nivel de confianza de que es parte de una categoría específica.
¿Cómo usar Multi-class Selfie Segmentation Model de MediaPipe en una imagen con 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 import numpy as np
Lo primero que haremos es importar mediapipe, luego de mediapipe.task.python importaremos vision, ya que estaremos empleando una tarea de visión por computador, la cual nos permitirá manejar imágenes o videos. Luego importaremos BaseOptions que básicamente se encargará de leer el modelo que se usará, así como configurarlo. Importamos OpenCV y finalmente Numpy.
# Especificar la configuración del ImageSegmenter
options = vision.ImageSegmenterOptions(
base_options=BaseOptions(model_asset_path="./path/selfie_multiclass_256x256.tflite"),
output_category_mask=True,
running_mode=vision.RunningMode.IMAGE
)
segmenter = vision.ImageSegmenter.create_from_options(options)
Y precisamente vamos a empezar con las opciones de configuración a través de vision.ImageSegmenterOptions. En primer lugar vamos a establecer el path del modelo que estaremos usando, en este caso selfie_multiclass_256x256.tflite. Luego tenemos a output_category_mask que estableceremos como True, y nos permitirá obtener una máscara de segmentación como una imagen uint8, donde cada valor de píxel indica la categoría ganadora.
Finalmente en running_mode vamos a ubicar vision.RunningMode.IMAGE, debido a que vamos a estar usando una imagen como entrada.
Una vez que tenemos las opciones de configuración, vamos a pasarlas al segmentador, para ello usamos vision.ImageSegmenter.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("./Inputs/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, transformamos 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.
# Aplicar el segmentador
segmentation_result = segmenter.segment(image_rgb)
print("segmentation_result", segmentation_result)
Para empezar con la segmentación tendremos que usar segmenter.segment, y especificar la imagen que deseamos analizar. A continuación extraemos las salidas del segmentador:
category_mask = segmentation_result.category_mask confidence_masks = segmentation_result.confidence_masks
segmentation_result.category_mask, nos permitirá extraer la máscara de categorías (category_mask) del resultado de la segmentación. Esta máscara es una imagen en formato uint8, donde cada valor de píxel indica la categoría de segmento a la que pertenece.
Para este caso, el modelo de segmentación multiclase nos permitirá extraer una matriz compuesta por valores comprendidos entre 0 y 6, de acuerdo a las categorías de segmentación.
segmentation_result.confidence_masks, nos devolverá 6 máscaras de confianza correspondientes a cada una de las categorías de segmentación que maneja el modelo.
A continuación vamos a pasar todas estas matrices a arrays de numpy, para poder visualizarlas y manipularlas.
# Convertir las máscaras en arrays de numpy # CONFIDENCE MASKS confidence_masks_np_bg = confidence_masks[0].numpy_view() confidence_masks_np_hair = confidence_masks[1].numpy_view() confidence_masks_np_body_skin = confidence_masks[2].numpy_view() confidence_masks_np_face_skin = confidence_masks[3].numpy_view() confidence_masks_np_clothes = confidence_masks[4].numpy_view() confidence_masks_np_others = confidence_masks[5].numpy_view() print(confidence_masks_np_bg.dtype) print(confidence_masks_np_hair.dtype) print(confidence_masks_np_body_skin.dtype) print(confidence_masks_np_face_skin.dtype) print(confidence_masks_np_clothes.dtype) print(confidence_masks_np_others.dtype) # CATEGORY MASK category_mask_np = category_mask.numpy_view() print(category_mask_np.dtype, category_mask_np.shape) print(np.unique(category_mask_np))
Como ejercicio de visualización en las líneas 28 a la 33, podemos transformar cada una de las máscaras de confianza a un array de numpy. Además podremos extraer el tipo de dato de cada una de ellos, que será float 32, y nos dará valores entre 0 y 1.
Si visualizamos todas estas máscaras de confianza, podremos ver algo como lo siguiente:
En el ejemplo, a la izquierda tenemos la imagen de entrada. Mientras que a la derecha tenemos las 6 máscaras de confianza, cada una para cada categoría. Colores más claros estarán más cerca del 1, por lo tanto habrá más confianza de que en esa área está cierta categoría. Mientras que colores más oscuros, o más cercanos a 0, representan la ausencia de dicha categoría.
De hecho, si empezamos por la sección superior izquierda, a la derecha de la imagen de entrada tenemos: máscara de confianza del fondo, cabello, piel del cuerpo. Mientras que en la sección inferior tenemos a: máscara de confianza de la piel de la cara, ropa y los accesorios, en este caso, los lentes.
En la línea 43 tenemos la transformación del category_mask a array de numpy. En la línea 44 podremos imprimir el tipo de dato (int8), y el ancho y alto que corresponderán a los datos de la imagen de entrada.
Mientras que en la línea 45 podremos imprimir los valores únicos que contiene category_mask, esto nos permitirá conocer con qué valores trabajar. Al imprimirlos obtendremos: [0 1 2 3 4 5 6]. Esto quiere decir que toda la imagen contiene unicamente esos valores. De hecho si quisiéramos visualizar la máscara, no podríamos distinguir entre esos valores, veríamos todo en color negro. Por esta razón tendremos que hacer una transformación de color a cada uno de los segmentos.
# Otorgar color a cada segmento #print(category_mask_np.shape) category_mask_bgr = cv2.cvtColor(category_mask_np, cv2.COLOR_GRAY2BGR) #print(category_mask_bgr.shape) category_mask_bgr[np.where(category_mask_np == 0)] = (255, 0, 0) # Azul: Fondo category_mask_bgr[np.where(category_mask_np == 1)] = (0, 255, 0) # Verde: Cabello category_mask_bgr[np.where(category_mask_np == 2)] = (0, 0, 255) # Rojo: Piel del cuerpo category_mask_bgr[np.where(category_mask_np == 3)] = (255, 255, 0) # Cian: Piel de la cara category_mask_bgr[np.where(category_mask_np == 4)] = (255, 0, 255) # Rosa: Ropa category_mask_bgr[np.where(category_mask_np == 5)] = (0, 255, 255) # Amarillo: Otros
En la línea 49 convertiremos la imagen de un solo canal (category_mask_np), a BGR para poder darle color. De hecho nos podemos ayudar de las impresiones de las líneas 48 y 50 para visualizar el cambio de 1 a los 3 canales.
De las líneas 52 a la 57, estableceremos un color para cada una de las categorías. Así tenemos que el fondo se visualizará en azul, el cabello en verde, la piel del cuerpo en rojo, la piel de la cara en celeste, la ropa en rosa y los accesorios en amarillo.
# Visualización
cv2.imshow("Image", image)
#cv2.imshow("confidence_masks_np_bg", confidence_masks_np_bg)
#cv2.imshow("confidence_masks_np_hair", confidence_masks_np_hair)
#cv2.imshow("confidence_masks_np_body_skin", confidence_masks_np_body_skin)
#cv2.imshow("confidence_masks_np_face_skin", confidence_masks_np_face_skin)
#cv2.imshow("confidence_masks_np_clothes", confidence_masks_np_clothes)
#cv2.imshow("confidence_masks_np_others", confidence_masks_np_others)
#cv2.imshow("category_mask_np", category_mask_np)
cv2.imshow("category_mask_bgr", category_mask_bgr)
cv2.waitKey(0)
cv2.destroyAllWindows()
Finalmente desde la línea 60, visualizaremos los resultados. A pesar de que podemos visualizar cada una de las matrices o imágenes, nos quedaremos con la imagen de entrada y la imagen resultante. Veamos:
A continuación tenemos el programa completo:
import mediapipe as mp
from mediapipe.tasks.python import vision
from mediapipe.tasks.python import BaseOptions
import cv2
import numpy as np
# Especificar la configuración del ImageSegmenter
options = vision.ImageSegmenterOptions(
base_options=BaseOptions(model_asset_path="./path/selfie_multiclass_256x256.tflite"),
output_category_mask=True,
running_mode=vision.RunningMode.IMAGE)
segmenter = vision.ImageSegmenter.create_from_options(options)
# Leer la imagen de entrada
image = cv2.imread("./Inputs/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 segmentador
segmentation_result = segmenter.segment(image_rgb)
print("segmentation_result:", segmentation_result)
category_mask = segmentation_result.category_mask
confidence_masks = segmentation_result.confidence_masks
# Convertir las máscaras en arrays de numpy
# CONFIDENCE MASKS
confidence_masks_np_bg = confidence_masks[0].numpy_view()
confidence_masks_np_hair = confidence_masks[1].numpy_view()
confidence_masks_np_body_skin = confidence_masks[2].numpy_view()
confidence_masks_np_face_skin = confidence_masks[3].numpy_view()
confidence_masks_np_clothes = confidence_masks[4].numpy_view()
confidence_masks_np_others = confidence_masks[5].numpy_view()
print(confidence_masks_np_bg.dtype)
print(confidence_masks_np_hair.dtype)
print(confidence_masks_np_body_skin.dtype)
print(confidence_masks_np_face_skin.dtype)
print(confidence_masks_np_clothes.dtype)
print(confidence_masks_np_others.dtype)
# CATEGORY MASK
category_mask_np = category_mask.numpy_view()
print(category_mask_np.dtype, category_mask_np.shape)
print(np.unique(category_mask_np))
# Otorgar color a cada segmento
#print(category_mask_np.shape)
category_mask_bgr = cv2.cvtColor(category_mask_np, cv2.COLOR_GRAY2BGR)
#print(category_mask_bgr.shape)
category_mask_bgr[np.where(category_mask_np == 0)] = (255, 0, 0) # Azul: Fondo
category_mask_bgr[np.where(category_mask_np == 1)] = (0, 255, 0) # Verde: Cabello
category_mask_bgr[np.where(category_mask_np == 2)] = (0, 0, 255) # Rojo: Piel del cuerpo
category_mask_bgr[np.where(category_mask_np == 3)] = (255, 255, 0) # Cian: Piel de la cara
category_mask_bgr[np.where(category_mask_np == 4)] = (255, 0, 255) # Rosa: Ropa
category_mask_bgr[np.where(category_mask_np == 5)] = (0, 255, 255) # Amarillo: Otros
# Visualización
cv2.imshow("Image", image)
#cv2.imshow("confidence_masks_np_bg", confidence_masks_np_bg)
#cv2.imshow("confidence_masks_np_hair", confidence_masks_np_hair)
#cv2.imshow("confidence_masks_np_body_skin", confidence_masks_np_body_skin)
#cv2.imshow("confidence_masks_np_face_skin", confidence_masks_np_face_skin)
#cv2.imshow("confidence_masks_np_clothes", confidence_masks_np_clothes)
#cv2.imshow("confidence_masks_np_others", confidence_masks_np_others)
#cv2.imshow("category_mask_np", category_mask_np)
cv2.imshow("category_mask_bgr", category_mask_bgr)
cv2.waitKey(0)
cv2.destroyAllWindows()
¿Cómo usar Multi-class Selfie Segmentation Model de MediaPipe en videos con Python?
Debido a que el programa para la segmentación en video presenta similitudes con la segmentación sobre una imagen, procederemos a explicar sus diferencias, así que vamos a por ello.
import mediapipe as mp
from mediapipe.tasks.python import vision
from mediapipe.tasks.python import BaseOptions
import cv2
import numpy as np
# Especificar la configuración del ImageSegmenter
options = vision.ImageSegmenterOptions(
base_options=BaseOptions(model_asset_path="./path/selfie_multiclass_256x256.tflite"),
output_category_mask=True,
running_mode=vision.RunningMode.VIDEO
)
segmenter = vision.ImageSegmenter.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.
# 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 segmentador
segmentation_result = segmenter.segment_for_video(frame_rgb, frame_timestamp_ms)
#print("segmentation_result", segmentation_result)
# Extraer la máscara de categorias
category_mask = segmentation_result.category_mask
# Máscara a Numpy
category_mask_np = category_mask.numpy_view()
# Otorgar color a cada segmento
category_mask_bgr = cv2.cvtColor(category_mask_np, cv2.COLOR_GRAY2BGR)
category_mask_bgr[np.where(category_mask_np == 0)] = (255, 0, 0) # Azul: Fondo
category_mask_bgr[np.where(category_mask_np == 1)] = (0, 255, 0) # Verde: Cabello
category_mask_bgr[np.where(category_mask_np == 2)] = (0, 0, 255) # Rojo: Piel del cuerpo
category_mask_bgr[np.where(category_mask_np == 3)] = (255, 255, 0) # Cian: Piel de la cara
category_mask_bgr[np.where(category_mask_np == 4)] = (255, 0, 255) # Rosa: Ropa
category_mask_bgr[np.where(category_mask_np == 5)] = (0, 255, 255) # Amarillo: Otros
# Visualización de las imágenes
cv2.imshow("Frame", frame)
cv2.imshow("category_mask_bgr", category_mask_bgr)
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 25 y 26, 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 29, 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 32, tendemos que usar segmenter.segment_for_video y a este entregarle el fotograma y la marca de tiempo calculada.
Una vez terminado el programa, probemos el modelo con un video:
A continuación tenemos el programa completo de la segmentación multiclase para video:
import mediapipe as mp
from mediapipe.tasks.python import vision
from mediapipe.tasks.python import BaseOptions
import cv2
import numpy as np
# Especificar la configuración del ImageSegmenter
options = vision.ImageSegmenterOptions(
base_options=BaseOptions(model_asset_path="./path/selfie_multiclass_256x256.tflite"),
output_category_mask=True,
running_mode=vision.RunningMode.VIDEO
)
segmenter = vision.ImageSegmenter.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 segmentador
segmentation_result = segmenter.segment_for_video(frame_rgb, frame_timestamp_ms)
#print("segmentation_result", segmentation_result)
# Extraer la máscara de categorias
category_mask = segmentation_result.category_mask
# Máscara a Numpy
category_mask_np = category_mask.numpy_view()
# Otorgar color a cada segmento
category_mask_bgr = cv2.cvtColor(category_mask_np, cv2.COLOR_GRAY2BGR)
category_mask_bgr[np.where(category_mask_np == 0)] = (255, 0, 0) # Azul: Fondo
category_mask_bgr[np.where(category_mask_np == 1)] = (0, 255, 0) # Verde: Cabello
category_mask_bgr[np.where(category_mask_np == 2)] = (0, 0, 255) # Rojo: Piel del cuerpo
category_mask_bgr[np.where(category_mask_np == 3)] = (255, 255, 0) # Cian: Piel de la cara
category_mask_bgr[np.where(category_mask_np == 4)] = (255, 0, 255) # Rosa: Ropa
category_mask_bgr[np.where(category_mask_np == 5)] = (0, 255, 255) # Amarillo: Otros
# Visualización de las imágenes
cv2.imshow("Frame", frame)
cv2.imshow("category_mask_bgr", category_mask_bgr)
if cv2.waitKey(1) & 0xFF == 27:
break
cap.release()
cv2.destroyAllWindows()
¿Cómo usar Multi-class Selfie Segmentation Model de MediaPipe en videostream con Python?
Finalmente veremos como podemos aplicar la segmentación multiclase sobre un video streaming o video en directo. Para ello tendremos que hacer unas cuantas modificaciones al programa anterior, veamos:
import mediapipe as mp
from mediapipe.tasks.python import vision
from mediapipe.tasks.python import BaseOptions
import cv2
import numpy as np
import time
segmentation_result_list = []
# Función callback para procesar los resultados
def segmentation_callback(result, output_image, timestamp_ms):
segmentation_result_list.append(result)
En la línea 8 definiremos una lista vacía que nos ayudará más adelante a almacenar los resultados de la segmentación. Mientras que en la línea 11, tenemos la definición de una función callback, la cual nos ayudará a procesar los mismos. Esta se llama automáticamente cada vez que el segmentador produce un nuevo resultado.
# Especificar la configuración del ImageSegmenter
options = vision.ImageSegmenterOptions(
base_options=BaseOptions(model_asset_path="./path/selfie_multiclass_256x256.tflite"),
output_category_mask=True,
running_mode=vision.RunningMode.LIVE_STREAM,
result_callback=segmentation_callback)
segmenter = vision.ImageSegmenter.create_from_options(options)
Para el apartado de opciones de configuración, tendremos que definir vision.RunningMode.LIVE_STREAM para el running_mode, y en result_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 segmentador
segmenter.segment_async(frame_rgb, time.time_ns() // 1_000_000)
if len(segmentation_result_list) > 0:
#print(f"segmentation_result_list: {segmentation_result_list}")
segmentation_result = segmentation_result_list[0]
# Extraer la máscara de categorias
category_mask = segmentation_result.category_mask
# Máscara a Numpy
category_mask_np = category_mask.numpy_view()
# Otorgar color a cada segmento
category_mask_bgr = cv2.cvtColor(category_mask_np, cv2.COLOR_GRAY2BGR)
category_mask_bgr[np.where(category_mask_np == 0)] = (255, 0, 0) # Azul: Fondo
category_mask_bgr[np.where(category_mask_np == 1)] = (0, 255, 0) # Verde: Cabello
category_mask_bgr[np.where(category_mask_np == 2)] = (0, 0, 255) # Rojo: Piel del cuerpo
category_mask_bgr[np.where(category_mask_np == 3)] = (255, 255, 0) # Cian: Piel de la cara
category_mask_bgr[np.where(category_mask_np == 4)] = (255, 0, 255) # Rosa: Ropa
category_mask_bgr[np.where(category_mask_np == 5)] = (0, 255, 255) # Amarillo: Otros
# Visualización de las imágenes
cv2.imshow("category_mask_bgr", category_mask_bgr)
segmentation_result_list.clear()
cv2.imshow("Frame", frame)
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 segmenter.segment_async, al cual tendremos que darle el frame que va a analizar junto con la marca de tiempo en milisegundos.
En la línea 36 simplemente realizamos una comparación de si existen resultados para visualizar las imágenes de salida.
Y en la línea 58 tendremos que asegurarnos de limpiar la lista de resultados
IMPORTANTE: La descripción en cuanto a la segmentación de imágenes en un videostreming nos dice lo siguiente: «El método segment_async envía datos de imagen en vivo (una imagen con una marca de tiempo única) para realizar la segmentación de imágenes… El método segment_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 segmentador de imágenes puede descartar imágenes de entrada si es necesario. En otras palabras, no se garantiza que haya una salida para cada imagen de entrada». (Fuente)
Veamos la aplicación del Multi-class Selfie Segmentation Model sobre un videostreaming:
Y por último, el programa completo:
import mediapipe as mp
from mediapipe.tasks.python import vision
from mediapipe.tasks.python import BaseOptions
import cv2
import numpy as np
import time
segmentation_result_list = []
# Función callback para procesar los resultados
def segmentation_callback(result, output_image, timestamp_ms):
segmentation_result_list.append(result)
# Especificar la configuración del ImageSegmenter
options = vision.ImageSegmenterOptions(
base_options=BaseOptions(model_asset_path="./path/selfie_multiclass_256x256.tflite"),
output_category_mask=True,
running_mode=vision.RunningMode.LIVE_STREAM,
result_callback=segmentation_callback)
segmenter = vision.ImageSegmenter.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 segmentador
segmenter.segment_async(frame_rgb, time.time_ns() // 1_000_000)
if len(segmentation_result_list) > 0:
#print(f"segmentation_result_list: {segmentation_result_list}")
segmentation_result = segmentation_result_list[0]
# Extraer la máscara de categorias
category_mask = segmentation_result.category_mask
# Máscara a Numpy
category_mask_np = category_mask.numpy_view()
# Otorgar color a cada segmento
category_mask_bgr = cv2.cvtColor(category_mask_np, cv2.COLOR_GRAY2BGR)
category_mask_bgr[np.where(category_mask_np == 0)] = (255, 0, 0) # Azul: Fondo
category_mask_bgr[np.where(category_mask_np == 1)] = (0, 255, 0) # Verde: Cabello
category_mask_bgr[np.where(category_mask_np == 2)] = (0, 0, 255) # Rojo: Piel del cuerpo
category_mask_bgr[np.where(category_mask_np == 3)] = (255, 255, 0) # Cian: Piel de la cara
category_mask_bgr[np.where(category_mask_np == 4)] = (255, 0, 255) # Rosa: Ropa
category_mask_bgr[np.where(category_mask_np == 5)] = (0, 255, 255) # Amarillo: Otros
# Visualización de las imágenes
cv2.imshow("category_mask_bgr", category_mask_bgr)
segmentation_result_list.clear()
cv2.imshow("Frame", frame)
if cv2.waitKey(1) & 0xFF == 27:
break
cap.release()
cv2.destroyAllWindows()
Y bien, hemos llegado al final de este tutorial. ¡Espero que te haya gustado y lo hayas encontrado útil!. Nos vemos en un siguiente post.
