SmartBrush: Text and Shape Guided Object Inpainting with Diffusion Model

• 논문 링크: https://arxiv.org/pdf/2212.05034

0. Abstract

• 멀티모달 이미지 인페인팅 (Multi-modal Image Inpainting) → 인페인팅 되는 콘텐츠에 대해 더욱 유연하고 유용한 제어를 제공
  • e.g. 물체에 대한 더욱 풍성한 특징을 설명하기 위해 텍스트 프롬프트를 사용할 수 있으며 마스크 (Mask)는 인페인팅 되는 물체의 모양을 제한하는데 사용
• 본 논문에서는 새로운 확산 (Diffusion) 기반 모델인 SmartBrush를 제안 → 텍스트와 모양 (shape)에 대한 가이드를 사용하여 빈 영역을 물체로 채워줌
• DALL-E 2나 Stable Diffusion과 같은 이전의 기법들은 텍스트 기반 인페인팅은 가능하지만 모양에 대한 가이드는 제공하지 않아서 생성된 물체 주변의 배경 텍스쳐를 변경하는 등의 결과를 보임
• 본 논문의 모델은 텍스트와 모양에 대한 가이드를 통합하여 정확한 제어가 가능하도록 함
• 배경을 잘 보존하기 위해 새로운 학습과 샘플링 전략을 제안 → 물체의 마스크 예측을 통해 확산 모델의 U-net을 강화
• 마지막으로 멀티 태스크 학습 전략을 사용 → 인페인팅과 text-to-image 생성을 동시에 학습

1. Introduction

• 본 논문에서 텍스트 설명과 인페인팅 될 물체의 모양을 조건으로 멀티모달 물체 인페인팅을 수행하는 방법을 제안 → 그림 1 참고

• 확산 모델 (e.g. Stable Diffusion, DALL-E, Imagen)은 인페인팅 문제에 적용될 수 있음
  • 배경 영역의 랜덤 노이즈를 확산의 역 과정 (reverse process) 동안 원래 이미지의 노이지 한 버전으로 대체
  • 그러나 이는 모델이 샘플링 동안 전체적인 이미지의 맥락을 볼 수 없으므로 원치 않는 결과를 도출할 수 있음
  • 이를 해결하기 위해 GLIDE나 Stable Inpainting은 이미지의 일부를 랜덤하게 지우고 관련된 이미지 캡션을 조건으로 지워진 영역을 복원하도록 모델에 대한 파인튜닝을 수행
• 이미지 인페인팅에서 발생할 수 있는 문제들
  • Text Misalignment
    • 삭제된 영역 (로컬 컨텐츠)와 전역적인 텍스트 설명 간의 의미적인 불일치에 의해 텍스트 프롬프트를 따르는 것이 아니라 배경으로 마스킹 된 영역을 채우는 결과를 만들 수 있음
  • Mask Misalignment
    • 멀티모달 이미지 인페인팅을 수행하기 위한 한가지 방법은 강력한 언어-비전 모델을 사용하는 것 (e.g. CLIP)
    • Blended diffusion은 CLIP을 사용하여 이미지 임베딩과 입력 텍스트 임베딩 간의 차이를 계산 → 이 차이를 사전학습된 확산 모델의 샘플링 과정에 주입
    • 그러나 CLIP 모델은 전역적이고 고차원의 이미지 특징을 추출하도록 학습됨 → 주어진 마스크에 따라 물체를 생성하는 방향으로 학습되지 않음 (그림 1의 Blended Diffusion 참고)
  • Background Preservation
    • 그림 1의 아랫줄에서 볼 수 있듯이 인페인팅 된 물체의 주변 배경에 왜곡을 생성
• 이런 문제들을 해결하기 위해 입력 마스크의 정밀도 계수 (Precision Factor)를 사용
  • 마스크 만을 입력으로 하는 것이 아니라 해당 마스크의 모양이 인페인팅 된 물체와 얼마나 가까운지를 나타내는 정보를 사용
  • 이를 위해 정확한 인스턴스 마스크에 가우시안 블러를 적용하여 미세한 (Fine) 마스크부터 거친 (Coarse) 마스크까지 다양한 유형의 마스크를 생성하고 마스크와 정밀도 계수를 사용하여 확산 모델을 학습
  • 해당 세팅에서는 거친 마스크를 사용하는 경우 해당 마스크 내부 어딘가에 원하는 물체가 포함되도록 하며 미세한 마스크를 사용하는 경우 물체의 모양을 정확히 제공하도록 함
