[StarGAN] pre-trained 모델로 커스텀 데이터셋 test하기

github.com/yunjey/stargan/issues/46 를 참고하였다.

 

 

구글 Colaboratory에서 진행한다. GPU는 필요 없다.

 

!apt-get install -y -qq software-properties-common python-software-properties module-init-tools
!add-apt-repository -y ppa:alessandro-strada/ppa 2>&1 > /dev/null
!apt-get update -qq 2>&1 > /dev/null
!apt-get -y install -qq google-drive-ocamlfuse fuse

from google.colab import auth
auth.authenticate_user()

from oauth2client.client import GoogleCredentials
creds = GoogleCredentials.get_application_default()

import getpass
!google-drive-ocamlfuse -headless -id={creds.client_id} -secret={creds.client_secret} < /dev/null 2>&1 | grep URL

vcode = getpass.getpass()
!echo {vcode} | google-drive-ocamlfuse -headless -id={creds.client_id} -secret={creds.client_secret}

우선 구글 드라이브와 연동하는 과정이다. 두번의 계정 인증이 필요하다.

 

 

cd /

!mkdir -p Gdrive
!google-drive-ocamlfuse Gdrive
!ls Gdrive

cd /Gdrive

'/' 경로 아래에 'Gdrive'라는 폴더를 생성하고 이동한다. 이 폴더가 구글 드라이브 폴더이다.

 

 

!git clone https://github.com/yunjey/StarGAN.git

cd StarGAN/

StarGAN repository를 clone하고 해당 폴더로 이동한다.

 

 

!bash download.sh celeba

CelebA 데이터셋을 다운로드 하는 과정인데, 갑자기 무슨 이유인지 Colab에서만 다운로드가 너무 느려서 20만장이 넘는 이미지들을 다운로드 받는게 며칠이 걸려도 불가했다.

그래서 이 과정은 패스했는데 제목과 같이 pretrain 모델을 커스텀 데이터셋으로 테스트하는 거라면 문제 없이 잘 진행된다.

 

 

!bash download.sh pretrained-celeba-128x128

StarGAN pretrained 모델을 다운받는다.

 

 

다음으로 test하는 명령을 수행하면 아마 'FileWriter'에 관련한 오류가 뜨게 된다.

이를 해결하기 위해 미리 'StarGAN/logger.py' 파일의 9번째 줄에서 'FileWriter' -> 'create_file_writer'로 수정한다. 

logger.py

 

 

test할 커스텀 데이터셋을 'StarGAN/data/custom/images' 경로에 저장한다.

원래 사용되는 이미지 파일의 형식을 살펴보니 .jpg 파일만 사용하고, 파일명은 '000001'과 같이 번호를 매겨 저장하기에 똑같은 형식으로 저장해주었다.

당연한 사실이지만, 측면이나 회전된 각도보다 정면으로 찍힌 사진이 성능이 훨씬 잘 나온다.

 

 

그리고 'StarGAN/data' 경로에 'list_attr_celeba_custom.txt' 파일을 만들어 커스텀 데이터셋의 각 이미지들을 레이블링한다.

list_attr_celeba_custom.txt

0.00MB

이 형식을 따라서 데이터를 추가해서 레이블링하면 된다.

Issue를 살펴보니까 레이블링은 수작업으로 하지 않고 레이블링 해주는 함수가 있다는 것 같기도 한데 나는 일단 간단하게 동작하는지 확인하기 위해서 수작업으로 했다.

그리고 test하는 명령에서 'Black_Hair, Blond_Hair, Brown_Hair, Male, Young' 이 5가지의 특성만 사용할 것이기 때문에 다른 특성들은 무조건 -1로 설정해두고 이 5가지 특성 값만 이미지에 맞게 수정해주었다.

이미지에 해당 특성이 들어있으면 1, 없으면 -1이다.

 

 

마지막으로 이미지들을 원본 크기로 그대로 저장하여 모델을 돌리면, 전처리 과정에서 얼굴이 아닌 다른 부분으로 자르는 경우가 발생해 성능이 낮게 나오는 경향이 있다.

따라서 커스텀 이미지들을 얼굴부분을 중심으로 crop하는 과정과 resize하는 과정이 필요하다.

crop은 수작업으로 하였고, resize는 StarGAN 모델이 128x128 크기의 이미지를 사용하기 때문에 128x128로 resize를 하

여 모델을 돌렸더니 결과물이 다음과 같이 나왔다.

실패한 결과물 (resize 오류)

이미 128x128인 이미지를 또 전처리하는 과정에서 더 작은 크기가 되서 이런 현상이 발생하는 것 같다. 

사진이 너무 작아 사진 밖에까지 합성결과가 튀어나오고, 합성이 제대로 되지 않은 것을 볼 수 있다.

따라서 200x200 정도로 resize하여 테스트하면 딱 적합한 결과물을 얻을 수 있었다.

 

test하는 명령은 다음과 같다.

# Custom 데이터셋으로 test
!python main.py --mode test --dataset CelebA --image_size 128 --c_dim 5 \
                --selected_attrs Black_Hair Blond_Hair Brown_Hair Male Young \
                --celeba_image_dir data/custom/images \
                --attr_path data/list_attr_celeba_custom.txt \
                --model_save_dir='stargan_celeba_128/models' \
                --result_dir='stargan_celeba_128_custom/results'

 

 

성공한 결과물

직접 해보면 성능이 굉장히 안좋을 줄 알았는데, 합성이 상당히 잘된 것을 확인할 수 있었다.