[Project] Personalized Image Enhancement 연구

게시 2025/01/08 업데이트 2025/01/13

By Jaemin Park 17 분읽는 시간

주제

두번째로 시작한 연구 주제로, PIE(Personalized Image Enhancement)는 이미지 개선(Image Enhancement) 과정에서 사용자 선호도(User Preference)를 반영하는 것이다.
품질이 좋지 않은(Blur, Low Light 등) 이미지 A가 있을 때, 이미지 개선 의 경우 하나의 GT image 를 설정하고 학습한다.
하지만 현실에서는 사람마다 이미지를 개선할 때 자신의 취향을 반영하고 싶어한다. 이러한 개인적인 이미지 취향을 사용자 선호도라고 정의한다.

기존의 PieNet, StarEnhancer 의 경우 MIT Adobe 5K Dataset에서 각 Expert의 이미지 개선 결과들을 개인 선호도 이미지 셋으로 설정한다.
그리고 각 Expert Image set 에 대해 사용자 선호도를 Preference Vector로 추출하는 것을 목표로 한다. PieNet의 경우 Triplet Loss를, StarEnhancer의 경우 Classfier 구조를 이용해 학습을 진행한다.
위와 같이 Preference Vector를 추출하는 Embedding Network를 Pre-training 한 후, Enhancer를 학습시키는 2-step Learning 방식을 취한다.

하지만 이러한 방식은 실제로 Preference vVctor가 직접적으로 사용자의 선호도를 내포하고 있고, 이것이 직접적으로 Enhancer에 우리가 생각하는 것처럼 영향을 준다고 확인하기 어렵다.
즉, Enhancer는 개선만 진행하고 preference vector가 중간에 style을 바꿔주는지, Enhancer 자체가 사용자 선호도까지 고려해서 학습되는 것인지 확인할 수 없다. 이는 두 모델의 구조가 모두 Black box 구조이기 때문이다.

따라서 본 연구 과정은 End to End 방식을 최대한 유지하면서 두 모델과 비슷하거나 좋은 성능을 내는 것은 목표로 시작하였다.

Log

2024.5-6.

목표
- 연구 주제 선정 및 논문 분석
- Base 모델 기초 설계
연구 주제 선정을 위한 논문 탐색
- Low-Light Image Enhancement
- Underwater Image Restoration
- Deraining
- Personalized Image Enhancemnet (PIE, 선정)
  Deraining, Underwater는 연구 논문의 수가 거의 없고 Restoration 분야 모델이 해당 task에서 좋은 성능을 발휘하는 것으로 확인됨.
  Low-Light의 경우, 당시 최고 성능이었던 Retinexformer가 오랫동안 최고 성능을 유지했으며, 24년 이후로 관련 논문이 많이 줄어듦.
  결론적으로 Deraining, Underwater는 인용할만한 자료가 부족하고, Low-Light는 더 좋은 성능의 모델을 내기 어려워 연구 난이도가 높아 PIE를 선정
관련 논문 분석
- PieNet, StarEnhancer, Masked Style Modeling
Base 모델 기초 잡기
- MIT Adobe 5K Dataset 및 Adobe Lightroom 활용
- Pienet, StarEnhancer 모델을 직접 학습시켜 성능 측정
역할 분배
- 나 : Encoder 성능 개선 연구, DataLoader, Dataset 관련 구현
- 팀원 : Enhancer 성능 개선 연구, Adobe Lightroom을 이용한 추가 Dataset 제작
문제점
- 학습 시간 : Pienet의 경우 Encoder, Enhancer 를 합쳐 약 2일, StarEnhancer는 3일 넘게 소요. 간단한 test를 위한 단축 방법이 필요함
- 학습이 도중에 중단되는 현상 : Python 3.7 버전에서 가끔 발생한다고 한다.
- 블루 스크린 : 알 수 없는 블루 스크린 발생 (DataLoader 중 메모리 문제로 추측)
- DataLoader에서 중단 현상 : 사진을 각 훈련 마다 약 20장씩 CPU에서 가져오는데, 이 과정에서 아예 가져오지 못하고 훈련이 멈춰버리는 현상
해결
- 배치 수와 가져오는 이미지 수를 줄여서 DataLoad, 블루 스크린 문제는 대부분 해결했으나 이로 인해 학습 시간이 증가해버리는 문제 발생
- 결국 GPU 성능의 한계라고 판단하고 학습 시간이 증가하더라도 수를 줄이는 방법을 선택
- 파이썬 버전 문제는 해결하지 못함. 파이썬 버전 변경 시 github 코드 자체가 안 돌아가서 코드 전체를 수정해야 하는데 비효율적임

