NIPA 2021 인공지능 온라인 경진대회, 한국인 헤어스타일 세그멘테이션 2등

NIPA 2021 인공지능 온라인 경진대회, 한국인 헤어스타일 세그멘테이션 2등

이번 포스팅에서는 ARTuna이라는 팀으로 4명의 팀원과 함께 나간 인공지능 온라인 경진대회의 후기를 작성하겠습니다.

• 김현우 : https://github.com/choco9966
• 코드 : https://github.com/choco9966/Korean-Hair-Segmentation

저희 팀은 카이스트 산업 및 시스템공학과인 저와 세이지리서치에서 근무하는 하헌진님, 아트랩에서 근무하시는 3분의 팀원 훈재님, 재희님, 대선님 이렇게 5명이 팀을 이루어서 대회를 진행했습니다.

본선

저희 팀이 참여한 대회는 이미지 분야에서도 한국인 헤어스타일 세그멘테이션 대회였습니다.

• 과학기술정보통신부 - 한국인 헤어스타일 세그멘테이션 모델 (2위)

10개의 대회에서도 해당 대회를 선택한 이유는 세그멘테이션을 전공하기에 관련 능력을 잘 활용할 수 있어서입니다. 특히, 아트랩에서도 세그멘테이션을 이용한 작업들을 해왔고 팀원인 하헌진님 또한 MICCAI 대회에서 세그멘테이션으로 수상한 경험이 있으신 분입니다.

과학기술정보통신부 - 한국인 헤어스타일 세그멘테이션

한국인 헤어스타일 세그멘테이션대회는 512x512 크기의 여러 헤어스타일에 대한 이미지를 Segmentation 하는 대회였습니다. 연구분야이기도 하고 활용 서비스 분야가 재밌었던게, 헤어스타일을 탐지해서 이를 GAN으로 새로운 헤어스타일로 변경하든지 염색 등 다양한 식으로 활용할 수 있기에 유의미했던 경진대회였습니다. (참고 : 해당 자료에 보이는 모든 이미지와 마스크는 경진대회에서 제공된 이미지는 아니었고 저작권의 문제로 AI-Hub에 있는 자료들로 생성을 했습니다. )

대회의 평가 함수로는 IoU 였고 입력 이미지의 경우는 image 사진 자체를 제공했지만, mask에 대한 정보는 polygon의 형태로 json이 제공되었습니다. 다음 그림과 같이 json 파일의 polygon 키가 있어서 각 [x, y] 좌표를 기록하고 이를 모두 합쳐서 mask를 생성하는 구조였습니다.

그렇기에 제출도 Polygon으로 좌표를 만들어서 제출을 해야합니다. 단, 이후에 다시 언급하겠지만 문제 상황이 메모리의 문제 때문인지 학습과 제출의 polygon은 무조건 하나만 있어야했습니다. 그렇기에, 양갈래로 polygon이 2개가 생기는 경우에는 랜덤하게 하나가 제공되는 형태였습니다.

EDA

첫번째 문제점은 레이블이 아예 안되거나 적게 된 유형들입니다. 이미지를 보시면 알겠지만 왼쪽의 2개의 이미지는 마스크의 픽셀이 아예 안보이고 (실제로는 6~10개의 픽셀로 거의 없었습니다.) 혹은 오른쪽과 같이 굉장히 작은 영역에 대해서만 픽셀이 있는 것을 확인할 수 있었습니다.

두번째 문제점은 과하게 마스킹이 된 유형들입니다. 위의 예시처럼 외각 영역은 머리카락이 없어도 마스킹되는 현상을 확인할 수 있습니다. 실제 이를 해결해보려고 CRF를 적용해봤지만, 잘 되지 않았고 관련해서도 연구가 많이 있을정도로 일반적인 현상이었습니다. 하지만, 해결하는 것은 쉽지 않았습니다..

• 2020-ECCV - Learning with Noisy Class Labels for Instance Segmentation. [Paper] [Code]
• 2020-NIPS - Disentangling Human Error from the Ground Truth in Segmentation of Medical Images. [Paper] [Code]
• 2020-MICCAI - Characterizing Label Errors: Confident Learning for Noisy-labeled Image Segmentation. [Paper] [Code]

