A recriação de imagens é um dos problemas clássicos da visão computacional, e visa restaurar as regiões mascaradas de uma imagem com conteúdo plausível e natural. O trabalho existente que emprega técnicas tradicionais de recriação de imagens, como Redes Adversárias Generativas (GANs) e Autoencoders Variacionais (VAEs), muitas vezes requer recursos auxiliares projetados manualmente, mas não entrega resultados satisfatórios. Nos últimos anos, os métodos baseados em difusão ganharam popularidade dentro da comunidade de visão computacional devido às suas capacidades notáveis de geração de imagens de alta qualidade, diversidade de saída e controle fino. As tentativas iniciais de empregar modelos de difusão para recriação de imagens orientadas por texto modificaram a estratégia de desruído padrão, amostrando as regiões mascaradas de um modelo de difusão pré-treinado e as áreas não mascaradas da imagem dada. Embora esses métodos tenham resultado em desempenho satisfatório em tarefas simples de recriação de imagens, eles lutaram com formas de máscara complexas, prompts de texto e conteúdo de imagem que resultaram em uma falta geral de coerência. A falta de coerência observada nesses métodos pode ser atribuída principalmente ao seu conhecimento perceptual limitado das fronteiras da máscara e do contexto da região da imagem não mascarada.

Apesar dos avanços, pesquisas e desenvolvimentos desses modelos nos últimos anos, a recriação de imagens ainda é um grande obstáculo para os desenvolvedores de visão computacional. As adaptações atuais de modelos de difusão para tarefas de recriação de imagens envolvem a modificação da estratégia de amostragem ou o desenvolvimento de modelos de recriação de imagens específicos, que muitas vezes sofrem de redução na qualidade da imagem e semântica inconsistente. Para enfrentar esses desafios e pavimentar o caminho para os modelos de recriação de imagens, neste artigo, vamos falar sobre o BrushNet, um framework novo e inovador de engenharia de duplo ramo que incorpora recursos de imagem mascarada em nível de pixel em qualquer modelo de difusão pré-treinado, garantindo coerência e resultados aprimorados em tarefas de recriação de imagens. O framework BrushNet introduz um novo paradigma no qual o framework divide as características da imagem e os latentes barulhentos em ramos separados. A divisão das características da imagem e dos latentes barulhentos diminui drasticamente a carga de aprendizado para o modelo e facilita a incorporação sutil de informações essenciais da imagem mascarada de forma hierárquica. Além do framework BrushNet, também vamos falar sobre o BrushBench e o BrushData, que facilitam a avaliação de desempenho baseada em segmentação e o treinamento de recriação de imagens, respectivamente.

Este artigo visa cobrir o framework BrushNet em profundidade e explorar o mecanismo, a metodologia, a arquitetura do framework, bem como sua comparação com frameworks de ponta. Vamos começar.

BrushNet: Recriação de Imagens com Difusão de Duplo Ramo

A recriação de imagens, um método que tenta restaurar as regiões faltantes de uma imagem enquanto mantém a coerência geral, tem sido um problema de longa data no campo da visão computacional e tem incomodado desenvolvedores e pesquisadores por alguns anos. A recriação de imagens encontra aplicações em uma ampla variedade de tarefas de visão computacional, incluindo edição de imagens e provas virtuais. Recentemente, modelos de difusão, como o Stable Diffusion e o Stable Diffusion 1.5, demonstraram capacidades notáveis de geração de imagens de alta qualidade e oferecem aos usuários a flexibilidade de controlar os controles semânticos e estruturais. O potencial notável dos modelos de difusão é o que tem levado os pesquisadores a recorrer a modelos de difusão para tarefas de recriação de imagens de alta qualidade que se alinham com os prompts de texto.

Os métodos empregados por frameworks de recriação de imagens orientadas por texto baseados em difusão tradicionais podem ser divididos em duas categorias, Modificação da Estratégia de Amostragem e Modelos de Recriação de Imagens Dedicados. O método de modificação da estratégia de amostragem modifica o processo de desruído padrão, amostrando as regiões mascaradas de um modelo de difusão pré-treinado e copiando as áreas não mascaradas da imagem dada em cada etapa de desruído. Embora as abordagens de modificação da estratégia de amostragem possam ser implementadas em modelos de difusão arbitrários, elas muitas vezes resultam em resultados de recriação de imagens incoerentes, pois têm conhecimento perceptual limitado das fronteiras da máscara e do contexto da região da imagem não mascarada. Por outro lado, os modelos de recriação de imagens dedicados aprimoram um modelo de recriação de imagens projetado especificamente, expandindo as dimensões do canal de entrada do modelo de difusão base para incorporar a imagem corrompida e as máscaras. Embora os modelos de recriação de imagens dedicados permitam que o modelo de difusão gere resultados mais satisfatórios com modelos especializados e conscientes de forma e conteúdo, eles podem ou não ser a melhor arquitetura para modelos de recriação de imagens.

