GAN을 이용한 염색 및 헤어 스타일 합성, <꽤 GAN찮은 헤어살롱>

제12회 빅데이터 연합동아리 BOAZ 컨퍼런스를 위해 진행한 프로젝트 소개글입니다.
발표 자료는 BOAZ slide share에서, 발표 영상은 BOAZ Youtube 계정에서, 관련 코드는 Github에서 확인하실 수 있습니다. (글씨를 클릭하면 해당 사이트로 이동합니다.)

 

어떤 머리가 어울릴지 미리 알아보고싶어!

'앞머리를 기르면 어떨까?', '염색은 처음인데.. 어울릴까?' 헤어 스타일은 많은 사람들의 평생 고민입니다. 미용실에 가기 전에 헤어 스타일을 미리 체험해볼 수 있다면 얼마나 좋을까요? 실제로 시중에는 소비자들의 이런 니즈를 반영한 어플이 많이 출시되어있습니다. 

 

어플 사용 결과

 

위가 바로 배우 손예진님의 사진을 어플에 사용한 결과입니다. 결과가 나쁘지 않죠? 하지만 헤어 스타일 별로 고정된 틀이 존재하기 때문에 다양한 분위기를 연출하기가 어렵습니다. 위의 사진을 예로 들자면, 빽빽한 앞머리, 시스루뱅과 같은 다양한 스타일의 앞머리를 만들수가 없다는 것이죠. 또한 어플이 제공하지 않는 헤어 스타일은 체험이 불가능한 점 역시 단점입니다. 

그리고 '이런 문제점을 딥러닝을 통해 해결할 수 있지 않을까?'라는 생각에서 저희 프로젝트가 시작되었습니다. 자세한 내용을 알아보기 전에, 먼저 저희 프로젝트의 결과를 공개하겠습니다.😊 중간 결과를 포함한 더 많은 이미지는 이 곳에서 확인하실 수 있습니다.

 

최종 결과 (1) - 염색

최종 결과 (2) - 헤어 스타일

 

 

프로젝트 목표

저희가 후보로 두고 실험해보았던, 또 최종적으로 선택한 모델은 모두 생성 모델, GAN(Generative Adversarial Networks) 입니다. 여러 GAN 모델 중에서도 '헤어 스타일 변형'이라는 태스크에 맞는 모델을 선택해야 했는데요, 여러 논문을 읽고 주변에 조언을 구한 결과 Clova AI에서 발표한 StarGAN-v2를 최종 모델로 선정하게 되었습니다. 

 

StarGAN-v2 Figure 4

 

위의 사진이 바로 StarGAN-v2 논문에 수록된 이미지 생성 결과입니다. 성별(Female, Male) 도메인으로 나뉘어 학습된 모델을 사용, 가장 윗줄의 이미지를 Source로, 가장 왼쪽의 이미지의 스타일을 입혀 생성한 이미지입니다. 헤어 스타일이 꽤나 자연스럽게 잘 입혀지는 것을 확인할 수 있습니다. 하지만 머리 스타일 외에도 머리 색, 피부 톤, 화장 등의 다른 특징까지 가져온다는 점이 저희 태스크에는 치명적인 단점이었습니다. 

v2에 앞서 StarGAN-v1의 pretrained 모델을 사용해보며 알게 된 단점이 한 가지 더 있었는데요, 바로 전반적인 사진의 특징을 모두 가져온다는 점이었습니다. 

 

StarGAN-v1 inference 결과 (왼쪽부터 차례로 원본/앞머리 없음/흑발/갈색/금발/백발/직모/파마 도메인)

 

위의 사진은 아이유님의 사진을 StarGAN-v1의 pretrained 모델의 인풋으로 하여, 각각 앞머리 없음/흑발/갈색/금발/백발/직모/파마 도메인으로 translate 한 결과입니다. pretrain에 이용된 사진(CelebA) 대부분이 서양인이었기 때문에 피부가 하얘지거나 코가 오똑해지는 등의 헤어 스타일 외 변화가 있었습니다. 특히 백발 도메인의 경우 대부분의 사진이 노인이었기 때문에 주름과 같은 다른 특징도 함께 해당 도메인의 특징으로 학습한 것을 확인할 수 있었습니다.

본격적인 프로젝트 시작 전 파악한 위와 같은 문제점을 바탕으로 저희는 다음과 같은 프로젝트 목표를 설정하였습니다.