2024.7-8.

목표
- 추가 논문 탐색 및 참고 자료 탐색
- 구현 및 실험 진행
- 결과 분석 및 성능 개선 방안 탐색
추가 논문 탐색
Encoder 부분에 중점을 두고 Deep Metric Learning 분야를 집중적으로 탐색
- Deep Metric Learning : 서로 비슷한 데이터는 가깝게, 다른 데이터는 멀게 배치되도록 거리 공간을 학습하는 기법
- 같은 스타일의 이미지는 거리 공간에서 더 가까워지도록, 다른 스타일의 이미지는 더 멀어지도록 학습
- Circle Loss, Center Loss, Triplet Loss, Angular Loss 등 활용
구현 및 실험
구현 부분은 기존 코드에서 필요한 부분을 수정, 추가하는 방식으로 진행하였음
- DataLoader, Dataset 구현 : 학습 이미지 수를 조정할 수 있도록 수정
- Triplet Loss 대신 Angular Loss, Center Loss 로 대체하여 후 실험 진행
- T-SNE 시각화 : Preference Vector 가 잘 분리되는지 확인
- T-SNE 구현 코드를 Classifier 분야의 Dataset으로 test를 수행, 오류가 있는지 확인
결과
- 사용자별 Preference Vector가 명확히 분리되지 않음
- 동일한 T-SNE 코드로 다른 데이터셋을 테스트한 결과, 시각화 코드 자체에는 문제가 없었음
- 기존 논문 대비 PSNR이 약 1.0 낮게 측정
- Loss를 변경하면서 여러가지 방면으로 진행한 실험 결과 +- 0.1 정도의 PSNR 변화로 큰 차이점을 발견할 수 없음.
개선 방안 탐색
- Metric Learning 기법의 최신 동향 및 논문 검토
- Metric Learning 외 다른 방식으로 Embedding Network를 구현할 수 있는 방법 모색
- 다수의 이미지를 더 적은 시간 내에 가져올 수 있는 최적화 방법 탐색

2024.9-11.

목표
- Deep Metric Learning 추가 실험
- Loss Function 확장
- 새로운 구조의 모델 설계 및 구현
- Encoder-Enhancer의 학습률을 차별화 실험
- Style 이미지 수의 변화를 통한 성능 실험
Deep Metric Learning 추가 실험
- 7-8월에 진행한 Loss 함수 외에 다양한 Loss 함수로 실험 진행
- 그러나 전반적으로 기대했던 성능 향상을 이루지 못함
Loss Function 추가
- Gram Loss, TV Loss를 추가한 뒤 가중치를 조정하며 여러 차례 실험 진행
- Gram Loss : PSNR 약 0.15 향상 (\lambda : 0.01)
- TV Loss : PSNR 약 0.05 향상 (\lambda : 0.01)
새로운 구조의 모델 설계 및 구현
- 모델 설계
  - U-net, Attention, Diffusion 등의 모델 구조 활용
  - AdaIN (Adaptive Instance Normalization) 기법 활용
  - StyleGAN, pixel2style2pixel 등 StyleTransfer 모델을 참고, 해당 모델들의 Style Feature 추출 방법을 변형해서 도입