Como demonstrado na imagem a seguir, os modelos de recriação de imagens dedicados fundem o latente da imagem mascarada, o latente barulhento, o texto e a máscara em uma etapa inicial. A arquitetura dos modelos de recriação de imagens dedicados influencia facilmente as características da imagem mascarada e impede que as camadas subsequentes na arquitetura UNet obtenham características puras da imagem mascarada devido à influência do texto. Além disso, lidar com a geração e a condição em um único ramo impõe uma carga adicional à arquitetura UNet, e como essas abordagens também requerem ajuste fino em diferentes variações do backbone de difusão, elas são frequentemente exaustivas em termos de tempo e têm limitada transferibilidade.

Pode parecer que adicionar um ramo adicional dedicado a extrair características da imagem mascarada pode ser uma solução adequada para os problemas mencionados acima, no entanto, os frameworks existentes muitas vezes resultam em extrair e inserir informações inadequadas quando aplicados diretamente à recriação de imagens. Como resultado, frameworks existentes, como o ControlNet, produzem resultados insatisfatórios quando comparados a modelos de recriação de imagens dedicados. Para enfrentar esse problema de forma eficaz, o framework BrushNet introduz um ramo adicional à rede de difusão original e cria uma arquitetura mais adequada para tarefas de recriação de imagens. A arquitetura e o design do framework BrushNet podem ser resumidos em três pontos.

1. Em vez de inicializar camadas de convolução aleatoriamente, o framework BrushNet implementa um codificador VAE para processar a imagem mascarada. Como resultado, o framework BrushNet consegue extrair as características da imagem para adaptação à distribuição UNet de forma mais eficaz.
2. O framework BrushNet incorpora gradualmente a camada de características UNet completa, camada por camada, na arquitetura UNet pré-treinada, uma abordagem hierárquica que permite controle denso por pixel.
3. O framework BrushNet remove a atenção cruzada de texto da componente UNet para garantir que apenas informações de imagem puras sejam consideradas no ramo adicional. Além disso, o modelo BrushNet também propõe implementar uma estratégia de mistura borrosa para obter melhor consistência, juntamente com um maior controle nas regiões não mascaradas da imagem.

BrushNet: Método e Arquitetura

A figura a seguir fornece uma visão geral do framework BrushNet.

Como pode ser observado, o framework emprega uma estratégia de duplo ramo para a inserção de orientação de imagem mascarada e usa operações de mistura com uma máscara borrosa para garantir melhor preservação das regiões não mascaradas. É importante notar que o framework BrushNet é capaz de ajustar a escala adicionada para alcançar controle flexível. Para uma imagem de entrada mascarada e a máscara, o modelo BrushNet produz uma imagem recriada. O modelo primeiro reduz a escala da máscara para acomodar o tamanho do latente e a imagem mascarada é alimentada como entrada no codificador VAE para alinhar a distribuição do espaço latente. O modelo então concatena o latente da imagem mascarada, o latente barulhento e a máscara reduzida e usa como entrada. As características que o modelo extrai são adicionadas à camada UNet pré-treinada após um bloco de convolução zero. Após o desruído, o modelo mistura a imagem mascarada e a imagem gerada com uma máscara borrosa.

Orientação de Imagem Mascarada

O framework BrushNet insere a característica da imagem mascarada na rede de difusão pré-treinada usando um ramo adicional, que separa explicitamente a extração de características da imagem mascarada do processo de geração de imagem. A entrada é formada pela concatenação do latente da imagem mascarada, do latente barulhento e da máscara reduzida. Para ser mais específico, o latente barulhento fornece informações para a geração de imagem durante o processo de geração atual e ajuda o framework a melhorar a coerência semântica da característica da imagem mascarada. O framework BrushNet então extrai o latente da imagem mascarada da imagem mascarada usando um Autoencoder Variacional. Além disso, o framework emprega interpolação cúbica para reduzir a escala da máscara em uma tentativa de garantir que o tamanho da máscara esteja alinhado com o latente da imagem mascarada e o latente barulhento. Para processar as características da imagem mascarada, o framework BrushNet implementa uma cópia do modelo de difusão pré-treinado e exclui as camadas de atenção cruzada do modelo de difusão. O motivo é que os pesos pré-treinados do modelo de difusão servem como uma forte prioridade para a extração das características da imagem mascarada e a exclusão das camadas de atenção cruzada garante que o modelo considere apenas informações de imagem puras dentro do ramo adicional. O framework BrushNet insere as características no modelo de difusão congelado camada por camada, permitindo controle denso por pixel hierárquico e também emprega camadas de convolução zero para estabelecer uma conexão entre o modelo BrushNet treinável e o modelo bloqueado, garantindo que o ruído prejudicial não tenha influência sobre os estados ocultos no modelo treinável durante as etapas iniciais de treinamento.

Operação de Mistura

