HR-VITONで、高解像度の仮想試着を実現する
今回ご紹介するのは、高解像度の仮想試着を実現するHR-VITONという技術です。試着条件ジェネレータ(Try-on Condition Generator)を設けることで、従来からあるワーピングとセグメンテーションマップの不整合による歪みを減少させています。
*この論文は、2022.6に提出されました。
2.HR-VITONとは?
従来、仮想試着は衣服の画像を人の体に合わせるワーピングと最終的な画像をガイドするセグメンテーションマップの生成2つによって構成されています。このとき、ワーピングモジュールとセグメンテーションマップは独立して動くため不整合が発生する場合があり、特に体の一部分が衣服を塞ぐとアーティファクト(歪み)が発生していました。
そこで、HR-VITONは、人の画像(I)と衣服の画像(c)を前処理(Pre-Process)した上で、試着条件ジェネレータ(Try-on Condition Generator)でワーピングとセグメンテーションマップの生成を融合して行うことでアーティファクトを減少させ、高解像度での仮想試着を実現しています。
3.コード
コードはGoogle Colabで動かす形にしてGithubに上げてありますので、それに沿って説明して行きます。自分で動かしてみたい方は、この「リンク」をクリックし表示されたノートブックの先頭にある**「Colab on Web」**ボタンをクリックすると動かせます。
まず、セットアップを行います。
#@title #**セットアップ** # githubからコードをコピー ! git clone https://github.com/cedro3/HR-VITON.git %cd HR-VITON # ライブラリーのインストール ! pip install tensorboardX ! pip install torchgeometry # 学習済みパラメータのダウンロード import gdown gdown.download('https://drive.google.com/uc?id=1XJTCdRBOPVgVTmqzhVGFAgMm2NLkw5uQ', 'mtviton_step_100000.pth', quiet=False) gdown.download('https://drive.google.com/uc?id=1BkSA8UJo-6eOkKcXTFOHK80Esc4vBmVC', 'gen_step_110000.pth', quiet=False) # テストデータのダウンロード gdown.download('https://drive.google.com/uc?id=1CcgCubhLc9iF6jGACdUgGhTDWMC7Gjzr', 'test.zip', quiet=False) ! unzip test.zip # 関数のインポート from function import * # warning非表示 import warnings warnings.filterwarnings("ignore") # ダウンロードフォルダ作成 ! mkdir download
! git clone https://github.com/cedro3/HR-VITON.git ! pip install tensorboardX ! pip install torchgeometry gdown.download('https://drive.google.com/uc?id=1XJTCdRBOPVgVTmqzhVGFAgMm2NLkw5uQ', 'mtviton_step_100000.pth', quiet=False) gdown.download('https://drive.google.com/uc?id=1BkSA8UJo-6eOkKcXTFOHK80Esc4vBmVC', 'gen_step_110000.pth', quiet=False) gdown.download('https://drive.google.com/uc?id=1CcgCubhLc9iF6jGACdUgGhTDWMC7Gjzr', 'test.zip', quiet=False) warnings.filterwarnings("ignore") |
VITON-HDデータセットの中から testデータ(2032個)のみを testフォルダに保存してあります。その中から、人と衣服のデータからそれぞれ seed_number に従ってランダムに5個選び、実行内容を記載した test_pairs.txt を作成します。
#@title #**人と衣服のランダム選択** import glob import random import shutil import os # 乱数シード設定 seed_number = 120#@param {type:"integer"} random.seed(seed_number) # imageのランダム選択 reset_folder('image') image_files = sorted(glob.glob('test/test/image/*.jpg')) cnt = len(image_files) num = random.sample(range(cnt),5) image_names = [] for i in num: shutil.copy(image_files[i], 'image/'+os.path.basename(image_files[i])) image_names.append(image_files[i]) image_names.sort() display_pic('image') # clothのランダム選択 reset_folder('cloth') cloth_files = sorted(glob.glob('test/test/cloth/*.jpg')) cnt = len(cloth_files) num = random.sample(range(cnt),5) cloth_names =[] for j in num: shutil.copy(cloth_files[j], 'cloth/'+os.path.basename(cloth_files[j])) cloth_names.append(cloth_files[j]) cloth_names.sort() display_pic('cloth') # txt作成 if os.path.isfile('test/test/test_pairs.txt'): os.remove('test/test/test_pairs.txt') f = open('test/test/test_pairs.txt', 'w', encoding='UTF-8') for image_name in image_names: for cloth_name in cloth_names: f.write(os.path.basename(image_name)+' ') f.write(os.path.basename(cloth_name)+'\n') f.close()
seed_number = 120#@param {type:"integer"} image_files = sorted(glob.glob('test/test/image/*.jpg')) num = random.sample(range(cnt),5) shutil.copy(image_files[i], 'image/'+os.path.basename(image_files[i])) image_names.append(image_files[i]) cloth_files = sorted(glob.glob('test/test/cloth/*.jpg')) num = random.sample(range(cnt),5) shutil.copy(cloth_files[j], 'cloth/'+os.path.basename(cloth_files[j])) cloth_names.append(cloth_files[j]) if os.path.isfile('test/test/test_pairs.txt'): os.remove('test/test/test_pairs.txt') f = open('test/test/test_pairs.txt', 'w', encoding='UTF-8') for image_name in image_names: for cloth_name in cloth_names: f.write(os.path.basename(image_name)+' ') f.write(os.path.basename(cloth_name)+'\n') |
それでは、表示された人と衣服のデータ5×5=25個の組み合わせで仮想試着を実行します。
#@title #**仮想試着** reset_folder('output') ! python3 test_generator.py --occlusion --tocg_checkpoint './mtviton_step_100000.pth'\ --gpu_ids 0\ --gen_checkpoint './gen_step_110000.pth'\ --datasetting unpaired\ --dataroot './test'\ --data_list './test/test_pairs.txt' clear_output() display_pic_png('output/mtviton_step_100000.pth/test/unpaired/generator/output')
! python3 test_generator.py --occlusion --tocg_checkpoint './mtviton_step_100000.pth'\ --gen_checkpoint './gen_step_110000.pth'\ --data_list './test/test_pairs.txt' display_pic_png('output/mtviton_step_100000.pth/test/unpaired/generator/output') |
試着画像を2種類のスタイル(バー、スクエア)の動画に変換します。
#@title #**動画の作成** import cv2 import glob reset_folder('movie1') reset_folder('movie2') result_files = sorted(glob.glob('output/mtviton_step_100000.pth/test/unpaired/generator/output/*.png')) cnt = 0 black = cv2.imread('black.jpg') for image_name in image_names: for cloth_name in cloth_names: left = cv2.imread(image_name) center = cv2.imread(cloth_name) right = cv2.imread(result_files[cnt]) tmp = cv2.hconcat([left, center]) img1 = cv2.hconcat([tmp, right]) cv2.imwrite('movie1/'+str(cnt).zfill(4)+'.jpg', img1) up = cv2.hconcat([black, center]) down = cv2.hconcat([left, right]) img2 = cv2.vconcat([up, down]) cv2.imwrite('movie2/'+str(cnt).zfill(4)+'.jpg', img2) cnt +=1 ! ffmpeg -y -r 1 -i movie1/%04d.jpg -vcodec libx264 -pix_fmt yuv420p -loglevel error output1.mp4 ! ffmpeg -y -r 1 -i movie2/%04d.jpg -vcodec libx264 -pix_fmt yuv420p -loglevel error output2.mp4 display_mp4('output1.mp4')
result_files = sorted(glob.glob('output/mtviton_step_100000.pth/test/unpaired/generator/output/*.png')) black = cv2.imread('black.jpg') for image_name in image_names: for cloth_name in cloth_names: left = cv2.imread(image_name) center = cv2.imread(cloth_name) right = cv2.imread(result_files[cnt]) tmp = cv2.hconcat([left, center]) img1 = cv2.hconcat([tmp, right]) cv2.imwrite('movie1/'+str(cnt).zfill(4)+'.jpg', img1) up = cv2.hconcat([black, center]) down = cv2.hconcat([left, right]) img2 = cv2.vconcat([up, down]) cv2.imwrite('movie2/'+str(cnt).zfill(4)+'.jpg', img2) ! ffmpeg -y -r 1 -i movie1/%04d.jpg -vcodec libx264 -pix_fmt yuv420p -loglevel error output1.mp4 ! ffmpeg -y -r 1 -i movie2/%04d.jpg -vcodec libx264 -pix_fmt yuv420p -loglevel error output2.mp4 display_mp4('output1.mp4') |
表示されているのは、バーです。
作成した動画をダウンロードします。square のチェックボックスにチェックを入れるとスクエアで、チェックを入れないとバーで動画をダウンロードします。ここでは、スクエアでダウンロードします。なお、ブラウザは google chrome でないと正常に動作しませんので、ご注意を。
#@title #**動画のダウンロード** from google.colab import files square = True #@param {type:"boolean"} if square == True: shutil.copy('output2.mp4', 'download/'+str(seed_number)+'_s.mp4') files.download('download/'+str(seed_number)+'_s.mp4') else: shutil.copy('output1.mp4', 'download/'+str(seed_number)+'.mp4') files.download('download/'+str(seed_number)+'.mp4')
from google.colab import files square = True #@param {type:"boolean"} shutil.copy('output2.mp4', 'download/'+str(seed_number)+'_s.mp4') files.download('download/'+str(seed_number)+'_s.mp4') shutil.copy('output1.mp4', 'download/'+str(seed_number)+'.mp4') files.download('download/'+str(seed_number)+'.mp4') |
仮想試着もかなり実用的なレベルになって来たようです。
では、また。
(オリジナルgithub)https://github.com/sangyun884/HR-VITON
(twitter投稿)
Google Colabgoogle colaboratoryGooglecolaboratoryHR-VITONTry-on Condition Generatorセグメンテーションマップワーピング仮想試着試着条件ジェネレータ高解像度