- 결과
  - U-net 과 AdaIN 병합 : 성능 저조
  - Attention 과 AdaIN 병합 : 준수한 성능을 보였으나 기존 성능을 초과하지 못함
  - pixel2style2pixel 의 pSpEncoder 변형 : PSNR 약 0.13 향상
  - Diffusion 구조 : 컴퓨터 성능 한계로 적용하지 못함
Encoder-Enhancer의 학습률 차별화 실험
- Encoder가 Enhancer보다 더 과적합되는 현상을 보여, Encoder 학습률을 낮추어 최적의 학습률 탐색
- Encoder LR = 1e-5 , Enhancer LR = 1e-4 에서 최고 성능, PSNR 약 0.12 향상
Style 이미지 수의 변화를 주며 실험
- 기존 style feature 추출 과정 : Style 이미지 1장 → Encoder → Style Feature
- 개선된 과정 : Style 이미지 3~5장 → Encoder → Style Feature
- Batch와 이미지 수를 결합하여 처리 시간 최적화
- PSNR이 약 0.1 향상
- Style 이미지 6장 이상에서는 Dataloader 멈춤 현상 발생 또는 처리 시간이 과도하게 증가

구현

Gram Loss

  
  def gram_matrix(img):
      (b, c, h, w) = img.size()
      img = img.view(b, c, h * w)
      gram = torch.bmm(img, img.transpose(1, 2))
        
      return gram / (c * h * w)

  def gram_loss(y_hat, style):
        
      # style img : batch, img-count, c, h, w
      if style.dim() > 4 : 
          B, N, C, H, W = style.size()
          y_hat = y_hat.unsqueeze(1).expand(-1, N, -1, -1, -1)

          y_hat = y_hat.reshape(B * N, C, H, W)
          style = style.reshape(B * N, C, H, W)

      y_hat = gram_matrix(y_hat)
      style = gram_matrix(style)
        
      loss = F.mse_loss(y_hat, style)
        
      return loss

Style Attention

  
  class SA_Block(nn.Module):
      def __init__(self, dim=256, dim_head=64, num_fmaps=4): 
          # x : dim, H, W
          # fmaps : num_fmaps, dim, h, w 
          # out : dim, H, W
          super(SA_Block, self).__init__()

          self.num_fmaps = num_fmaps
          self.dim_head = dim_head

          self.to_q = nn.Linear(dim, dim_head * self.num_fmaps, bias=False)
          self.to_k = nn.Conv2d(in_channels=dim * self.num_fmaps, out_channels=dim_head * self.num_fmaps, kernel_size=1, groups=self.num_fmaps, bias=False)
          self.to_v = nn.Conv2d(in_channels=dim * self.num_fmaps, out_channels=dim_head * self.num_fmaps, kernel_size=1, groups=self.num_fmaps, bias=False)
          self.rescale = nn.Parameter(torch.ones(self.num_fmaps, 1, 1))

          self.ffn = nn.Linear(dim_head * self.num_fmaps, dim, bias=True)
      def forward(self, x, fmaps):
          """
          x : b c h w (4, 256, 32, 32)
          fmaps : b n c h w (4, 4, 256, 32, 32)
          # N : Style image 개수
          """
          B, N, C, H, W = fmaps.size() # [4, 4, 256, 32, 32]
          fmaps = fmaps.view(B, N*C, H, W) # (4, 4*256, 32, 32) = 4, 1024, 32, 32
          x_in = x.permute(0,2,3,1).reshape(B, H*W, C) # (4, 32*32, 256) = 4, 1024, 256

          q_in = self.to_q(x_in) # (4, 32*32, 256) = [4, 1024, 256]
          k_in = self.to_k(fmaps).permute(0,2,3,1).reshape(B, H*W, -1) # (4, 64*4, 32, 32) -> (4, 32, 32, 64*4) -> (4, 32*32, 64*4) = [4, 1024, 256]
          v_in = self.to_v(fmaps).permute(0,2,3,1).reshape(B, H*W, -1) 

          d = C // self.num_fmaps
          q = q_in.view(B, H*W, self.num_fmaps, d).permute(0, 2, 1, 3)  # (4, 4, 32*32, 64) = [4, 4, 1024, 64]
          k = k_in.view(B, H*W, self.num_fmaps, d).permute(0, 2, 1, 3)  
          v = v_in.view(B, H*W, self.num_fmaps, d).permute(0, 2, 1, 3)  

          q = q.transpose(-2, -1) # [4, 4, 64, 1024]
          k = k.transpose(-2, -1)
          v = v.transpose(-2, -1)
          q = F.normalize(q, dim=-1, p=2)
          k = F.normalize(k, dim=-1, p=2)

          attn = (k @ q.transpose(-2, -1)) # [4, 4, 64, 64]
          attn = attn * self.rescale
          attn = attn.softmax(dim=-1) 

          out = attn @ v # (4, 4, 32*2, 32*32) = [4, 4, 64, 1024] = (B, HEAD, C, H*W)
          out = out.permute(0, 3, 1, 2) # B, H*W, HEAD, C = 4, 1024, 4, 64
          out = out.reshape(B, H*W, self.num_fmaps*self.dim_head) # 4, 1024, 4*64
          out = out + q_in # 4, 1024, 256
          out = self.ffn(out) # 4, 1024, 256
          out = out + q_in # 4, 1024, 256
          out = out.view(B, H, W, C).permute(0, 3, 1, 2)

          return out