Como mencionado anteriormente, realizar a operação de mistura no espaço latente redimensiona as máscaras, o que muitas vezes resulta em várias imprecisões, e o framework BrushNet encontra um problema semelhante quando redimensiona a máscara para combinar com o tamanho do espaço latente. Além disso, é importante notar que as operações de codificação e decodificação nos Autoencoders Variacionais têm operações limitadas inerentes e podem não garantir a reconstrução completa da imagem. Para garantir que o framework reconstrua uma imagem consistente e completa da região não mascarada, trabalhos existentes implementaram diferentes técnicas, como copiar as regiões não mascaradas da imagem original. Embora a abordagem funcione, ela muitas vezes resulta em falta de coerência semântica na geração dos resultados finais. Por outro lado, métodos como a adoção de operações de mistura latente enfrentam dificuldades em preservar as informações desejadas nas regiões não mascaradas.

Controle Flexível

A arquitetura do framework BrushNet o torna uma escolha adequada para integrações plug and play inerentemente em vários modelos de difusão pré-treinados e permite preservação de escala flexível. Como o framework BrushNet não altera os pesos do modelo de difusão pré-treinado, os desenvolvedores têm a flexibilidade de integrá-lo como um componente plug and play com um modelo de difusão ajustado, permitindo adoção e experimentação fáceis com modelos pré-treinados. Além disso, os desenvolvedores também têm a opção de controlar a escala de preservação das regiões não mascaradas, incorporando as características do modelo BrushNet no modelo de difusão congelado com o peso w dado, que determina a influência do framework BrushNet na escala de preservação, oferecendo aos desenvolvedores a capacidade de ajustar os níveis desejados de preservação. Finalmente, o framework BrushNet permite que os usuários ajustem a escala de borramento e decidam se implementar ou não a operação de borramento, facilitando assim ajustes flexíveis e controle fino sobre o processo de recriação de imagens.

BrushNet: Implementação e Resultados

Para analisar seus resultados, o framework BrushNet propõe o BrushBench, um conjunto de dados de recriação de imagens baseado em segmentação com mais de 600 imagens, cada uma acompanhada de uma máscara anotada por humanos e anotação de legenda. As imagens no conjunto de dados de referência são distribuídas uniformemente entre imagens naturais e artificiais e também garantem distribuição uniforme entre diferentes categorias, permitindo uma avaliação justa em diferentes categorias. Para aprimorar a análise das tarefas de recriação de imagens ainda mais, o framework BrushNet categoriza o conjunto de dados em duas partes distintas com base nos métodos utilizados: segmentação baseada em recriação e máscaras de pincel.

Comparações Quantitativas

A tabela a seguir compara o framework BrushNet com modelos de recriação de imagens baseados em difusão existentes no conjunto de dados BrushBench com o Stable Diffusion como modelo base.

Como pode ser observado, o framework BrushNet demonstra eficiência notável em termos de preservação de regiões mascaradas, alinhamento de texto e qualidade de imagem. Além disso, modelos como o Stable Diffusion Inpainting, o HD-Painter, o PowerPaint e outros demonstram desempenho forte em tarefas de recriação de imagens internas, embora falhem em replicar seu desempenho em tarefas de recriação de imagens externas, especialmente em termos de alinhamento de texto e qualidade de imagem. No geral, o framework BrushNet entrega os resultados mais fortes.

Além disso, a tabela a seguir compara o framework BrushNet com modelos de recriação de imagens baseados em difusão existentes no conjunto de dados EditBench, e o desempenho é comparável ao observado no conjunto de dados BrushBench. Os resultados indicam que o framework BrushNet entrega desempenho forte em uma ampla variedade de tarefas de recriação de imagens com diferentes tipos de máscaras.

Comparações Qualitativas

A figura a seguir compara qualitativamente o framework BrushNet com métodos de recriação de imagens existentes, com resultados que abrangem imagens de inteligência artificial e naturais em diferentes tarefas de recriação de imagens, incluindo recriação de imagens com máscaras aleatórias, recriação de imagens com máscaras de segmentação internas e recriação de imagens com máscaras de segmentação externas.

Como pode ser observado, o framework BrushNet entrega resultados notáveis na coerência da região não mascarada e nas regiões coerentes e consegue realizar com sucesso a conscientização das informações de fundo devido à implementação da abordagem de desacoplamento de duplo ramo. Além disso, o ramo não tocado do modelo de difusão pré-treinado também fornece a vantagem de cobrir diferentes domínios de dados, como anime e pintura, resultando em melhor desempenho em diferentes cenários.

Pensamentos Finais

Neste artigo, falamos sobre o BrushNet, um framework novo e inovador de engenharia de duplo ramo que incorpora recursos de imagem mascarada em nível de pixel em qualquer modelo de difusão pré-treinado, garantindo coerência e resultados aprimorados em tarefas de recriação de imagens. O framework BrushNet introduz um novo paradigma no qual o framework divide as características da imagem e os latentes barulhentos em ramos separados. A divisão das características da imagem e dos latentes barulhentos diminui drasticamente a carga de aprendizado para o modelo e facilita a incorporação sutil de informações essenciais da imagem mascarada de forma hierárquica. Além do framework BrushNet, também vamos falar sobre o BrushBench e o BrushData, que facilitam a avaliação de desempenho baseada em segmentação e o treinamento de recriação de imagens, respectivamente.