• 여기서 중요한 특성 하나는 바운딩 박스와 같이 거친 마스크의 경우 인페인팅 된 영역 내부의 배경이 원본 이미지의 배경을 유지해야 한다는 것
  • 이를 달성하기 위해 본 논문은 모델이 마스킹 영역을 인페인팅 하도록 할 뿐 아니라 정규화 손실 (Regularization loss)을 사용하여 생성할 물체의 인스턴스 마스크를 예측하도록 함
  • 테스트 시에는 거친 마스크를 예측된 마스크로 대체하여 샘플링 동안 최대한 배경이 유지되도록 함 (그림 1의 두번째 행 참고)
• 본 논문의 모델을 몇몇의 어려운 물체 인페인팅 문제에서 평가 → 본 논문의 기법이 몇몇 물체 인페인팅 데이터 셋에서 SOTA 성능 달성
  • DALL-E 2, Stable Inpainting 등과 성능 비교
• 본 논문의 contribution
  • 텍스트와 모양으로 가이드하는 물체 인페인팅 확산 모델을 제안 → 다른 정밀도의 물체 마스크를 조건으로 하여 물체 인페인팅을 위한 새로운 단계의 제어 달성
  • 거친 입력 마스크에서 이미지 배경을 유지하기 위해 모델은 물체의 마스크를 예측하도록 학습됨 → 합성된 물체의 주변을 기존 배경으로 유지하도록 함
  • 멀티 태스크 학습 전략 사용 → 물체 인페인팅과 text-to-image 생성을 동시에 학습

2. Preliminary: Diffusion Model

• 입력 이미지 $x_0$가 주어졌을 때 순방향 확산 Markov process를 통해 이미지에 scheduled variance $\beta_t$를 통해 시간 스텝 $t$만큼 노이즈를 추가

 

• $T$는 전체 스텝의 수
• 만약 $T \rightarrow \infin$이면 출력 $x_t$는 등방성 가우시안 (Isotropic Gaussian)이 됨

 

• 정의된 Markov process는 우리가 $x_t$를 closed form으로 얻을 수 있도록 함

 

• 랜덤 노이즈로부터 이미지를 생성하기 위해 확산의 과정을 역방향으로 수행하도록 학습
  • 신경망 네트워크 $p_{\theta}$를 학습하여 조건부 분포를 근사하도록 함

• • $\mu_{\theta}$는 식 (2)에서 얻어진 $x_{t-1} = \frac{1}{\sqrt{a_t}} (x_t - \frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t} \epsilon_t})$를 예측하도록 학습됨
• 우리는 이미 학습 중에 $x_t$를 가지고 있으므로 $\mu_\theta$를 학습하는 대신 $\epsilon_t$를 예측하기 위해 $\epsilon_{\theta}$를 학습할 수 있음 → 확산 모델 학습을 위한 목적함수가 아래와 같음

• 테스트시에는 랜덤 노이즈 $x_T \sim \mathcal{N}(0,I)$부터 시작해서 반복적으로 모델 $\epsilon_\theta$를 적용하여 $t=0$이 될 때까지 $x_t$로부터 $x_{t-1}$을 취득
• DDIM, PNDM 등 더욱 효율적인 샘플링 기법을 사용하여 샘플링 속도를 증가시키고 classifier free guidance를 사용하여 샘플의 품질을 향상시킴
• 조건부 확산 모델 (e.g. text-to-image 와 인페인팅 모델)의 조건부 정보는 손실 함수의 변화 없이 네트워크 $\epsilon_{\theta}$에 입력될 수 있음
  • 모델은 더 높은 품질의 조건부 이미지를 생성하기 위해 조건을 잘 사용하도록 학습됨

3. Our Approach

• 이미지 $x$, 텍스트 프롬프트 $d$, 바이너리 마스크 $m$ ($x$에서 변경되어야 할 영역)이 주어졌을 때 우리의 목표는 이미지 $\widetilde{x}$를 생성하는 것 → $\widetilde{x}$는 입력 $x$와 동일해야함
• 마스크 영역 안에 생성된 $\widetilde{x} \odot m$은 텍스트 프롬프트 $d$와 마스크 $m$과 잘 관련되어야함

3.1. Text and Shape Guided Diffusion

