El inpainting de imágenes es uno de los problemas clásicos en visión por computadora, y tiene como objetivo restaurar las regiones enmascaradas en una imagen con contenido plausible y natural. Los trabajos existentes que emplean técnicas de inpainting de imágenes tradicionales como Redes Antagonistas Generativas o GAN, y AutoCodificadores Variacionales o VAE, a menudo requieren características auxiliares diseñadas a mano, pero al mismo tiempo, no entregan resultados satisfactorios. En los últimos años, los métodos basados en difusión han ganado popularidad dentro de la comunidad de visión por computadora debido a sus capacidades de generación de imágenes de alta calidad, diversidad de salida y control fino. Los intentos iniciales de emplear modelos de difusión para inpainting de imágenes guiadas por texto modificaron la estrategia de ruido estándar al muestrear las regiones enmascaradas desde un modelo de difusión preentrenado, y las regiones no enmascaradas desde la imagen dada. Aunque estos métodos resultaron en un rendimiento satisfactorio en tareas de inpainting de imágenes simples, lucharon con formas de máscara complejas, textos y contenidos de imagen que resultaron en una falta general de coherencia. La falta de coherencia observada en estos métodos se puede deber principalmente a su limitado conocimiento perceptual de los límites de la máscara y el contexto de la región de la imagen no enmascarada.

A pesar de los avances, la investigación y el desarrollo de estos modelos en los últimos años, el inpainting de imágenes sigue siendo un obstáculo importante para los desarrolladores de visión por computadora. Las adaptaciones actuales de los modelos de difusión para tareas de inpainting de imágenes implican modificar la estrategia de muestreo o el desarrollo de modelos de difusión específicos para inpainting, que a menudo sufren de una reducción en la calidad de la imagen y una semántica inconsistente. Para abordar estos desafíos y allanar el camino para los modelos de inpainting de imágenes, en este artículo, hablaremos sobre BrushNet, un marco de ingeniería dual de rama novedoso que incorpora características de imagen enmascarada a nivel de píxel en cualquier modelo de difusión preentrenado, garantizando así la coherencia y el resultado mejorado en tareas de inpainting de imágenes. El marco de BrushNet introduce un nuevo paradigma bajo el cual el marco divide las características de la imagen y el latente ruidoso en ramas separadas. La división de las características de la imagen y los latentes ruidosos disminuye drásticamente la carga de aprendizaje para el modelo y facilita una incorporación matizada de la información esencial de la imagen enmascarada de manera jerárquica. Además del marco de BrushNet, también hablaremos sobre BrushBench y BrushData que facilitan la evaluación del rendimiento basada en segmentación y el entrenamiento de inpainting de imágenes, respectivamente.

Este artículo tiene como objetivo cubrir el marco de BrushNet en profundidad, y exploramos el mecanismo, la metodología, la arquitectura del marco junto con su comparación con marcos de estado del arte. Así que comencemos.

BrushNet: Inpainting de imágenes con difusión de doble rama

El inpainting de imágenes, un método que intenta restaurar las regiones faltantes de una imagen mientras mantiene la coherencia general, ha sido un problema de larga data en el campo de la visión por computadora, y ha molestado a desarrolladores e investigadores durante varios años. El inpainting de imágenes encuentra sus aplicaciones en una amplia variedad de tareas de visión por computadora, incluyendo edición de imágenes y pruebas virtuales. Recientemente, los modelos de difusión como Stable Diffusion, y Stable Diffusion 1.5 han demostrado una capacidad notable para generar imágenes de alta calidad, y proporcionan a los usuarios la flexibilidad para controlar los controles semánticos y estructurales. El potencial notable de los modelos de difusión es lo que ha llevado a los investigadores a recurrir a los modelos de difusión para tareas de inpainting de imágenes de alta calidad que se alinean con los textos de entrada.

