Nov 04 2024
0

Cómo segmentar el cabello con Mediapipe en Python

Por Administrador

¡Hola, hola Omesitos!.

Seguimos analizando los modelos que nos otorga Mediapipe Image Segmentation. En el anterior post vimos como usar el Selfie Segmentation Model, y en esta ocasión veremos como segmentar el cabello con: Hair Segmentation Model. Así que… ¡Vamos a por ello!.

CONTENIDO

  • ¿Qué es la segmentación de cabello en visión por computador?
  • ¿Cómo usar la segmentación de cabello en Mediapipe?
  • Hair segmentation model
  • ¿Cómo segmentar cabello en una imagen con Mediapipe en Python?
  • ¿Cómo segmentar cabello en videos con Mediapipe en Python?
  • ¿Cómo usar la segmentación de cabello en videostream con Mediapipe en Python?

¿Qué es la segmentación de cabello en visión por computador?

En visión por computador, la segmentación de cabello es un proceso que permite identificar los pixeles dentro de una imagen o video que corresponden al cabello, diferenciando dicha región del fondo y de la persona.

De hecho, la segmentación de cabello permite obtener una máscara en donde se podrán diferenciar los segmentos correspondientes al cabello de otros elementos de la imagen, como la piel, el fondo, o accesorios. Esta máscara segmentada facilita el análisis y la manipulación del cabello en aplicaciones de visión por computadora, como el cambio de color, la aplicación de filtros o efectos visuales.

A continuación podremos ver el resultado de la aplicación de la segmentación de cabello. A la izquierda tenemos a la imagen de entrada, mientras que a la derecha, la máscara binaria, cuyo color blanco representa la presencial del cabello, y el color negro, el fondo.

¿Cómo usar la segmentación de cabello en 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:

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 si 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.

Hair segmentation model

Este modelo de segmentación toma la imagen de una persona, identifica el área del cabello en su cabeza y genera un mapa de segmentación para esta región específica. Esto permite usar el modelo para cambiar el color del cabello o aplicar otros efectos visuales. La salida del modelo clasifica cada píxel en las siguientes categorías:

  • 0 – background (fondo): todo lo que no es cabello en la imagen.
  • 1 – hair (cabello): los píxeles correspondientes al cabello.

Este modelo necesita como entrada una imagen de mínimo 512×512 pixeles, es decir una imagen cuadrada. Para descargar el modelo de segmentación de cabello tendremos que dar clic sobre él, y listo.

Además de estas especificaciones tenemos el apartado de Model Card, donde podremos obtener más información del modelo como: uso, limitaciones, métricas, etc.

Y bien, una vez que tengamos el modelo, vamos a pasar con la programación.

¿Cómo segmentar cabello en una imagen con Mediapipe en 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/hair_segmenter.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 hair_segmenter.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, 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.

# 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 de cabello:

category_mask = segmentation_result.category_mask
confidence_mask_1 = segmentation_result.confidence_masks[0]
confidence_mask_2 = segmentation_result.confidence_masks[1]

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 de cabello nos permitirá extraer dos matrices con las confianzas asociadas al cabello y al fondo veamos:

segmentation_result.confidence_masks[0], nos devolverá el nivel de confianza correspondiente al fondo de la imagen con valores entre 0 y 1, donde aquellos valores más cercanos a 1 corresponderán a un mayor nivel de confianza de la presencia del fondo.

segmentation_result.confidence_masks[1], nos devolverá el nivel de confianza correspondiente al cabello en la imagen con valores entre 0 y 1, donde aquellos valores más cercanos a 1 corresponderán a un mayor nivel de confianza de la presencia del cabello.

# Convertir la máscara en una vista de numpy
category_mask_np = category_mask.numpy_view()
print(np.unique(category_mask_np))

confidence_mask_np_1 = confidence_mask_1.numpy_view()
confidence_mask_np_2 = confidence_mask_2.numpy_view()

Tanto para category_mask como para las confidence_masks les aplicaremos numpy_view(). Esto para convertir a ambos en arrays de numpy. Este es un paso fundamental ya que nos permitirá manipular, analizar y visualizar las imágenes y datos de salida.

En la línea 30 podremos imprimir los valores únicos que contiene category_mask, esto nos permitirá conocer con qué valores trabajar. Si imprimimos veremos que obtenemos: [0 1]. Esto quiere decir que toda la imagen contiene unicamente valores de 0 (fondo) o 1 (cabello). Si visualizamos, toda la imagen parecerá estar en negro, por la similitud de colores de 0 y 1, así que tendremos que cambiar a color a la región del cabello.