◾ 동양인에 최적화된 모델을 구축할 것.
◾ 다양한 헤어 스타일의 이미지 생성이 가능할 것. 
◾ 헤어 스타일 외의 부분 (ex. 머리색, 피부색, 화장 등)의 변형은 최소화 할 것.

위와 같은 목표를 달성하기 위해 어떤 모델을, 어떻게 사용했는지 확인해볼까요?

 

전체 프로세스

 

<꽤 GAN찮은 헤어살롱> 프로세스

 

<꽤 GAN찮은 헤어살롱> 모델의 프로세스는 위와 같습니다. 1) 염색과 2) 헤어 스타일 변형, 두 가지의 헤어 스타일 체험이 가능하고 각각은 서로 다른 프로세스를 거칩니다.

StarGAN-v2와 SEAN(Image synthesis with Semantic region-adaptive normalization), 크게 두 개의 모델을 사용하였고, StarGAN-v2는 새로운 헤어 스타일의 이미지를 생성할 때, SEAN은 헤어 스타일 외 변화를 복구하거나 염색을 하기 위해 사용하였습니다. 지금부터는 각각의 모델에 대해 간단히 소개해드리도록 하겠습니다.

 

(1) StarGAN-v2

StarGAN-v2는 지난 19년 12월에 Clova AI팀에서 공개한 모델로, 하나의 모델로 여러 도메인,  여러 스타일의 이미지를 생성할 수 있는 Multi-modal, domain 모델입니다. 여러 헤어 도메인의 다양한 스타일 사진을 생성해야 하는 저희의 태스크에 딱 맞는 모델이었죠.🥰

 

StarGAN-v2 Figure 2

 

이미지를 생성할 수 있는 방법이 여러가지라는 점도 매력적이었습니다. StarGAN-v2는 이미지와 스타일 벡터를 인풋으로, 새로운 이미지를 생성합니다. 그리고 스타일 벡터를 구하는 방법에는 1) Reference 이미지를 Style encoder에 넣어 도메인에 맞는 스타일 벡터를 구하는 방법과 2) 랜덤 벡터를 Mapping network에 넣어 도메인에 맞는 스타일 벡터를 구하는 방법, 2가지가 있습니다. StarGAN-v2에 대한 보다 자세한 내용은 2020/03/23 - [ML & DL/GAN] - 도메인과 스타일, 모두 잡았다! StarGAN v2을 참고해주세요.

 

(2) SEAN (Image synthesis with Semantic region-adaptive normalization)

 

SEAN Figure 4 - Generator

 

SEAN은 Segmentation map을 바탕으로 스타일 벡터를 추출하고, 이미지를 reconstruct 하는 모델입니다. 역시나 생성 모델이고요!

지금까지 프로젝트의 전반적인 개요를 알아보았으니  이제 어떤 데이터셋을 이용해 모델을 학습했고, 어떤 아이디어로 프로젝트 목표를 달성하려 했는지 알아보겠습니다!

 

데이터셋과 라벨링

도메인 별로 학습되는 StarGAN-v2의 특성상, 데이터 수집에 앞서 헤어 스타일 도메인을 정해주어야 했습니다. 이에 저희 팀은 색 - 검은색&갈색, 금발 / 앞머리 - 유, 무 / 스타일 - 긴 생머리, 긴 파마머리(롱펌), 단발머리, 쇼트커트, 총 8개의 도메인으로 헤어 스타일을 분류하였습니다. 

이제 필요한 건 각 도메인에 해당되는 충분한 양의 이미지! 특히나 동양인에 최적화된 모델을 만들기 위해 가장 필요했던 것은 많은 양의 동양인 사진이었습니다. 다행히 seeprettyface.com이라는 사이트에서 StyleGAN을 통해 생성된 8만 장의 동양인 사진을 확보할 수 있었습니다. (의도한 것은 아니었지만 데이터마저도 GAN으로 생성된...😅) 여기에 금발, 긴 파마머리와 같이 개수가 부족한 도메인에 대해 포털 사이트, SNS 크롤링을 통해 사진을 보충하였습니다. (약 2000장)

 

데이터셋 예시

 

