Nov 25 2024
0

¿Cómo segmentar distintos objetos en una imagen? MediaPipe – DeepLab v3 – Python

Por Administrador

¡Hola, hola Omesito!.

Te doy la bienvenida a un nuevo tutorial, en el que seguimos analizando los modelos de Mediapipe Image Segmentation. En anteriores blog posts ya hemos tratado varios de los modelos que nos ofrece MediaPipe para este campo, sin embargo hoy veremos el último modelo, llamado: Multi-class selfie segmentation model. Así que… ¡Vamos a por ello!.

CONTENIDO

  • Modelo DeepLab v3 de Mediapipe
  • ¡Vamos con el código! Segmentar múltiples objetos con DeepLab v3 de Mediapipe en Python

Modelo DeepLab v3 de Mediapipe

El modelo DeepLab v3 es una herramienta de machine learning para visión artificial, diseñada específicamente para realizar segmentación semántica en imágenes. Este modelo tiene la capacidad de identificar y clasificar diferentes objetos dentro de una imagen, asignando a cada píxel una etiqueta que indica su categoría correspondiente.

Las categorías que este modelo puede segmentar son:

  • 0: Fondo
  • 1: Avión
  • 2: Bicicleta
  • 3: Pájaro
  • 4: Barco
  • 5: Botella
  • 6: Autobús
  • 7: Auto
  • 8: Gato
  • 9: Silla
  • 10: Vaca
  • 11: Mesa de comedor
  • 12: Perro
  • 13: Caballo
  • 14: Motocicleta
  • 15: Persona
  • 16: Planta en maceta
  • 17: Oveja
  • 18: Sofá
  • 19: Tren
  • 20: Televisor

Como se puede observar, cada categoría tiene asignado un número único que estará presente en la máscara de segmentación generada por el modelo. Esta máscara permite distinguir claramente entre las diferentes categorías presentes en la imagen, asignando un color único a cada segmento para facilitar su interpretación visual.

Como podemos visualizar, a la izquierda tenemos la imagen de entrada, en el centro la imagen segmentada, mientras que a la derecha tenemos resaltadas las categorías segmentadas. Podemos ver que en la imagen segmentada, el fondo se ha coloreado de morado, la bicicleta de limón y la persona de marrón. Y esto precisamente es lo que haremos en este tutorial.

En la documentación de MediaPipe, podemos ver la información de DeepLab-v3. Este modelo necesita una imagen de entrada de 257×257, es decir, imágenes que tengan una relación de aspecto de 1. He estado probando con imágenes más grandes, manteniendo la relación de aspecto, y se obtienen buenos resultados. Lo importante está en respetar la relación de aspecto y no usar imágenes más pequeñas que lo recomendado.

Para descargar el modelo, simplemente vamos a tener que dar clic sobre su nombre.

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

Y antes de pasar con el código, debemos tener instalado Mediapipe, para ello nos ayudaremos de pip:

pip install mediapipe

¡Vamos con el código! Segmentar múltiples objetos con DeepLab v3 de 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/deeplab_v3.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 deeplab_v3.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.

# Categorías
categories ={0: 'background',
            1: 'aeroplane',
            2: 'bicycle',
            3: 'bird',
            4: 'boat',
            5: 'bottle',
            6: 'bus',
            7: 'car',
            8: 'cat',
            9: 'chair',
            10: 'cow',
            11: 'diningtable',
            12: 'dog',
            13: 'horse',
            14: 'motorbike',
            15: 'person',
            16: 'pottedplant',
            17: 'sheep',
            18: 'sofa',
            19: 'train',
            20: 'tv'}

# Asignar colores únicos a las categorías
category_colors = np.random.randint(0, 255, size=(len(categories), 3), dtype="uint8")
#print(category_colors)
# ------------------------------------------------

De la línea 16 a la 36 vamos a especificar las categorías que el modelo es capaz de segmentar, esto lo haremos a través de un diccionario, en donde cada clave corresponderá a una clase. Y estas claves además, nos servirán para identificar las clases en la máscara de categorías que obtendremos más adelante, una vez se aplique el segmentador.

En la línea 39 vamos a generar colores aleatorios para cada una de las categorías. Esto nos servirá para más adelante colorear los segmentos.

# APLICANDO EL MODELO SOBRE UNA IMAGEN
# Leer la imagen de entrada
image = cv2.imread("./Inputs/image.jpg")

# Convertir la imagen a RGB
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.

# Obtener los resultados del segmentador
segmentation_result = segmenter.segment(image_rgb)
#print(segmentation_result)

Para empezar con la segmentación tendremos que usar segmenter.segment, y especificar la imagen que deseamos analizar.

A continuación vamos a extraer la máscara de categorías, ya que con esta es con la que vamos a trabajar.

# Convertir la máscara de categorías en un array de Numpy
category_mask = segmentation_result.category_mask
category_mask_np = category_mask.numpy_view()
print(np.unique(category_mask_np))

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 nos permitirá obtener hasta 21 categorías, por lo que la matriz podría estar compuesta de valores comprendidos entre 0 y 21.

Con category_mask.numpy_view(), podremos transformar la máscara de confianza a un array de numpy, que nos servirá para poder manipular la imagen. Una vez transformado, podremos visualizarlo:

A la izquierda tenemos la imagen de entrada, y a la derecha la máscara de categorías resultante. Como los valores que se pueden obtener están entre 0 y 21, es muy difícil poder diferenciarlas, debido a que van a presentar colores muy cercanos al 0, es decir a negro.

Con print(np.unique(category_mask_np)), vamos a poder imprimir en consola a las categorías presentes en la imagen.

A continuación vamos a pasar la máscara a tres canales. Esto nos ayudará a darle color a la matriz de categorías:

# Transformar el category mask a 3 canales
category_mask_bgr = cv2.cvtColor(category_mask_np, cv2.COLOR_GRAY2BGR)
#print(category_mask_bgr.shape)

# Pintar cada segmento con su color correspondiente
for category_id in np.unique(category_mask_np):
    color = category_colors[category_id]
    category_mask_bgr[np.where(category_mask_np == category_id)] = color

En la línea 61, pasamos la imagen a BGR, con ello podremos colorear la máscara de categorías. De la línea 65 a la 67, iteramos por cada una de las categorías encontradas por el segmentador y coloreamos los segmentos con el color aleatorio correspondiente que generamos en la línea 39. Vamos a obtener algo como lo siguiente:

# Aplicar transparencia
alpha = 0.5
final_image = cv2.addWeighted(image, 1 - alpha, category_mask_bgr, alpha, 0)

En la línea 71, sumamos la imagen de entrada y la máscara segmentada coloreada, añadiéndole cierta transparencia. Obtendremos:

A continuación crearemos una imagen en color negro donde irán los nombres de cada una de las categorías. Y estas se resaltarán cuando estén presentes en la imagen:

# VISUALIZAR CATEGORÍAS
black_image = np.zeros((430, 200, 3), dtype="uint8")

y_offset = 20
font_scale = 0.6
line_thickness = 2

for category_id, name in categories.items():
    #print(category_id, name)
    if category_id in np.unique(category_mask_np):
        color = tuple(map(int, category_colors[category_id]))
    else:
        color = (128, 128, 128)

    cv2.putText(black_image,
                f"{category_id}: {name}",
                (30, y_offset),
                cv2.FONT_HERSHEY_SIMPLEX,
                font_scale,
                color,
                line_thickness,
                cv2.LINE_AA)
    y_offset += 20
# Visualización
cv2.imshow("Image", image)
#cv2.imshow("category_mask_np", category_mask_np)
#cv2.imshow("category_mask_bgr", category_mask_bgr)
cv2.imshow("final_image", final_image)
cv2.imshow("black_image", black_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

Entonces, vamos a crear una imagen en color negro en la línea 74. Luego iteramos cada una de las categorías, y le añadiremos el color generado aleatoriamente cuando estas aparezcan como predichas.

Desde la línea 97, a la 101 visualizamos las tres imágenes: imagen de entrada, imagen con los segmentos coloreados y la imagen en negro con cada categoría.

Veamos los resultados:

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/deeplab_v3.tflite"),
    output_category_mask=True,
    running_mode=vision.RunningMode.IMAGE)
segmenter = vision.ImageSegmenter.create_from_options(options)

# Categorías
categories ={0: 'background',
            1: 'aeroplane',
            2: 'bicycle',
            3: 'bird',
            4: 'boat',
            5: 'bottle',
            6: 'bus',
            7: 'car',
            8: 'cat',
            9: 'chair',
            10: 'cow',
            11: 'diningtable',
            12: 'dog',
            13: 'horse',
            14: 'motorbike',
            15: 'person',
            16: 'pottedplant',
            17: 'sheep',
            18: 'sofa',
            19: 'train',
            20: 'tv'}

# Asignar colores únicos a las categorías
category_colors = np.random.randint(0, 255, size=(len(categories), 3), dtype="uint8")
#print(category_colors)
# ------------------------------------------------

# APLICANDO EL MODELO SOBRE UNA IMAGEN
# Leer la imagen de entrada
image = cv2.imread("./Inputs/image_005.jpg")

# Convertir la imagen a RGB
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)

# Convertir la máscara de categorías en un array de Numpy
category_mask = segmentation_result.category_mask
category_mask_np = category_mask.numpy_view()
print(np.unique(category_mask_np))

# Transformar el category mask a 3 canales
category_mask_bgr = cv2.cvtColor(category_mask_np, cv2.COLOR_GRAY2BGR)
#print(category_mask_bgr.shape)

# Pintar cada segmento con su color correspondiente
for category_id in np.unique(category_mask_np):
    color = category_colors[category_id]
    category_mask_bgr[np.where(category_mask_np == category_id)] = color

# Aplicar transparencia
alpha = 0.5
final_image = cv2.addWeighted(image, 1 - alpha, category_mask_bgr, alpha, 0)

# VISUALIZAR CATEGORÍAS
black_image = np.zeros((430, 200, 3), dtype="uint8")

y_offset = 20
font_scale = 0.6
line_thickness = 2

for category_id, name in categories.items():
    #print(category_id, name)
    if category_id in np.unique(category_mask_np):
        color = tuple(map(int, category_colors[category_id]))
    else:
        color = (128, 128, 128)

    cv2.putText(black_image,
                f"{category_id}: {name}",
                (30, y_offset),
                cv2.FONT_HERSHEY_SIMPLEX,
                font_scale,
                color,
                line_thickness,
                cv2.LINE_AA)
    y_offset += 20
# Visualización
cv2.imshow("Image", image)
#cv2.imshow("category_mask_np", category_mask_np)
#cv2.imshow("category_mask_bgr", category_mask_bgr)
cv2.imshow("final_image", final_image)
cv2.imshow("black_image", black_image)
cv2.waitKey(0)
cv2.destroyAllWindows()

Cabe destacar que también probé con otras imágenes, en las que no se obtuvo muy buenos resultados. Estos resultados erróneos se los puede observar porque no siguen la forma del objeto, o no se determina correctamente la categoría:

Esto puede ser por la naturaleza de las imágenes con las que se entrenó el segmentador, y también podría estar relacionado con los colores del objeto y la similitud con el fondo de la imagen.

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 *