• 본 논문의 기법은 이미지 내에 존재하는 인스턴스로부터의 텍스트나 모양 정보를 사용
  • 해당 데이터셋들은 주석이 달린 마스크들 $\{m_i\}_{i=1}^N$을 포함하며 숫자 $N$은 주석의 수이고 각 마스킹 영역 $x \odot m_i$는 하나의 물체만을 포함
  • 각 마스크에 대해서는 관련된 클래스 라벨을 가짐 → e.g. 모자, 고양이, …
• 순방향 과정에서는 이미지 $x$에 대해 랜덤하게 segmentation 마스크 $m$을 그리고 이에 대한 클래스 텍스트인 라벨 $c$를 정의
• $x_0 = x$로 정의하고 모든 픽셀이 아니라 마스크 영역에만 노이즈를 추가

• $\epsilon \sim \mathcal{N} (0,I)$이고 $t$는 순방향 과정의 시간 스텝

 

• $x_t, m, c$를 모델의 입력으로 하여 모델이 배경의 정보를 사용하면서 마스킹된 영역 $x_0 \odot m$에 대상을 생성하도록 함 → 노이지한 이미지인 $x_t$로부터 노이즈 $\epsilon$를 예측하도록 네트워크 $\epsilon_\theta$를 학습

• 추론시에는 마스킹 영역에 랜덤 가우시안 노이즈를 생성 → $x_T = \epsilon \odot m + x_0 \odot (1-m)$
  • $T$는 샘플링 스텝의 수
• 그리고 확산 과정을 역으로 수행하며 인페인팅 된 결과 $x_0$를 학습

3.2. Shape Precision Control

• 만약 정확한 인스턴스 마스크를 학습을 위한 마스크로 사용하는 경우 모델은 정확히 입력 마스크의 모양만을 따르도록 학습됨
• 사용자가 정확한 마스크나 (e.g. 고양이의 모양을 정확히 따르는 마스크) 거친 마스크 (e.g. 바운딩 박스)를 사용하는 모든 경우를 커버하기 위해 다른 정밀도로 마스크를 생성하고 사용
• 이를 위해 학습 중에 랜덤하게 원본 마스크의 형태를 변형
• 정확한 인스턴스 마스크 $m$이 있을 때 마스크 정밀도 지표 $s \sim [0, S]$를 사용하고 해당 지표를 위한 파라미터 정의

• $k_s$는 가우시안 커널 사이즈를 나타내며 $\sigma_s$는 커널의 표준편차를 나타냄
• 만약 $s=0$이면 마스크는 변경되지 않고 정확한 인스턴스 마스크를 사용
• 만약 $s=S$이면 마스크 $m_s$는 인스턴스 마스크 $m$에 대한 바운딩 박스가 될 것 → 상세한 모양의 정보를 잃게 됨

 

• 학습 중에 각 학습 샘플 (물체)을 위해 세밀한 마스크부터 거친 마스크까지 마스크의 세트 $\{m_s, s\}$를 사용하고 정밀도 지표 $s$로 확산 모델의 조건으로 사용

 

• 마스크의 샘플들은 그림 5에서 살펴볼 수 있음

3.3. Background Preservation

• 추론 동안 확산 모델은 마스킹 된 이미지에 대해 디노이징을 수행하고 주어진 텍스트 프롬프트에 따라 물체를 생성
• 결과적으로 입력 마스크가 거친 경우 마스킹 영역의 배경이 변하게 됨
  • 예를 들어 네모난 마스킹 영역 안에 고양이를 만든다고 할 때 네모 박스 영역안에 생성할 고양이 외에 다른 픽셀들도 변함
  • 이상적으로는 배경을 보존해야하지만 거친 마스크 내부의 어디에 모델이 원하는 물체를 생성해야하는지 알 수 없음
• 이 문제를 해결하기 위해 마스크 정밀도 정보를 사용
  • 확산 네트워크가 거친 마스크 $m_s$로부터 정확한 인스턴스 마스크 $m$을 예측하도록 확산 네트워크를 학습

• • $H$는 segmentation을 위한 적절한 어떤 기준도 될 수 있으며 본 논문에서는 DICE 손실함수를 사용
  • $i.e. H(X,Y)=1-\frac{2|X \bigcap Y|}{|X| + |Y|}$
  • 확산 모델에 추가적인 출력채널을 추가 → 인스턴스 마스크의 예측을 포함하도록 함