이렇게 모은 8만 2천여 장의 사진은 학습을 위해 머리 색, 앞머리 유무, 헤어 스타일에 따라 라벨링 되어야 했습니다. 사람이 직접 하기에는 방대한 양이었기 때문에 모델을 사용하였습니다. 처음에는 Resnet34를 이용하고자 했으나, 헤어 스타일 분류에 어려움을 보여 어텐션이 도입된 Resnet34_CBAM 모델을 이용하였습니다.

인간지능👩🏻‍💻을 이용해 1500장가량을 라벨링 하였고, 좌우반전 augmentation을 통해 3000장 정도를 확보한 뒤 학습을 진행했습니다. 그 결과 validation 기준 색 95%, 앞머리 99%, 스타일 93%의 accuracy를 보이는 모델을 얻을 수 있었고, 성공적으로 라벨링을 할 수 있었습니다. 라벨링 결과는 다음과 같습니다.

 

라벨링 결과 (왼쪽부터) 색 - 검은색&갈색, 금발 / 앞머리 - 유, 무 / 스타일 - 긴 생머리, 긴 파마머리, 단발, 숏컷

 

◾ 색         | 검은색 & 갈색 79109장    금발 2850장    기타
◾ 앞머리 | 있음 59986장    없음 22173장
◾ 스타일 | 긴 생머리 19286장    긴 파마머리 9673장    단발 7084장    숏컷 28426장    기타

 

모델 학습과 Inference

(1) StarGAN-v2

저희 팀은 따로 GPU를 가지고 있지 않아서 Google Cloud에서 제공하는 Tesla V100을 사용했는데요, 돈이 많이 들고 한 번 학습에 2~3일가량의 시간이 걸리다 보니 많은 조건 하에서 실험을 해보지는 못했습니다.😥 최종적으로 채택된 모델은 256 x 256 이미지를 배치 사이즈 7로, 총 100,000 iterations 동안 학습시킨 모델이었습니다. 그럼 이제 학습 결과를 한 번 확인해볼까요?

가. Reference 이미지를 이용해 이미지를 생성한 경우

 

(왼쪽부터 차례로) 앞머리 있음 / 앞머리 없음 / 긴 생머리 도메인으로 변환한 결과

(왼쪽부터 차례로) 롱 펌 / 단발 / 숏컷 도메인으로 변환한 결과

 

 

나. Random vector를 이용한 경우

 

(왼쪽 위부터 차례로) 앞머리 있음 / 앞머리 없음 / 긴 생머리 (오른쪽 위부터 차례로) 롱 펌 / 단발/ 숏 컷 도메인으로의 변환

 

위에서 언급했던 StarGAN-v2의 두 가지 이미지 생성 방법으로 이미지를 생성한 결과인데요, 결과가 꽤 괜찮지 않나요? 같은 도메인이더라도 다양한 스타일의 데이터가 존재했던 앞머리 有, 단발 도메인은 Source 이미지가 비교적 다양한 헤어 스타일을 가진 이미지로 바뀐 것을 보실 수 있습니다. 하지만 이미지 전체의 스타일을 가져오고 입히는 모델의 특성상 Source 이미지의 머리 색, 피부톤, 화장 등까지 바꿔놓기 때문에 조금은 다른 사람으로 느껴지기도 합니다.

 

금발 도메인으로 변형한 결과

 

StarGAN-v2의 문제점이 또 한 가지 있었는데요, 바로 데이터의 개수가 생성되는 이미지의 퀄리티를 좌우한다는 점이었습니다. 저희가 가지고 있던 금발 이미지의 경우 2850장으로, 제대로 된 학습을 보장할 수 있는 최소한의 데이터 수 10,000장에 비해 턱없이 부족한 숫자였습니다. 검은색 외의 머리색을 가진 동양인 데이터를 10,000장 이상 확보하기가 어렵다는 점, 확보해야 할 색상 도메인이 아주 많다는 점에서 StarGAN-v2는 염색에 적합한 모델이 아니었죠.

따라서 헤어 스타일 외에도 머리 색, 피부톤, 화장과 같은 다른 스타일까지 가져온다는 점, StarGAN-v2가 염색에 적합한 모델이 아니라는 점을 보완하기 위해 저희는 새로운 모델 SEAN(Image synthesis with Semantic region-adaptive normalization)을 도입하게 되었습니다.

 

(2) SEAN(Image synthesis with Semantic region-adaptive normalization)

