[논문 리뷰] ShadowDiffusion: When Degradation Prior Meets Diffusion Model for Shadow Removal (2022)

 

 

기존의 딥러닝 방법으로 복원된 이미지는 degradation prior embedding과 모델링 용량의 부족으로 인해 여전히 boundary에 대한 어려움을 겪는다. 이 논문에서는 이를 해결하기 위해 image와 degradation을 모두 통합하는 unified diffusion framework를 제안한다. 먼저, ShadowDiffusion이라는 shadow degradation 모델을 제안한다. 이를 통해 복원을 위한 강력한 baseline 역할이 가능한 degradation prior과 diffusive generative prior를 모두 사용하여 과정을 개선시킬 수 있다. 또한, ShadowDiffusion은 추정된 shadow mask를 diffusion generator에 점진적으로 세분화하여 추가시켜줌으로써 그림자 제거 즉, clean image 생성에 성능을 높인다. 

 

#  Introduction 

그림자 제거 방법의 기존 방법의 대부분은 그림자 제거를 위한 discriminative model들을 학습하는 데 중점을 두었지만, natural image의 기본 분포를 모델링하는 것은 복원 과정에서 간과되고 있다. 따라서 그림자 제거 결과는 일반적으로 심각한 경계 아티팩트와 나머지 그림자 패턴을 포함한다. GAN과 같은 adversarial loss가 이러한 문제를 완화할 수 있지만, 이러한 모델들은 굉장히 까다로운 fine tuning이 필요하며, 그림자와 나머지 부분의 경계를 구분하기 힘들게 만든다.  반대로, DDPM과 같은 다양한 확산 모델은 이미지 픽셀 분포를 모델링하는 데 더 강력하며, 이는 시각적 품질을 크게 향상시킬 수 있는 큰 잠재력을 제공하고 고품질 이미지 복원에 도움이 된다. 그러나 이 논문 이전에 그림자 제거 작업에 diffusion 모델이 사용된 연구는 없다. 

기존 그림자 제거 방법의 두 가지 주요 제한은 다음과 같다. 

 

   1. 물리적 특성을 반영하는 shadow degradation prior는 활용되지 않음 

          - 복잡한 빛(광원)에 의한 그림자를 선형적이고 균일한 degradation으로 통합하는 것은 제한적이다. 

   2. 대부분의 모델은 사전학습모델이나 데이터 세트에 의해 추정된 그림자 마스크를 입력으로 사용해야한다. 

          - 일반적으로 예측된 그림자 마스크는 부정확 

#  Contributions

   1. 그림자 제거 작업을 위한 첫 번째 Diffusion 기반 모델을 제안한다.  

          -  Shadow-free image와 refined shadow mask를 구하기 위한 새로운 dynamic mask-aware diffusion model(DMDM)을 도입한다.

   2. Shadow degradation prior을 DMDM의 반복적 과정에 통합하는 확산 샘플링 전략을 제안한다. 

 

   3. 다양한 그림자 데이터셋에서 성능을 큰 폭으로 능가하였으며, 이후 다양한 이미지 향상 작업으로 일반화 될 수 있다. 

 

#  ShadowDiffusion 

먼저, shadow degradation model을 제안한다. 이후, dynamic mask-aware diffusion model(DMDM)과 train 과정을 제안하며, 이 과정에서 점진적인 shadow mask refinement를 통해 shadow-free image를 예측할 수 있다. 마지막으로, 확산 생성 모델과 shadow degradation prior를 통합한 DMDM을 기반으로 unrolling-inspired diffusive sampling 방법을 제안한다.

# - 1  Shadow Degradation Model 

이미지 y의 그림자 영역은 부분적 또는 완전한 폐색(occlusion)에 의해 발생한다. 일반적으로 조도 저하(illumination degradation)는 공간적으로 다양해야 하며 그림자 마스크 정보에 크게 의존해야 한다. 따라서, 다음과 같은 새로운 shadow degradation model을 제안한다.

여기서, h는 픽셀 단위의 illumination degradation map을 나타내며, 이는 shadow mask m과 illumination weight w로 분해될 수 있다. shodow mask m은 shadow region이 1이고 나머지가 0인 shadow locations을 나타낸다.

해당 shadow degradation model의 이점은 아래와 같다.

• shadow image y는 픽셀 단위 degradation map h에서 shadow-free image x의 불균일한 조명 변환으로 모델링될 수 있다. 이러한 degradation은 그림자 제거를 위한 균일한 degradation보다 더 풍부한 정보를 제공한다.
• shadow mask m은 shadow degradation h에 직접적이고 중요한 영향을 미치는 그림자 위치 정보를 제공하며, 따라서 원하는 shadow-free image x^의 추정에 상당한 영향을 미친다. 제안된 shadow degradation model은 보조 작업으로 mask refinement를 포함할 수 있고, refinement와 shadow-free image 복원을 서로 유익하게 만든다

# - 2  Dynamic Mask-Aware Diffusion Model (DMDM)

#-1를 통해 shadow region의 정확한 위치를 파악할 수 있다. 이후, shadow-free image와 refined mask를 추구하는 공동 작업으로 shadow removal 작업을 리모델링한다. 이 작업에서 mask refinement는 shadow-free image generation과 함께 점진적으로 shadow mask를 정제하기 위한 diffusion generator의 보조 작업이 된다. 불변 조건의 기본 이미지를 생성하는 이전의 conditional diffusion-based image restoration works와 다르게, shadow-free image와 refined mask(조건)을 점진적으로 생성하는 dynamic mask-aware diffusion model(DMDM)을 제안한다.