• 추론 동안 확산 모델의 예측을 통해 거친 마스크 $m_s$ 내부에 물체를 예측하도록 함
  • 먼저 거친 마스크 $m_s$를 확산 모델에 입력하고 디노이징 수행을 위해서 예측 마스크를 도출
  • 예측된 마스크로 마스크 영역 내 어디에 물체가 생성될지 알 수 있다면 생성될 물체 주변의 배경을 보존하는데 도움이 될 수 있음

3.4. Training Strategy

• 식 (8), (9)를 결합하여 최종 학습 목적 함수를 다음과 같이 정의

• $\lambda$는 두 손실함수 간의 균형을 설정하는 하이퍼파라미터로 본 논문에서는 $\lambda=0.01$로 설정

 

• 본 모델은 사전학습 된 text-to-image 생성 모델을 기반으로 함 (e.g. Stable Diffusion, Imagen) → 학습 속도 향상 (Stable Diffusion text-to-image 모델 v1.2를 기반으로 파인튜닝 수행)
• 조건을 사용하며 (그림 2) 손실함수는 $\mathcal{L}_{total}$ (식 10)을 사용
• 모델이 단일 단어와 짧은 문장 텍스트 모두에 대응할 수 있도록 적용
  • 텍스트 불일치 (text misalignment)를 피하고 로컬 마스크 컨텐츠와 텍스트 설명을 일치시키기 위해 segmentation과 이에 해당하는 라벨을 사용
  • 하지만 경험적으로 이런 카테고리 텍스트만을 사용하는 것은 성능 저하의 원인이 될 수 있음
  • 그러므로 BLIP 모델을 사용 → 이미지에 대해 더 풍부하고 긴 캡션을 사용
  • 학습 동안 랜덤하게 해당 마스크에 대해 segmentation 라벨이나 BLIP 캡션을 사용
  • 이에 따라 모델은 단일 단어 텍스트와 짧은 문장 텍스트를 모두 추론에 사용할 수 있음
• 멀티 태스크 학습
  • 주요한 인페인팅 문제와 기반이 되는 text-to-image 생성 문제를 동시에 학습하는 멀티 태스크 학습 전략을 제안
  • LAION-Aesthetics v2 5+ subset의 이미지/텍스트 쌍 사용
  • text-to-image를 위해 입력 마스크가 전체 이미지를 덮도록 하고 이를 특별한 형태의 인페인팅 케이스로 정의

4. Experimental Evaluation

4.1. Experimental Setup

• 식 10의 전체 손실함수에서 $\lambda=0.01$로 설정하고 배치 크기는 1024로 설정
• 인페인팅 문제와 text-to-image 생성 문제의 비율은 각각 80% - 20%로 설정
• A100 GPU 8개로 20K 스텝만큼 학습 (Stable Inpainting의 경우 A100 GPU 256개로 440K 스텝만큼 학습했음)

Baselines

• 현재 최고 성능의 이미지 인페인팅 기법들을 베이스라인으로 사용 → Blended Diffusion, GLIDE, Stable Diffusion, Stable Inpainting, DALL-E 2

Testing Datasets

• 두개의 유명한 segmentation 데이터셋들 통해 성능 평가 → OpenImages, MSCOCO
• MSCOCO 테스트 데이터의 각 이미지에서 2개의 마스크를 샘플링 → 총 테스트 이미지의 수는 9,311개
• OpenImages에서는 512 보다 큰 해상도의 이미지들만 사용하고 1개의 마스크 사용 → 13,400개
• 입력 프롬프트는 segmentation 클래스 라벨을 그대로 사용

Evaluation Metrics

• 첫번째 성능 평가 지표 → Frechet Inception Distance (FID)
  • 하지만 전역적인 (Global) FID는 이미지의 일부인 마스크 영역에 대한 생성 성능 평가에는 적절하지 않음
  • 이에 따라 바운딩 박스로 이미지를 자르고 지역적인 부분의 FID를 계산하는 Local FID를 사용
• 두번째 성능 지표 → CLIP Score
  • 텍스트와 이미지의 일치도 평가

4.2. Text and Shape Guided Inpainting

• 제안된 SmartBrush는 텍스트와 모양에 대한 가이드를 기반으로 물체 뿐 아니라 “석양에 물든 하늘”과 같은 일반적인 장면도 인페인팅 가능
• 물체 인페인팅에 대한 2가지 사용 예시를 고려
  1. 정확한 물체 마스크
  2. 바운딩 박스 마스크
• 1의 경우 주어진 마스크의 모양대로 물체를 생성할 것으로 기대되고 2의 경우 물체가 박스 안에만 있도록 하고 모양은 규정하지 않음
• 정량적인 결과는 표 1, 표 2를 통해 살펴볼 수 있음

 