Los métodos empleados por los marcos de inpainting de imágenes guiadas por texto basados en difusión tradicionales se pueden dividir en dos categorías, Modificación de la estrategia de muestreo y Modelos de inpainting dedicados. El método de modificación de la estrategia de muestreo modifica el proceso de ruido estándar al muestrear las regiones enmascaradas desde un modelo de difusión preentrenado, y copiar y pegar las regiones no enmascaradas desde la imagen dada en cada paso de ruido. Aunque los enfoques de modificación de la estrategia de muestreo pueden implementarse en modelos de difusión arbitrarios, a menudo resultan en resultados de inpainting incoherentes, ya que tienen un conocimiento perceptual limitado de los límites de la máscara y el contexto de la región de la imagen no enmascarada. Por otro lado, los modelos de inpainting dedicados ajustan un modelo de inpainting de imágenes diseñado específicamente ampliando las dimensiones del canal de entrada del modelo de difusión base para incorporar la imagen dañada y las máscaras. Mientras que los modelos de inpainting dedicados permiten que el modelo de difusión genere resultados más satisfactorios con modelos especializados que son conscientes de la forma y el contenido, puede que no sea el mejor diseño arquitectónico para los modelos de inpainting de imágenes.

Como se demuestra en la siguiente imagen, los modelos de inpainting dedicados fusionan el latente de la imagen enmascarada, el latente ruidoso, el texto y la máscara en una etapa temprana. El diseño arquitectónico de dichos modelos de inpainting dedicados fácilmente influye en las características de la imagen enmascarada y evita que las capas posteriores en la arquitectura de UNet obtengan características de imagen enmascarada puras debido a la influencia del texto. Además, manejar la generación y la condición en una sola rama impone una carga adicional en la arquitectura de UNet, y como estos enfoques también requieren ajuste fino en diferentes variaciones del modelo de difusión base, estos enfoques a menudo son agotadores en términos de tiempo y tienen una transferibilidad limitada.

Puede parecer que agregar una rama adicional dedicada a extraer características de la imagen enmascarada podría ser una solución adecuada para los problemas mencionados anteriormente, sin embargo, los marcos existentes a menudo resultan en la extracción e inserción de información inadecuada cuando se aplican directamente al inpainting. Como resultado, los marcos existentes como ControlNet producen resultados insatisfactorios en comparación con los modelos de inpainting dedicados. Para abordar este problema de la manera más efectiva posible, el marco de BrushNet introduce una rama adicional a la red de difusión original y crea así una arquitectura más adecuada para tareas de inpainting de imágenes. El diseño y la arquitectura del marco de BrushNet se pueden resumir en tres puntos.

1. En lugar de inicializar las capas de convolución aleatoriamente, el marco de BrushNet implementa un codificador VAE para procesar la imagen enmascarada. Como resultado, el marco de BrushNet puede extraer las características de la imagen para adaptarse a la distribución de UNet de manera más efectiva.
2. El marco de BrushNet incorpora gradualmente la capa de características completa de UNet capa por capa en la arquitectura de UNet preentrenada, un enfoque jerárquico que permite el control denso por píxel.
3. El marco de BrushNet elimina la atención cruzada de texto de la componente de UNet para garantizar que se considere solo la información de la imagen pura en la rama adicional. Además, el marco de BrushNet también propone implementar una estrategia de mezcla difuminada para lograr una mejor coherencia junto con un rango más amplio de controlabilidad en las regiones no enmascaradas de la imagen.

BrushNet: Método y Arquitectura

La siguiente figura nos da una visión general del marco de BrushNet.

