Segmentación universal de imágenes con Mask2Former y OneFormer

'Segmentación de imágenes con Mask2Former y OneFormer'

Esta guía presenta Mask2Former y OneFormer, 2 redes neuronales de vanguardia para la segmentación de imágenes. Los modelos están ahora disponibles en 🤗 transformers, una biblioteca de código abierto que ofrece implementaciones fáciles de usar de modelos de vanguardia. En el camino, aprenderás sobre la diferencia entre las diversas formas de segmentación de imágenes.

Segmentación de imágenes

La segmentación de imágenes es la tarea de identificar diferentes “segmentos” en una imagen, como personas o coches. Más técnicamente, la segmentación de imágenes es la tarea de agrupar píxeles con diferentes semánticas. Consulta la página de tareas de Hugging Face para una breve introducción.

La segmentación de imágenes se puede dividir en gran medida en 3 sub tareas: segmentación de instancias, segmentación semántica y segmentación panóptica, con numerosos métodos y arquitecturas de modelos para realizar cada sub tarea.

• La segmentación de instancias es la tarea de identificar diferentes “instancias”, como personas individuales, en una imagen. La segmentación de instancias es muy similar a la detección de objetos, excepto que nos gustaría producir un conjunto de máscaras de segmentación binarias, en lugar de cajas delimitadoras, con etiquetas de clase correspondientes. Las instancias a menudo también se llaman “objetos” o “cosas”. Ten en cuenta que las instancias individuales pueden superponerse.
• La segmentación semántica es la tarea de identificar diferentes “categorías semánticas”, como “persona” o “cielo”, de cada píxel en una imagen. A diferencia de la segmentación de instancias, no se hace distinción entre instancias individuales de una categoría semántica dada; solo queremos obtener una máscara para la categoría “persona”, por ejemplo, en lugar de para las personas individuales. Las categorías semánticas que no tienen instancias individuales, como “cielo” o “hierba”, a menudo se denominan “cosas”, para hacer la distinción con “objetos” (nombres geniales, ¿verdad?). Ten en cuenta que no es posible la superposición entre categorías semánticas, ya que cada píxel pertenece a una categoría.
• La segmentación panóptica, introducida en 2018 por Kirillov et al., tiene como objetivo unificar la segmentación de instancias y la segmentación semántica, haciendo que los modelos simplemente identifiquen un conjunto de “segmentos”, cada uno con una máscara binaria y una etiqueta de clase correspondiente. Los segmentos pueden ser tanto “objetos” como “cosas”. A diferencia de la segmentación de instancias, no es posible la superposición entre diferentes segmentos.

La figura a continuación ilustra la diferencia entre las 3 sub tareas (tomada de esta publicación del blog).

• Generación de activos 3D IA para el desarrollo de juegos #3
• Usando LoRA para Ajuste Fino de Difusión Eficiente y Estable
• Generación de Activos 2D IA para el Desarrollo de Juegos #4

En los últimos años, los investigadores han propuesto varias arquitecturas que estaban típicamente adaptadas a la segmentación de instancias, semántica o panóptica. La segmentación de instancias y panóptica se resolvían típicamente mediante la generación de un conjunto de máscaras binarias + etiquetas correspondientes por instancia de objeto (muy similar a la detección de objetos, excepto que se genera una máscara binaria en lugar de una caja delimitadora por instancia). Esto a menudo se llama “clasificación de máscara binaria”. Por otro lado, la segmentación semántica se resolvía típicamente haciendo que los modelos generen un solo “mapa de segmentación” con una etiqueta por píxel. Por lo tanto, la segmentación semántica se trataba como un problema de “clasificación por píxel”. Algunos modelos populares de segmentación semántica que adoptan este paradigma son SegFormer, del cual escribimos una extensa publicación en el blog, y UPerNet.

Segmentación universal de imágenes

Afortunadamente, desde alrededor de 2020, las personas comenzaron a proponer modelos que pueden resolver las 3 tareas (segmentación de instancias, semántica y panóptica) con una arquitectura unificada, utilizando el mismo paradigma. Esto comenzó con DETR, que fue el primer modelo que resolvió la segmentación panóptica utilizando un paradigma de “clasificación de máscara binaria”, tratando las clases de “objetos” y “cosas” de manera unificada. La innovación clave fue tener un decodificador Transformer que generara un conjunto de máscaras binarias + clases de manera paralela. Esto se mejoró posteriormente en el artículo MaskFormer, que demostró que el paradigma de “clasificación de máscara binaria” también funciona muy bien para la segmentación semántica.

Mask2Former amplía esto a la segmentación de instancias al mejorar aún más la arquitectura de la red neuronal. Por lo tanto, hemos evolucionado de arquitecturas separadas a lo que los investigadores ahora llaman arquitecturas de “segmentación universal de imágenes”, capaces de resolver cualquier tarea de segmentación de imágenes. Curiosamente, todos estos modelos universales adoptan el paradigma de “clasificación de máscara”, descartando por completo el paradigma de “clasificación por píxel”. A continuación se muestra una figura que ilustra la arquitectura de Mask2Former (tomada del artículo original).

En resumen, una imagen se envía primero a una columna vertebral (que, en el artículo, puede ser ResNet o Swin Transformer) para obtener una lista de mapas de características de baja resolución. A continuación, estos mapas de características se mejoran utilizando un módulo decodificador de píxeles para obtener características de alta resolución. Por último, un decodificador Transformer toma un conjunto de consultas y las transforma en un conjunto de máscaras binarias y predicciones de clases, condicionado a las características del decodificador de píxeles.

