HR-VITONで、高解像度の仮想試着を実現する

1.はじめに

　今回ご紹介するのは、高解像度の仮想試着を実現するHR-VITONという技術です。試着条件ジェネレータ（Try-on Condition Generator）を設けることで、従来からあるワーピングとセグメンテーションマップの不整合による歪みを減少させています。

＊この論文は、2022.6に提出されました。

2.HR-VITONとは？

　従来、仮想試着は衣服の画像を人の体に合わせるワーピングと最終的な画像をガイドするセグメンテーションマップの生成２つによって構成されています。このとき、ワーピングモジュールとセグメンテーションマップは独立して動くため不整合が発生する場合があり、特に体の一部分が衣服を塞ぐとアーティファクト（歪み）が発生していました。

　そこで、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

#@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

　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()

#@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()

　それでは、表示された人と衣服のデータ５×５＝２５個の組み合わせで仮想試着を実行します。

#@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')

#@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')

　　試着画像を２種類のスタイル（バー、スクエア）の動画に変換します。

#@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')

#@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')

　表示されているのは、バーです。

　作成した動画をダウンロードします。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')

#@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')

　仮想試着もかなり実用的なレベルになって来たようです。

　では、また。

（オリジナルgithub）https://github.com/sangyun884/HR-VITON

（twitter投稿）

ブログを書きました！

　今回ご紹介するのは、高解像度の仮想試着を実現するHR-VITONという技術です。
　新たに試着条件ジェネレータを設けることで、従来からあるワーピングとセグメンテーションマップの不整合による歪みを減少させています。

ブログ：https://t.co/lhlTlm24oQ pic.twitter.com/S5Zv9MlGzS
— cedro (@jun40vn) July 4, 2022

初めまして。
いつも興味深い技術を配信くださり、楽しくブログを拝見させていただいております。
取り扱われている技術に関してお伺いしたいことがあり、コメントいたしました。

cedro様はどのようにしてブログに掲載されているような技術をgithub上で探しておられるのでしょうか。先日、私も画像のアニメーション化に挑戦したく、githun上で探したのですがブログに掲載されいますような学習済みの重みやデモコードまで公開してくださっているgithubのページを見つけることができませんでした。cedro様はどのように新しい技術をgithun上で探しておられるのでしょうか。

上記に関しましてお教えいただけると幸いです。よろしくお願い致します。