# ---------- CAMBIAR EL COLOR DEL CABELLO ----------
hair_color = (255, 0, 0)
#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 == 1)] = hair_color

# Combinar la imagen original y la capa de color
final_image = cv2.addWeighted(image, 1, category_mask_bgr, 0.6, 0)
# --------------------------------------------------

De la línea 35 a la 44 procederemos a hacer el cambio de color a la máscara que obtenemos de la segmentación de cabello. En la línea 36 establecemos el color al que queremos cambiar en BGR. En mi caso he escogido el color azul, pero podrías escoger otro color.

En la línea 38 convertiremos la imagen con un solo canal (category_mask_np), a BGR para poder darle color. De hecho nos podemos ayudar de las impresiones de las líneas 37 y 38 para conocer el cambio de 1 a los 3 canales.

En la línea 40, tomaremos la ubicación de los 1 de la máscara (región correspondiente al cabello), y los transformaremos al color que establecimos en la línea 36. De hecho si visualizamos category_mask_bgr, obtendríamos:

A la izquierda podemos ver la imagen de entrada, y a la derecha la máscara obtenida del segmentador de cabello. Ahora solo nos queda sumar estas dos imágenes. Esto precisamente lo haremos en la línea 43.

cv2.imshow("image", image)
#cv2.imshow("category_mask_np", category_mask_np)
#cv2.imshow("confidence_mask_np_1", confidence_mask_np_1)
#cv2.imshow("confidence_mask_np_2", confidence_mask_np_2)
#cv2.imshow("category_mask_bgr", category_mask_bgr)
cv2.imshow("final_image", final_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

Finalmente desde la línea 46, 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.

En las imágenes resultantes podemos darnos cuenta de lo siguiente:

  • Las máscaras no se adaptan perfectamente a pequeñas porciones de cabello, pero se adaptan bien cuando hay diferenciación entre el color de cabello y el fondo.
  • Si existe similitud entre el color del cabello y el fondo o accesorios, puede que esos elementos sean segmentados como cabello también.
  • Se probó el segmentador de cabello con dos personas en la imagen de entrada y sí pudo identificar el cabello de ambas.
  • Si queremos darle color al cabello, debemos tomar en cuenta el color base del mismo, puesto que podrían haber cambios el el color resultante. Esto lo podemos apreciar en la última imagen en donde a pesar de que se aplicó el color azul, la mujer con cabello rosa obtuvo un color diferente a la mujer con cabello castaño.

A continuación tenemos el programa completo sobre la segmentación de cabello sobre una imagen:

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/hair_segmenter.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)

# Aplicar el segmentador
segmentation_result = segmenter.segment(image_rgb)
print("segmentation_result", segmentation_result)

category_mask = segmentation_result.category_mask
confidence_mask_1 = segmentation_result.confidence_masks[0]
confidence_mask_2 = segmentation_result.confidence_masks[1]

# Convertir la máscara en una vista de numpy
category_mask_np = category_mask.numpy_view()
print(np.unique(category_mask_np))

confidence_mask_np_1 = confidence_mask_1.numpy_view()
confidence_mask_np_2 = confidence_mask_2.numpy_view()

# ---------- CAMBIAR EL COLOR DEL CABELLO ----------
hair_color = (255, 0, 0)
#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 == 1)] = hair_color

# Combinar la imagen original y la capa de color
final_image = cv2.addWeighted(image, 1, category_mask_bgr, 0.6, 0)
# --------------------------------------------------