Como se puede observar, el marco emplea una estrategia de rama dual para la inserción de la guía de la imagen enmascarada y utiliza operaciones de mezcla con una máscara difuminada para garantizar una mejor preservación de las regiones no enmascaradas. Es digno de destacar que el marco de BrushNet es capaz de ajustar la escala agregada para lograr un control flexible. Para una imagen de entrada enmascarada y la máscara dada, el modelo de BrushNet produce una imagen inpainting. El modelo primero reduce la máscara para adaptar el tamaño del latente y la imagen enmascarada se alimenta como entrada al codificador VAE para alinear la distribución del espacio latente. El modelo luego concatena el latente de la imagen enmascarada, el latente ruidoso y la máscara reducida, y lo utiliza como entrada. Las características que el modelo extrae se agregan a la capa de UNet preentrenada después de un bloque de convolución cero. Después del ruido, el modelo mezcla la imagen enmascarada y la imagen generada con una máscara difuminada.

Guía de la imagen enmascarada

El marco de BrushNet inserta la característica de la imagen enmascarada en la red de difusión preentrenada utilizando una rama adicional, que separa explícitamente la extracción de características de la imagen enmascarada del proceso de generación de la imagen. La entrada se forma concatenando el latente de la imagen enmascarada, el latente ruidoso y la máscara reducida. Para ser más específicos, el latente ruidoso proporciona información para la generación de la imagen durante el proceso de generación actual y ayuda al marco a mejorar la coherencia semántica de la característica de la imagen enmascarada. El marco de BrushNet extrae el latente de la imagen enmascarada de la imagen enmascarada utilizando un AutoCodificador Variacional. Además, el marco emplea interpolación cúbica para reducir la máscara en un intento de garantizar que el tamaño de la máscara se alinee con el latente de la imagen enmascarada y el latente ruidoso. Para procesar las características de la imagen enmascarada, el marco de BrushNet implementa una copia del modelo de difusión preentrenado y excluye las capas de atención cruzada del modelo de difusión. La razón es que los pesos preentrenados del modelo de difusión sirven como un fuerte prior para la extracción de las características de la imagen enmascarada y excluir las capas de atención cruzada garantiza que el modelo solo considere la información de la imagen pura en la rama adicional. El marco de BrushNet inserta las características en la arquitectura de UNet preentrenada capa por capa, lo que permite un control denso por píxel jerárquico y también emplea capas de convolución cero para establecer una conexión entre el modelo de BrushNet trainable y el modelo bloqueado, garantizando que el ruido dañino no tenga influencia sobre los estados ocultos en el modelo trainable durante las etapas de entrenamiento inicial.

Operación de mezcla

Como se mencionó anteriormente, realizar la operación de mezcla en el espacio latente cambia el tamaño de las máscaras que a menudo resulta en varias inexactitudes y el marco de BrushNet se enfrenta a un problema similar cuando cambia el tamaño de la máscara para adaptarlo al tamaño del espacio latente. Además, es digno de destacar que las operaciones de codificación y decodificación en los AutoCodificadores Variacionales tienen operaciones limitadas inherentes y pueden no garantizar la reconstrucción completa de la imagen. Para garantizar que el marco reconstruya una imagen coherente y completa de la región no enmascarada, los trabajos existentes han implementado diferentes técnicas como copiar las regiones no enmascaradas desde la imagen original. Aunque el enfoque funciona, a menudo resulta en una falta de coherencia semántica en la generación de los resultados finales. Por otro lado, otros métodos como adoptar operaciones de mezcla latente enfrentan dificultades para preservar la información deseada en las regiones no enmascaradas.

Control flexible

El diseño arquitectónico del marco de BrushNet lo hace una opción adecuada para integraciones de plug and play inherentemente con varios modelos de difusión preentrenados y permite una escala de preservación flexible. Dado que el marco de BrushNet no altera los pesos del modelo de difusión preentrenado, los desarrolladores tienen la flexibilidad de integrarlo como un componente de plug and play con un modelo de difusión ajustado, lo que permite una adopción y experimentación fáciles con modelos preentrenados. Además, los desarrolladores también tienen la opción de controlar la escala de preservación de las regiones no enmascaradas incorporando las características del modelo de BrushNet en el modelo de difusión congelado con el peso w dado que determina la influencia del marco de BrushNet en la escala de preservación, lo que ofrece a los desarrolladores la capacidad de ajustar los niveles de preservación deseados. Finalmente, el marco de BrushNet permite a los usuarios ajustar la escala de difuminado y decidir si implementar o no la operación de difuminado, lo que facilita la personalización de la escala de preservación de las regiones no enmascaradas, lo que permite ajustes flexibles y un control fino sobre el proceso de inpainting de la imagen.