• SmartBrush가 두 데이터셋에 대해 모든 평가 기준에서 가장 뛰어난 성능을 보임
• 그림 3은 베이스라인 모델들과 SmartBrush의 인페인팅 예시를 시각화 한 결과

 

• Amazon Mechanical Turk를 통해 유저 스터디 수행
  • 300명 이상의 작업자가 다음에 질문들에 대해 응답
    • 어떤 결과가 물체의 마스크를 가장 잘 따랐는가
    • 어떤 결과가 입력 텍스트 설명을 가장 잘 따랐는가?
    • 어떤 결과가 가장 자연스럽고 실제적으로 보이는가?
  • 조사의 결과는 그림 4에서 살펴볼 수 있음 
    • 50% 이상의 사용자가 각 질문에 대해 본 논문의 모델을 최고로 선정

4.3. Mask Precision Control

• 실생활 문제에서 사용자들은 인페인팅을 하고 싶은 물체의 마스크를 정확하게 제공하지 않는 경우가 많음
• 이에 따라 SmartBrush는 그림 5와 같이 다양한 형태의 마스크를 사용 → 가우시안 블러를 반복적으로 적용해서 세밀한 마스크부터 거친 마스크까지 얻을 수 있음

• Stable Diffusion의 경우 이와 같은 방식으로 학습되지 않았으므로 마스크의 종류에 영향을 받지 않고 일관적으로 결과가 좋지 않음
• Stable Inpainting은 마스크의 크기에 따라 물체의 크기만 바뀌고 마스크의 모양은 따르지 않음
• 본 논문의 기법은 세밀한 마스크를 사용할 때는 마스크의 모양을 엄격하게 따르고 거친 마스크가 주어진 경우 대략적으로만 마스크의 모양을 따름
• 마지막 열처럼 박스 모양의 마스크가 주어졌을 때 본 논문의 기법은 박스 어딘가에 생성을 수행

4.4. Background Preservation

• 물체를 인페인팅하기 위해서, 특히 박스 모양의 마스크가 주어졌을 때 인페인팅 될 물체는 마스크 영역의 일부만을 차지하기 떄문에 배경을 보존하는 것이 중요함
• 그림 6은 박스 모양의 마스크가 주어졌을 때 배경을 보존하는 다른 기법들의 결과를 비교

• 배경 유지를 위한 정규화를 제외한 경우 (w/o bg pres)와 DALL-E 2는 마스크 내부에서 물체가 아닌 픽셀들에 변화가 발생
• SmartBrush의 경우 물체 마스크 예측을 사용하여 (그림 7 참고) 샘플링 동안 예측된 마스크를 사용하여 더욱 나은 배경 유지를 수행

4.5. Ablation Study

• 본 논문의 요소들을 하나씩 제거하며 MSCOCO 데이터의 테스트 셋과 바운딩 박스 마스크로 테스트 수행 → 표3 참고

• 제안된 배경 유지 기법을 사용하는 경우 더 많은 픽셀들이 보존되면서 정량적인 결과가 향상된 것을 확인
• 만약 마스크 정밀도 컨디션을 적용하지 않는 경우 생성된 대상이 잘 제어 되지 않으며 Local FID (LFID)의 하락을 유발
• BLIP이나 멀티 태스크 학습이 없는 경우 발생하는 성능 저하 또한 해당 기법들이 물체 인페인팅 문제에 유용하다는 것을 보임

5. Conclusion, Limitation and Future Work

• 기존의 텍스트와 모양 기반 이미지 인페인팅 모델은 3가지 문제점을 가짐
  • Mask misalignment, text misalignment, background preservation
• 본 논문에서는 새로운 학습 방식을 제안
  • 텍스트와 모양에 대한 가이드를 통해 text misalignment 문제를 해결
  • 다른 단계의 마스크들을 만드는 방법을 사용하여 생성의 제어 정밀도를 다양하게 함
  • 마지막으로 추가적인 학습 손실 함수를 도입하여 모델이 입력 마스크에서 물체의 마스크를 예측하도록 하고 이를 통해 마스크에 불필요한 변화가 생기지 않도록 함
• 본 기법의 한계점
  • 큰 그림자가 있는 경우 → 물체의 그림자가 물체의 마스크를 초과하는 경우 (e.g. 사람의 그림자가 길어서 바운딩 박스가 모든 그림자를 커버하지 못함)