cv2.imshow("image", image)
#cv2.imshow("category_mask_np", category_mask_np)
#cv2.imshow("confidence_mask_np_1", confidence_mask_np_1)
#cv2.imshow("confidence_mask_np_2", confidence_mask_np_2)
#cv2.imshow("category_mask_bgr", category_mask_bgr)
cv2.imshow("final_image", final_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

¿Cómo segmentar cabello en videos con Mediapipe en Python?

Debido a que el programa para la segmentación en video presenta similitudes con la segmentación de cabello 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/hair_segmenter.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)

     # Convertir la máscara en una vista de numpy
     category_mask = segmentation_result.category_mask
     confidence_mask_1 = segmentation_result.confidence_masks[0]
     confidence_mask_2 = segmentation_result.confidence_masks[1]

     # Visualizar las máscaras
     category_mask_np = category_mask.numpy_view()

     confidence_mask_np_1 = confidence_mask_1.numpy_view()
     confidence_mask_np_2 = confidence_mask_2.numpy_view()

     # ------- CAMBIAR EL COLOR DEL CABELLO -------
     hair_color = (255, 0, 0) # Color en BGR, puede probar otros colores
     category_mask_bgr = cv2.cvtColor(category_mask_np, cv2.COLOR_GRAY2BGR)
     category_mask_bgr[np.where(category_mask_np == 1)] = hair_color

     # Combinar la imagen original y la capa de color
     final_frame = cv2.addWeighted(frame, 1, category_mask_bgr, 0.6, 0)
     # --------------------------------------------
     # Visualización de las imágenes
     cv2.imshow("final_frame", final_frame)
     cv2.imshow("Frame", 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 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 de cabello 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/hair_segmenter.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)

     # Convertir la máscara en una vista de numpy
     category_mask = segmentation_result.category_mask
     confidence_mask_1 = segmentation_result.confidence_masks[0]
     confidence_mask_2 = segmentation_result.confidence_masks[1]

     # Visualizar las máscaras
     category_mask_np = category_mask.numpy_view()

     confidence_mask_np_1 = confidence_mask_1.numpy_view()
     confidence_mask_np_2 = confidence_mask_2.numpy_view()

     # ------- CAMBIAR EL COLOR DEL CABELLO -------
     hair_color = (255, 0, 0) # Color en BGR, puede probar otros colores
     category_mask_bgr = cv2.cvtColor(category_mask_np, cv2.COLOR_GRAY2BGR)
     category_mask_bgr[np.where(category_mask_np == 1)] = hair_color

     # Combinar la imagen original y la capa de color
     final_frame = cv2.addWeighted(frame, 1, category_mask_bgr, 0.6, 0)
     # --------------------------------------------
     # Visualización de las imágenes
     cv2.imshow("final_frame", final_frame)
     cv2.imshow("Frame", frame)

     if cv2.waitKey(1) & 0xFF == 27:
          break
cap.release()
cv2.destroyAllWindows()

¿Cómo usar la segmentación de cabello en videostream con Mediapipe en Python?

Finalmente veremos como podemos aplicar la segmentación de cabello 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/hair_segmenter.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]
          # Convertir la máscara en una vista de numpy
          category_mask = segmentation_result.category_mask
          confidence_mask_1 = segmentation_result.confidence_masks[0]
          confidence_mask_2 = segmentation_result.confidence_masks[1]

          # Visualizar las máscaras
          category_mask_np = category_mask.numpy_view()

          confidence_mask_np_1 = confidence_mask_1.numpy_view()
          confidence_mask_np_2 = confidence_mask_2.numpy_view()

          # ------- CAMBIAR EL COLOR DEL CABELLO -------
          hair_color = (255, 0, 255) # Color en BGR, puede probar otros colores
          category_mask_bgr = cv2.cvtColor(category_mask_np, cv2.COLOR_GRAY2BGR)
          category_mask_bgr[np.where(category_mask_np == 1)] = hair_color

          # Combinar la imagen original y la capa de color
          final_frame = cv2.addWeighted(frame, 1, category_mask_bgr, 0.6, 0)
          # --------------------------------------------
          # Visualización de las imágenes
          cv2.imshow("final_frame", final_frame)
     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 60 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 de Hair Segmentation model sobre un videostreaming:

Y por último, el programa completo para aplicar el Hair Segmentation Model sobre un video streaming:

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/hair_segmenter.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]
          # Convertir la máscara en una vista de numpy
          category_mask = segmentation_result.category_mask
          confidence_mask_1 = segmentation_result.confidence_masks[0]
          confidence_mask_2 = segmentation_result.confidence_masks[1]

          # Visualizar las máscaras
          category_mask_np = category_mask.numpy_view()

          confidence_mask_np_1 = confidence_mask_1.numpy_view()
          confidence_mask_np_2 = confidence_mask_2.numpy_view()

          # ------- CAMBIAR EL COLOR DEL CABELLO -------
          hair_color = (255, 0, 255) # Color en BGR
          category_mask_bgr = cv2.cvtColor(category_mask_np, cv2.COLOR_GRAY2BGR)
          category_mask_bgr[np.where(category_mask_np == 1)] = hair_color

          # Combinar la imagen original y la capa de color
          final_frame = cv2.addWeighted(frame, 1, category_mask_bgr, 0.6, 0)
          # --------------------------------------------
          # Visualización de las imágenes
          cv2.imshow("final_frame", final_frame)
     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.

Referencias

Publicado enBlog MediaPipe

Deja una respuesta

Tu dirección de correo electrónico no será publicada. Los campos obligatorios están marcados con *