BrushNet: Implementación y Resultados

Para analizar sus resultados, el marco de BrushNet propone BrushBench, un conjunto de datos de inpainting de imágenes basado en segmentación con más de 600 imágenes, con cada imagen acompañada de una máscara anotada por humanos y una anotación de subtítulo. Las imágenes en el conjunto de datos de BrushBench se distribuyen uniformemente entre imágenes naturales y artificiales, y también garantiza una distribución uniforme entre diferentes categorías, lo que permite una evaluación justa en diferentes categorías. Para mejorar aún más el análisis de las tareas de inpainting, el marco de BrushNet categoriza el conjunto de datos en dos partes distintas en función de los métodos utilizados: segmentación basada en y máscaras de pincel.

Comparación cuantitativa

La siguiente tabla compara el marco de BrushNet con los modelos de inpainting de imágenes basados en difusión existentes en el conjunto de datos de BrushBench con el modelo de Stable Diffusion como modelo base.

Como se puede observar, el marco de BrushNet demuestra una eficiencia notable en la preservación de la región enmascarada, la alineación del texto y la calidad de la imagen. Además, los modelos como Stable Diffusion Inpainting, HD-Painter, PowerPaint y otros demuestran un rendimiento fuerte en tareas de inpainting de imágenes dentro de la imagen, aunque no logran replicar su rendimiento en tareas de inpainting fuera de la imagen, especialmente en términos de alineación del texto y calidad de la imagen. En general, el marco de BrushNet entrega los resultados más fuertes.

Además, la siguiente tabla compara el marco de BrushNet con los modelos de inpainting de imágenes basados en difusión existentes en el conjunto de datos de EditBench, y el rendimiento es comparable al observado en el conjunto de datos de BrushBench. Los resultados indican que el marco de BrushNet entrega un rendimiento fuerte en una amplia gama de tareas de inpainting de imágenes con diferentes tipos de máscaras.

Comparación cualitativa

La siguiente figura compara cualitativamente el marco de BrushNet con los métodos de inpainting de imágenes existentes, con resultados que cubren imágenes de inteligencia artificial y naturales en diferentes tareas de inpainting, incluyendo inpainting de máscara aleatoria, inpainting de máscara de segmentación dentro y fuera de la imagen.

Como se puede observar, el marco de BrushNet entrega resultados notables en la coherencia de la región no enmascarada y las regiones coherentes, y logra con éxito la conciencia de la información de fondo gracias a la implementación del enfoque de desacoplamiento de doble rama. Además, la rama no modificada del modelo de difusión preentrenado también proporciona la ventaja de cubrir diferentes dominios de datos como anime y pintura de manera mejor, lo que resulta en un mejor rendimiento en diferentes escenarios.

Pensamientos finales

En este artículo, hemos hablado sobre BrushNet, un marco de ingeniería dual de rama novedoso que incorpora características de imagen enmascarada a nivel de píxel en cualquier modelo de difusión preentrenado, garantizando así la coherencia y el resultado mejorado en tareas de inpainting de imágenes. El marco de BrushNet introduce un nuevo paradigma bajo el cual el marco divide las características de la imagen y el latente ruidoso en ramas separadas. La división de las características de la imagen y los latentes ruidosos disminuye drásticamente la carga de aprendizaje para el modelo y facilita una incorporación matizada de la información esencial de la imagen enmascarada de manera jerárquica. Además del marco de BrushNet, también hablaremos sobre BrushBench y BrushData que facilitan la evaluación del rendimiento basada en segmentación y el entrenamiento de inpainting de imágenes, respectivamente.