Tenga en cuenta que Mask2Former aún necesita ser entrenado en cada tarea por separado para obtener resultados de vanguardia. Esto ha sido mejorado por el modelo OneFormer, que obtiene un rendimiento de vanguardia en las 3 tareas al entrenar solo en una versión panóptica del conjunto de datos (!), agregando un codificador de texto para condicionar el modelo en entradas “instancia”, “semántica” o “panóptica”. Este modelo también está disponible a partir de hoy en transformers. Es aún más preciso que Mask2Former, pero tiene una mayor latencia debido al codificador de texto adicional. Consulte la figura a continuación para obtener una descripción general de OneFormer. Utiliza Swin Transformer o el nuevo modelo DiNAT como base.

Inferencia con Mask2Former y OneFormer en Transformers

El uso de Mask2Former y OneFormer es bastante sencillo y muy similar a su predecesor MaskFormer. Instanciemos un modelo Mask2Former desde el centro entrenado en el conjunto de datos panóptico COCO, junto con su procesador. Tenga en cuenta que los autores lanzaron nada menos que 30 puntos de control entrenados en diversos conjuntos de datos.

from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation

processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")

A continuación, carguemos la conocida imagen de gatos del conjunto de datos COCO, en la cual realizaremos la inferencia.

from PIL import Image

url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)
image

Preparamos la imagen para el modelo utilizando el procesador de imágenes y la pasamos por el modelo.

inputs = processor(image, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)

El modelo devuelve un conjunto de máscaras binarias y logaritmos de clases correspondientes. Los resultados sin procesar de Mask2Former se pueden postprocesar fácilmente utilizando el procesador de imágenes para obtener las predicciones finales de segmentación de instancia, semántica o panóptica:

prediction = processor.post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
print(prediction.keys())

En la segmentación panóptica, la predicción final contiene 2 cosas: un mapa de segmentación de forma (altura, ancho) donde cada valor codifica el ID de instancia de un píxel dado, así como una segments_info correspondiente. El segments_info contiene más información sobre los segmentos individuales del mapa (como su ID de clase/categoría). Tenga en cuenta que Mask2Former genera propuestas de máscaras binarias de forma (96, 96) para mayor eficiencia y el argumento target_sizes se utiliza para redimensionar la máscara final al tamaño original de la imagen.

Visualicemos los resultados:

from collections import defaultdict
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib import cm

def draw_panoptic_segmentation(segmentation, segments_info):
    # obtener el mapa de colores utilizado
    viridis = cm.get_cmap('viridis', torch.max(segmentation))
    fig, ax = plt.subplots()
    ax.imshow(segmentation)
    instances_counter = defaultdict(int)
    handles = []
    # para cada segmento, dibujar su leyenda
    for segment in segments_info:
        segment_id = segment['id']
        segment_label_id = segment['label_id']
        segment_label = model.config.id2label[segment_label_id]
        label = f"{segment_label}-{instances_counter[segment_label_id]}"
        instances_counter[segment_label_id] += 1
        color = viridis(segment_id)
        handles.append(mpatches.Patch(color=color, label=label))
        
    ax.legend(handles=handles)

draw_panoptic_segmentation(**panoptic_segmentation)

Aquí, podemos ver que el modelo es capaz de detectar los gatos individuales y los mandos a distancia en la imagen. Por otro lado, la segmentación semántica simplemente crearía una única máscara para la categoría “gato”.

Para realizar inferencia con OneFormer, que tiene una API idéntica excepto que también toma una entrada adicional de texto como prompt, consulte el cuaderno de demostración .

Afinando Mask2Former y OneFormer en Transformers

Para afinar Mask2Former/OneFormer en un conjunto de datos personalizado para instancia, segmentación semántica y panóptica, consulte nuestros cuadernos de demostración . MaskFormer, Mask2Former y OneFormer comparten una API similar, por lo que la actualización desde MaskFormer es fácil y requiere cambios mínimos.

Los cuadernos de demostración utilizan MaskFormerForInstanceSegmentation para cargar el modelo, mientras que deberá cambiar a Mask2FormerForUniversalSegmentation o OneFormerForUniversalSegmentation . En el caso del procesamiento de imágenes para Mask2Former, también deberá cambiar a Mask2FormerImageProcessor . También puede cargar el procesador de imágenes utilizando la clase AutoImageProcessor que se encarga automáticamente de cargar el procesador correcto correspondiente a su modelo. Por otro lado, OneFormer requiere un OneFormerProcessor , que prepara las imágenes, junto con una entrada de texto, para el modelo.

¡Eso es todo! Ahora conoces la diferencia entre segmentación de instancia, semántica y panóptica, así como cómo utilizar arquitecturas “universales” como Mask2Former y OneFormer utilizando la biblioteca transformers de 🤗.

Esperamos que hayas disfrutado esta publicación y hayas aprendido algo. No dudes en hacernos saber si estás satisfecho con los resultados al afinar Mask2Former o OneFormer.

Si te gustó este tema y quieres aprender más, te recomendamos los siguientes recursos:

• Nuestros cuadernos de demostración para MaskFormer , Mask2Former y OneFormer , que brindan una visión más amplia sobre la inferencia (incluyendo visualización) así como el afinamiento en datos personalizados.
• Las [demostraciones en vivo] para Mask2Former y OneFormer disponibles en el Hugging Face Hub, que puedes usar para probar rápidamente los modelos en entradas de muestra de tu elección.