먼저, 기존의 conditional diffusion model을 다시 보면, x에 대한 diffusion process q(x_1:T |x_0) 를 수정하지 않고 conditional reverse process pθ(x_0:T | y)를 학습함으로써, sample image는 y에 대한 데이터 분포에 대해 높은 충실도(fidelity)를 갖게 된다.  

 

DMDM의 학습 중에서는, triplet data distribution(shadow-free image x, shadow image y, initial shadow mask m`)으로부터 (x0, y, m`) ~ q(x, y, m`)을 샘플링한다. 이러한 reverse process 학습 알고리즘은 아래와 같다. 

 

dynamic mask-aware reverse process

Shadow-free image x0에서 중간 텍스트 용어를 직접 샘플링하기 위해 가우스 확산 프로세스를 한계화할 수 있습니다. 여기서 βt는 노이즈 스케줄, αt = 1-βt, αθt = Qti = 1αi, ~ N(0, I)은 x0과 동일한 차원을 갖습니다. 노이즈 제거기 θ는 다음과 같이 섀도 이미지 y, 중간 변수 xt 및 시간 스텝 t를 입력으로 사용하여 노이즈 맵과 정제된 마스크 mt를 예측합니다:

Shadow mask 정보는 shadow-free image generation에 크게 의존하기 때문에 shadow-free image prediction과 mask refinement를 공동으로 수행하는 모델을 구축한다. Refined mask를 예측하기 위해 하나의 1 × 1 conv layer와 하나의 시그모이드 함수를 가진 mask prediction head를 €_θ의 마지막 layer 뒤에 추가한다. Diffusive objective function은 다음과 같다.

또한, 학습 단계의 shadow, shadow-free image pairs에 따라, refined mask의 합리성(rationality of the refined mask)을 제한하기 위해 ground truth shadow mask를 기준으로 사용할 수 있다.

여기서 gt shadow mask m_gt는 shadow, shadow-free image 사이의 residual map을 이진화하여 얻을 수 있다.

Ground Truth Shadow Mask

Hybrid objective function L_total은 위의 loss의 결합으로 얻어지며, 이는 DMDM 내의 denoiser €_θ의 학습을 진행시킨다.

L_mask 앞의 계수는 각 항의 균형을 위한 가중치 계수

# - 3  Unrolling-Inspired Diffusive Sampling 

Shadow degradation model(algorithm 2)을 기반으로, Shadow removal 작업을 degradation prior guided model로 공식화할 수 있다. Regularization terms(항)은 MAP(Maximum A Posteriori) 프레임워크에서 학습 가능한 conditional generative diffusion model에 의해 추론될 수 있다. 이때문에 사람이 사전 확률을 설정할 필요가 없다. 

 

Initial mask ~m이 부정확 할 수 있기 때문에 shadow mask m은 원하는 shadow-free image optimization과 함께 반복적으로 refine될 수 있으며, 이는 joint image-mask regularizer로 볼 수 있고 아래의 objective (8)을 최소화함으로써 얻을 수 있다.

 여기서 z와 v는 (8)을 제한된 문제로 변환하는 보조 변수이고 R(·)은 가정된 joint image and mask priors을 포착하는 regularizer이다. 여기서 사전 훈련된 degradation estimation network에 의해 degradation matrix h = D(y, ~m)로 추정할 수 있다고 가정한다. Equality constraints을 다루기 위해 두 개의 quadratic penalty 항을 사용하고 objective는 아래와 같이 다시 작성할 수 있다.

φ1과 φ2는 penalty parameters이다. ( φ1 = φ2 로 설정)

Half-quadratic splitting(HQS)과 같은 기존의 variable splitting algorithms을 사용하여 아래와 같이 세 가지 하위 문제를 반복적으로 해결함으로써 optimization 문제(equation. 9)를 해결할 수 있다:

여기서 (11)과 (12)는 closed-form의 솔루션을 갖는 quadratic penalty 항이 있는 최소 제곱 문제이다.

참고로, (13)과 (14)는 기존 연구의 data consistency(DC) steps를 참고하여 input variable과 reconstructed variable 사이의 정보를 공유한다. x와 m의 업데이트 과정은 DMDM의 샘플링 과정을 통해 해결할 수 있으며, 이는 Gθ(·)로 표시된다. (# - 2 참조)

아래의 Algorithm 2는 diffusion based unrolling의 전체 과정이며, 여기서 Gθ(·)는 2-3줄에 해당한다. Gθ의 샘플링은 추론 단계를 빠르게 하기 위해 DDIM의 diffusive sampling 전략을 따른다.

Single-stage model을 반복적으로 포워딩하는 것과 달리 diffusion model은 점진적 생성 프로세스를 통해 unrolling optimization problem을 해결하기 위한 자연스러운 아키텍처라는 점에 주목한다. 본 프레임워크는 inference time의 변화가 거의 없이 diffusion model에 degradation priors를 통합할 수 있다. Shadow-free image는 diffusion model을 기반으로 천천히 복원되는 반면, extra degradation prior는 shadow-free image 생성을 크게 가속화하고 반복을 실제 shadow-free data manifold에 가깝게 만들 수 있다. 

 

paper : https://arxiv.org/abs/2212.04711