세번째 문제점은 마스킹이 아예 잘못된 경우였습니다. 보시면 알겠지만, 마스킹이 아예 이상한게 들어오거나 오른쪽처럼 마스킹이 회전되어서 나오는 경우가 있었습니다. 그리고 데이터에 같은 사람이 헤어스타일이나 각도만 다르게해서 나오는 이미지가 있는데, 한 사람의 마스킹 정보가 잘못되면 그 사람과 관련된 마싕 정보는 전부 잘못되었습니다. (이 점에서 헤어스타일이 다른데 같은 사진은 GAN 등의 모델을 통해서 만든게 아닐까 생각도 들었습니다.)

네번째 문제점은 양갈래의 머리입니다. 대회에서 가장 이해가 안되고, 어려운 점 중 하나였는데 Polygon을 제공도 하나만 해주고 제출도 하나만 할 수 있게 만든 점입니다. 그래서 양갈래의 머리경우 2개 이상의 Polygon으로 구성되어야하는데, 해당 부분이 없어서 여러개의 Polygon 중에서 하나가 제공되는 형태였습니다. (대부분 가장 큰 영역이 주어지기는 하지만, 몇몇 이미지의 경우는 빨갠색 박스처럼 작은 영역만 주기도 했습니다. ) sumission 시에도 Polygon을 하나만 제공해야하기에 초창기에는 모든 Polygon 영역을 계산해서 가장 큰 영역을 제출했는데, 이에 대해서 어떤 후처리를 했는지에 대해서는 뒤에서 설명드리도록 하겠습니다.

마지막 문제점은 유사한 이미지로 인한 학습 이미지의 증가입니다. 데이터의 수를 보시면 512x512로 큰 이미지가 19만장으로 굉장히 많았습니다. 실제 V100 2대로 학습셋 90%, 검증셋 10%으로 학습한 결과 24시간 이상이 걸렸습니다. 대회기간이 2주도 안되는 시간이어서, 모든 데이터셋을 가지고 학습을 하는 것에는 무리가 있었습니다.

Train

본격적으로 위의 문제점들을 어떻게 해결해나아갔는지 소개해드리도록 하겠습니다.

먼저, 레이블이 없거나 적게된 유형과 마스킹이 많이 된 유형은 마스킹된 영역의 크기를 계산해서 제거했습니다. 아래의 코드처럼 npixels라는 변수에 마스크 이미지에서 1을 가지는 개수를 계산하고, 이를 기반으로 zscore를 통해서 노이즈의 영역을 제거했습니다. 이를 기반으로 전처리 했을때 남성의 짧은 머리처럼 정상적인 이미지도 제거되는 문제가 발생하기는 했지만, 노이즈가 있는 이미지가 훨씬 많이 제거되어서 성능 향상이 있었습니다.

다음으로 마스킹이 아예 잘못된 유형은 모델을 통해서 제거를 진행했습니다. 기본적인 방법은 K-Fold 검증에서 사용하는 것처럼 학습셋으로 검증셋의 모든 이미지에 대해서 평가를 진행하고, 이를 기반으로 IoU가 낮은 이미지를 찾아서 제거하는 것입니다. 하지만, 학습셋에 노이즈가 많기에 학습이 제대로 안되는 문제가 있고 학습시간이 36시간으로 제한이 있어서 이 부분에 많은 시간투자를 할 수가 없다는 점입니다. 그래서 이를 해결하기 위해서 CO-Teaching 기반의 Segmentation인 CO-Segmentation을 이용해서 노이즈에 강건하게 모델을 만들고, 학습은 3에포크 이내로 학습시켜서 모든 이미지에 대해 IoU를 계산하였습니다.

그 결과 위의 히스토그램처럼 0.3 부분부터 1까지 IoU가 잘 분포하는 것을 확인할 수 있습니다. IoU가 낮은 이미지에 대해서는 Pseudo Label을 통해서 새로 레이블을 할수도 있었지만, 학습이미지가 많은 관계로 제거하고 본격적인 학습을 진행했습니다.

다음 문제로 1개의 PolyGon에 의해서 양갈래의 머리가 제대로 학습되지 않는 문제를 Self-Mix를 통해서 해결했습니다. 위의 그림처럼 Object의 마스크가 있는 영역에서 일부분을 잘라서 이미지의 다른 영역에 붙이는 것입니다. 제대로 잘려서 붙이게 될 경우 위의 그림처럼 양갈래의 효과를 기대할 수 있지만, 실제로는 랜덤하게 붙여지기 때문에 양갈래가 아니거나 이상한 모습으로 형성되기도 했습니다. 하지만, 애초에 양갈래 형식이 없다가 있는 것처럼 학습되기에 성능적인 측면에서 소폭 향상이 있었고 테스트 이미지의 양갈래에 대해서도 어느정도 결과가 나오는 것을 확인했습니다.

