Virtual Try on - 4

11 Apr 2021 in Virtual Try On

모델 학습

학습 데이터를 준비하였기 때문에 직접 학습을 하였다. Image Segmentation Model은 FCN, DeepLabv3, PSPNet 등이 있지만 Unet이 적당한 성능과 적당한 크기를 갖고 있다고 생각하여 Unet 모델을 사용하였다.
기본적인 Unet을 학습시킨 결과 기본적인 옷에 대해서는 결과가 양호했지만 흰 옷인 경우 결과가 이상하게 나왔다.

이러한 문제를 해결하기 위해 다른 방법을 찾아야 했고, Unet이 GAN으로도 많이 사용된다고하여 Discriminator 모델을 추가하여 GAN 방식으로 학습을 시켰다. 추가한 Discriminator 모델은 Patch GAN을 사용하였다.
Unet을 GAN 방식으로 학습한 결과 흰 옷인 경우 기존 방식보다 더 좋은 결과를 확인할 수 있었다.

하지만 흰 옷이 아닌 경우 기존 방식에서는 없던 결과물에 노이지가 발생하였다.

이 문제는 모폴리지 연산(열림 연산)을 사용하여 해결할 수 있었다.

또한 배경색이 흰색이 아닌 경우 다음과 같이 배경을 포함하여 모두 옷이라고 판단는 문제가 있었다.

이 문제가 발생한 이유가 Train data로 사용한 모든 이미지의 배경색이 흰색이기 때문에 발생했다고 생각하여 배경색만을 임의로 추가한 이미지를 Train dataset에 추가하였다. 임의로 만든 배경색은 쇼핑몰에서 볼 수 있는 배경색으로 지정하였다.

배경색 추가 코드

from PIL import Image 
import numpy as np
import cv2
import os
import argparse

# 지정한 배경색 만들기
def get_background_color(red, green, blue, img_width, img_high):
    r = np.ones((img_high, img_width, 1), dtype=np.uint8) * red
    g = np.ones((img_high, img_width, 1), dtype=np.uint8) * green
    b = np.ones((img_high, img_width, 1), dtype=np.uint8) * blue
    background = cv2.merge((b,g,r))
    return background

# iamge 폴더, mask 폴더
img_dir = 'imgs'
mask_dir = 'masks'

img_width = 192
img_high = 256
width=256
hight=256
diff = (width - img_width)//2

# 각각의 Color 이름과 RGB 값
Color = ['brown', 'gray', 'orange']
R = [226, 231, 215]
G = [211, 226, 79]
B = [190, 222, 41]

for img_name in os.listdir(img_dir):
    img = cv2.imread(img_dir + '\\' + img_name, cv2.IMREAD_COLOR)
    mask = cv2.imread(mask_dir + '\\' + img_name, cv2.IMREAD_COLOR) 
    mask_ori = mask.copy()
    mask_ori = cv2.cvtColor(mask_ori, cv2.COLOR_BGR2GRAY)

    mask = mask[:, diff:diff+img_width,:]
    mask_not = np.where(mask > 100,0,1)      # 0과 1로 변경 (옷 부분이 0, 배경이 1)
    mask_not = np.array(mask_not, np.uint8)  

    mask = np.where(mask > 100,1,0) # 0과 1로 변경 (옷 부분이 1, 배경이 0)
    mask = np.array(mask, np.uint8)
    for i in range(len(Color)):
        image = img[:, diff:diff+img_width,:]
        background = get_background_color(R[i], G[i], B[i], img_width, img_high)
        result_img = mask*image + mask_not*background  # 기존 이미지에서 옷 부분만 유지하고, 새로운 배경에서 배경만 유지
        img[:, diff:diff+img_width,:] = result_img

        # cv2.imshow("R", img)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()
        name = img_name.split('.png')[0] + '_' + Color[i] + '.png'
        
        cv2.imwrite(img_dir + '/' + name, img)
        cv2.imwrite(mask_dir + '/' + name, mask_ori)

다음과 같이 기존 흰색 배경만 있던 이미지에서 배경을 추가하여 데이터셋을 증강시켰다.

새로운 이미지를 추가한 데이터셋으로 학습을 다시하였더니 배경색이 흰색이 아니여도 결과가 제대로 나오는 것을 확인할 수 있었다.