El relleno de imágenes es uno de los problemas clásicos en visión por computadora, y tiene como objetivo restaurar las regiones enmascaradas de una imagen con contenido plausible y natural. Los trabajos existentes que emplean técnicas tradicionales de relleno de imágenes, como Redes Adversarias Generativas (GAN) y AutoCodificadores Variacionales (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 el relleno de imágenes guiado por texto modificaron la estrategia de desenoiseo estándar mediante la muestra de las regiones enmascaradas de un modelo de difusión preentrenado, y las áreas no enmascaradas de la imagen dada. Aunque estos métodos resultaron en un desempeño satisfactorio en tareas simples de relleno de imágenes, 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 atribuir principalmente a su conocimiento perceptual limitado 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 relleno 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 relleno de imágenes implican modificar la estrategia de muestreo o el desarrollo de modelos de relleno de imágenes específicos, 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 hacia adelante para los modelos de relleno de imágenes, en este artículo, hablaremos sobre BrushNet, un marco de ingeniería de doble rama novel que incorpora características de imagen enmascarada a nivel de píxel en cualquier modelo de difusión preentrenado, garantizando así la coherencia y un resultado mejorado en tareas de relleno de imágenes. El marco de BrushNet introduce un nuevo paradigma en el que el marco divide las características de la imagen y los latentes ruidosos en ramas separadas. La división de las características de la imagen y los latentes ruidosos disminuye la carga de aprendizaje para el modelo de manera drástica, y facilita una incorporación matizada de 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 relleno 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, así como su comparación con marcos de estado del arte. Así que comencemos.

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

El relleno de imágenes, un método que intenta restaurar las regiones enmascaradas de una imagen mientras mantiene la coherencia general, ha sido un problema persistente en el campo de la visión por computadora, y ha molestado a desarrolladores e investigadores durante varios años. El relleno 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. La capacidad 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 relleno de imágenes de alta calidad que se alinean con los textos de entrada.

Los métodos empleados por los marcos de relleno de imágenes guiado por texto basados en difusión tradicionales se pueden dividir en dos categorías, Modificación de la estrategia de muestreo y Modelos de relleno de imágenes dedicados. El método de modificación de la estrategia de muestreo modifica el proceso de desenoiseo estándar mediante la muestra de las regiones enmascaradas de un modelo de difusión preentrenado, y copia y pega las áreas no enmascaradas de la imagen dada en cada paso de desenoiseo. Aunque los enfoques de modificación de la estrategia de muestreo se pueden implementar en modelos de difusión arbitrarios, a menudo resultan en resultados de relleno de imágenes 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 relleno de imágenes dedicados ajustan un modelo de relleno de imágenes diseñado específicamente ampliando las dimensiones del canal de entrada del modelo de difusión base para incorporar la imagen corrupta y las máscaras. Aunque los modelos de relleno de imágenes dedicados permiten que el modelo de difusión genere resultados más satisfactorios con modelos especializados en forma y contenido, puede que no sea el mejor diseño arquitectónico para los modelos de relleno de imágenes.

Como se demuestra en la siguiente imagen, los modelos de relleno de imágenes 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 relleno de imágenes dedicados influye fácilmente en las características de la imagen enmascarada y evita que las capas posteriores en la arquitectura 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 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 a 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 relleno de imágenes. Como resultado, marcos existentes como ControlNet producen resultados insatisfactorios en comparación con los modelos de relleno de imágenes 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 relleno 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 adaptarlas a la distribución de UNet de manera más efectiva.
2. El marco de BrushNet incorpora gradualmente la capa de características de UNet completa, capa por capa, en la arquitectura de UNet preentrenada, un enfoque jerárquico que permite un 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 modelo 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 doble rama 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 rellenada. 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 desenoiseo, 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 la 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 la exclusión de 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 capa de difusión congelada 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 perjudicial no tenga influencia sobre los estados ocultos en la copia trainable durante las etapas iniciales de entrenamiento.

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, lo que a menudo resulta en varias inexactitudes, y el marco de BrushNet enfrenta un problema similar cuando reduce la máscara para adaptar el 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 de 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 la adopción de 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 en varios modelos de difusión preentrenados, y permite una preservación flexible de la escala. 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 deseables. 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 permite una personalización fácil de la escala de preservación de las regiones no enmascaradas, dejando espacio para ajustes flexibles y un control fino sobre el proceso de relleno de imágenes.

BrushNet: Implementación y Resultados

Para analizar sus resultados, el marco de BrushNet propone BrushBench, un conjunto de datos de relleno de imágenes basado en segmentación con más de 600 imágenes, cada una acompañada de una máscara y una anotación de captura humana. Las imágenes en el conjunto de datos de referencia 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 relleno de imágenes, 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 segmentación y máscaras de pincel.

Comparación cuantitativa

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

Como se puede observar, el marco de BrushNet demuestra una eficiencia notable en la preservación de las regiones enmascaradas, la alineación de texto y la calidad de la imagen. Además, los modelos como Stable Diffusion Inpainting, HD-Painter, PowerPaint y otros demuestran un desempeño fuerte en tareas de relleno de imágenes dentro de la imagen, aunque no logran replicar su desempeño en tareas de relleno de imágenes fuera de la imagen, especialmente en términos de alineación de 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 relleno de imágenes basados en difusión existentes en el conjunto de datos de EditBench, y el desempeño es comparable al observado en el conjunto de datos de BrushBench. Los resultados indican que el marco de BrushNet entrega un desempeño fuerte en una amplia variedad de tareas de relleno 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 relleno de imágenes existentes, con resultados que cubren inteligencia artificial y imágenes naturales en diferentes tareas de relleno de imágenes, incluyendo relleno de máscaras aleatorias, relleno de máscaras de segmentación dentro y relleno de máscaras de segmentación fuera.

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, lo que resulta en un mejor desempeño en diferentes escenarios.

Pensamientos finales

En este artículo, hemos hablado sobre BrushNet, un marco de ingeniería de doble rama novel que incorpora características de imagen enmascarada a nivel de píxel en cualquier modelo de difusión preentrenado, garantizando así la coherencia y un resultado mejorado en tareas de relleno de imágenes. El marco de BrushNet introduce un nuevo paradigma en el que el marco divide las características de la imagen y los latentes ruidosos en ramas separadas. La división de las características de la imagen y los latentes ruidosos disminuye la carga de aprendizaje para el modelo de manera drástica, y facilita una incorporación matizada de 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 relleno de imágenes, respectivamente.