전체 일정

5월~6월 : 주제 선정, 논문 탐색, 기존 논문 실험 결과 확인, 역할 분배
7월~8월 : Base 모델 구현 및 실험, Metric Learning 연구 및 적용, 시각화를 통한 문제 분석, 최신 논문 탐색,
9월 : Metric Learning 추가 실험, U-net 기반 새로운 구조 모델 설계
10월 : Loss 함수 확장(Gram, TV), 학습룔 차별화 실험, Style Transfer 분야의 Feature Extractor, AdaIN 도입
11월 : Attention 기반 모델 설계, Style 이미지 수 변화 실험

총평

이번 연구는 성과를 얻지 못한 실패한 연구였다. 성능 향상에 크게 기여하지 못했고, 모델 설계와 연구 방향성에도 아쉬움이 남았다.
그럼에도 불구하고 많은 것을 배운 과정이었다.
논문을 읽고 분석하는 능력이 향상되었으며, Deep Metric Learning, Style Transfer 등 다양한 분야의 이론적 지식과 연구 방법을 습득했다.
수식을 이해하고 이를 코드로 구현하는 능력도 많이 향상되었고, Git clone 후 부분 수정하는 방식에서 벗어나 처음부터 끝까지 구현하는 경험을 통해 Pytorch에 대한 전체적인 이해도를 높일 수 있었다.
교수님과의 지속적인 상담을 통해 연구의 방향을 설정하고 새로운 방법을 고안하는 과정도 중요한 배움이 되었다.

이번 연구 실패의 주요 원인은 구현 능력 부족, 설계 검토 미흡, 방향 설정이었다.
특히 설계 검토 단계에서 설계를 완료한 후 구현 중 실수를 발견하고 다시 설계하는 일이 반복되었다. 이러한 점을 개선하면 다음 연구에서 시간을 크게 단축할 수 있을 것이다.

Reference

핵심 참고

관련 조사

DeepLPF
DPE
Harmonizer
MAXIM
LPTN
N-pair Loss
Retinexformer
RSFNet
CSRNet
DreamStyler
HDRNet
Puff-Net
Multi-Similarity Loss
Style Transformer for Image Inversion and Editing
Local Color Distributions Prior for Image Enhancement
Inversion-Based Style Transfer with Diffusion Models
Enhancement by Your Aesthetic - An Intelligible Unsupervised Personalized Enhancer for Low-Light Images
Close Imitation of Expert Retouching for Black-and-White Photography
Image Deomoireing with Learnable Banpass Filters
Exposure: A White-Box Photo Post-Processing Framework

Develop, Project