마지막 문제로 많은 학습 데이터에 의해서, 실험 속도가 오래 걸리는 문제를 해결하기 위해서 샘플링을 사용했습니다. 같은 사람의 이미지에서 각도만 다르게해서 여러개의 예시가 존재했기에 샘플링을 통해서 몇개만 추출하였고, 실제 20% 정도를 추출해서 실험할 경우 학습속도는 배로 빨라지지만 성능은 소폭 감소할 정도로 큰 차이가 없었습니다. 그리고 마지막 결과를 제출할때에는 전체 데이터를 모두 이용해서 학습해서 성능을 높이는 전략을 사용했습니다. 그 결과 제출 횟수가 106회로 남들보다 훨씬 많은 제출 횟수를 달성했습니다. (참고로 저는 샘플링을 이용했지만 다른 팀원은 256x256의 이미지로 줄여서 학습을 진행했고 이또한 속도는 1/4배 가까이 줄어드는 반면 성능 차이는 크게 없었습니다.)

그 외로 염색된 이미지에 대해서 IoU가 낮았는데, 이를 해결해주기 위해서 다음과 같이 랜덤하게 Coloring Augmentation 도 적용해가면서 못맞추는 이미지를 잘 맞추도록 노력해주었습니다.

Inference

Infernece 부분은 여러 모델을 학습할 시간이 없었기에 하나의 모델을 이용해서 최대한 앙상블을 하려고 노력했습니다. 그래서 Epoch10, 11, 12의 값을 모두 모아서 앙상블을 진행하면서 TTA로 Horizontal Flip을 적용하였습니다. 이때, 원본 이미지의 효과를 더 주기 위해서 원본 이미지에 가중치를 좀 더 주고 앙상블을 진행했습니다. 그 결과 되게 좋은 성능을 달성했지만, 아직 2등까지 가기에는 조금의 한계가 있었습니다.

이를 해결해주기 위한 필살기가 한개의 PolyGon만으로 두개의 PolyGon을 제출한 방법입니다. 실제 초중반의 경우 제출을 2개의 Polygon이 나올경우 크기를 계산해서 가장 큰 영역을 제출하였습니다. 하지만 후반부에는 2개~3개의 가장 큰 영역을 추출하고 둘을 잇는 가장 짧은 선을 만들었습니다. 이렇게 할 경우 빨간색 오차가 발생하지만, 파란색 부분에서 성능이 올라가서 마지막 성능 향상을 이루었던 전략입니다. (이때, 2번째 영역의 경우 잘못 마스킹된 영역도 많아서 1번째 다음 영역의 경우 크기를 최소한 몇으로 할지 등에 대해서 하이퍼파라미터 실험을 많이 진행했습니다. 남성과 여성 이미지에 대해서 다른 하이퍼파라미터를 사용하기도 했습니다. )

Result

최종적으로 0.9429의 IoU로 1등 팀과 소폭의 차이로 2등을 달성할 수 있었습니다.

7월 1일까지 오전까지만 해도 6등이었는데, 다들 포기하지 않아서 7월 2일에 대회 종료 9분남기고 2등을 달성할 수 있었습니다.

후기

바쁜일이 끝나고 대회끝난지 2달뒤에 후기를 작성했는데, 아직도 기억날 정도로 쉽지 않았던 대회였습니다. 많은 학습 데이터, 레이블의 노이즈, 그리고 2주도 안되는 짧은 기간 모두 쉽지 않은 상황이었습니다. 대회 거의 중후반까지는 3등안에도 못들어서 수상밖이었는데, 다른 분들이 좋은 아이디어 내주고 포기하지 않아서 2등으로 마무리할 수 있었던 것 같습니다.

이 대회하면서 스트레스 받느라 새치도 생기고 그랬는데, 좋은 결과 생기고 나중에 아트랩분들이 발표도 너무 잘해주셔서 전체 이미지 20 팀 중에서 3등까지 기록할 수 있었습니다. 이번 대회 통해서 연구 분야인 Sematic Segmentation 에 대해서도 자신감이 붙을 수 있었습니다. 수상을 할 수 있게 도와준 팀원들과 주최측에 감사인사드립니다. ARTuna팀 만세 ~

분석과 관련한 코드는 https://github.com/choco9966/Korean-Hair-Segmentation 깃허브 페이지를 참고하시기 바랍니다.