SEAN을 사용하기에 앞서, SEAN은 Segmentation map을 이용하는 모델이기 때문에, segmentation map을 구하는 과정이 선행되어야 했습니다. 이 과정의 경우, https://github.com/zllrunning/face-parsing.PyTorch에서 제공하는 pretrained 모델을 활용하였습니다. 또한 SEAN의 Generator 역시 Official Github에서 제공하는 pretrained 모델을 사용하였습니다.

위에서 언급하였듯이, SEAN 모델을 통해 얻고자 했던 것은 헤어 스타일 외 변화를 보정하는 것과 데이터 확보 없이 다양한 색에 대한 염색을 성공시키는 것이었는데요, 저희는 두 목표를 각기 다른 아이디어를 통해 해결할 수 있었습니다.

 

가. 헤어 스타일 외 변화 보정

 

Architecture

헤어 스타일 외의 변화를 보정하기 위해 저희가 선택한 프로세스는 다음과 같습니다.

① StarGAN v2를 이용해 헤어스타일을 바꾸어준다.
② 원본 사진과 1번에서 새롭게 생성한 사진의 Segmentation map을 구한다.
③ 원본 사진의 스타일 벡터를 2번에서 구한 원본 사진의 Segmentation map을 기반으로 추출한다.
④ 생성한 사진의 Segmentation map에 3번에서 구한 스타일 벡터를 입혀 이미지를 reconstruct 한다.

 

나. 염색

 

Architecture

다음은 염색 프로세스입니다. 염색의 경우 헤어 스타일 변환과 달리 이미지의 geometric 한 변화가 필요 없기 때문에 SEAN 모델만을 단독으로 사용하였습니다.

① Src 이미지와 Ref 이미지의 Segmentation map을 구한다.
② 1번에서 구한 Segmentation map을 기반으로 각각 사진의 스타일 벡터를 추출한다.
③ Ref 이미지의 스타일 벡터 중 헤어 영역의 스타일 벡터를 Src 이미지의 스타일 벡터의 헤어 영역 스타일 벡터와 바꿔치기한다.
④ 바꿔치기한 스타일 벡터를 Src 이미지의 Segmentation map에 reconsturct 한다.

 

최종 결과와 아쉬운 점

위 프로세스를 따라 생성된 최종 결과물을 공개합니다..! 😄

 

최종 결과 (1) - 염색

최종 결과 (2) - 헤어 스타일

 

어떤가요? 염색은 정말 흠잡을 곳 없이 잘 되는 것 같습니다. 헤어 스타일은 기대한 만큼의 드라마틱한 변화는 아니지만 도메인에 해당하는 이미지를 잘 생성한 것을 볼 수 있습니다. 좀 더 드라마틱한 헤어 스타일 변화, 헤어 스타일 외 부분 변화 감소, 깨끗한 이미지 생성 등이 남은 과제가 될 것 같습니다.

 

프로젝트를 마무리하며

이렇게 딥러닝 입문 1년 차, 딥린이의 첫 딥러닝 프로젝트가 마무리되었는데요👏🏻, 사실 이 프로젝트는 8월 1일 자로 끝이 났었고, 이 글은 8월 말에 80% 이상 완성이 되어 있었던 글입니다. 쓰고 나니 프로젝트가 너무 빈약해 보이기도 하고 글도 재미가 없는 것 같아서 임시저장만 해두다가... 깃헙을 보시고 연락 주시는 분들이 계시기도 했고,  작고 부족한 것이 많은 프로젝트였지만 기록으로 남기고 사람들과 공유해보고자 글을 올리게 되었습니다.

저는 프로젝트 마무리 이후 좋은 기회가 생겨서 훌륭하신 분들과 같은 주제로 연구를 이어나가고 있습니다. 그러면서 프로젝트 전반과 코드에 대해 다시 살펴보고 피드백을 받을 시간이 있었는데, 결국 StarGAN을 Segmentation 마스크만을 생성하기 위해 사용했다는 점, 코드가 깔끔하지 않다는 점 등 불과 두세 달 밖에 지나지 않은 프로젝트임에도 불구하고 많은 부족한 점을 발견할 수 있었습니다. 사실 그래서 글 올리는 걸 더 망설이게 된 것도 있었네요 😅

부족한 프로젝트지만 누군가에게 조금이라도 흥미 있고, 도움이 되는 내용이었으면 좋겠습니다. 궁금하신 점이 있으시다면 댓글 남겨주시면 답변드리도록 하겠습니다. 읽어주셔서 감사합니다.