私的AI研究会 > PAMA
PAMAで、スタイル転送をやってみる
機械学習で画像に別の画像のスタイルを転送するのを、上記サイトの手順に従って検証してみる。
現在の「Google Colaboratory」環境で動作する。Progressive Attentional Manifold Alignment(PAMA) 概念図
#@title セットアップ
%%capture
# githubからコードをコピー
! git clone https://github.com/cedro3/PAMA.git
%cd PAMA
# 学習パラメータのダウンロード
! pip install --upgrade gdown
import gdown
# original
gdown.download('https://drive.google.com/uc?id=1rPB_qnelVVSad6CtadmhRFi0PMI_RKdy', 'original_PAMA.zip', quiet=False)
! unzip original_PAMA.zip
# consistency
gdown.download('https://drive.google.com/uc?id=1IrggOiutiZceJCrEb24cLnBjeA5I3N1D', 'PAMA_without_color.zip', quiet=False)
! unzip PAMA_without_color.zip
# color
gdown.download('https://drive.google.com/uc?id=1HXet2u_zk2QCVM_z5Llg2bcfvvndabtt', 'PAMA_1.5_color.zip', quiet=False)
! unzip PAMA_1.5_color.zip
# content
gdown.download('https://drive.google.com/uc?id=13m7Lb9xwfG_DVOesuG9PyxDHG4SwqlNt', 'PAMA_1.5_content.zip', quiet=False)
! unzip PAMA_1.5_content.zip
import os
import shutil
# downloadフォルダ作成
os.makedirs('download', exist_ok=True)
#学習済みパラメータの設定関数
def select_param(select):
# checkpointsフォルダーリセット
if os.path.isdir('checkpoints'):
shutil.rmtree('checkpoints')
if select == 'original':
shutil.copytree('original_PAMA', 'checkpoints')
if select == 'consistency':
shutil.copytree('PAMA_without_color', 'checkpoints')
shutil.copy('original_PAMA/encoder.pth', 'checkpoints/encoder.pth')
if select == 'color':
shutil.copytree('PAMA_1.5_color', 'checkpoints')
shutil.copy('original_PAMA/encoder.pth', 'checkpoints/encoder.pth')
if select == 'content':
shutil.copytree('PAMA_1.5_content', 'checkpoints')
shutil.copy('original_PAMA/encoder.pth', 'checkpoints/encoder.pth')
print('checkpoint is '+select)
# 画像連結関数
from PIL import Image
def get_concat_h(im1, im2):
dst = Image.new('RGB', (im1.width + im2.width, im1.height))
dst.paste(im1, (0, 0))
dst.paste(im2, (im1.width, 0))
return dst
def edit_pic(file1, file2):
f1 = Image.open(file1)
f2 = Image.open(file2)
f1_resize = f1.resize((f2.width, f2.height))
get_concat_h(f1_resize, f2).save('disp.jpg')
#@title 画像のスタイル転送
# 設定
checkpoint = 'consistency' #@param ["original", "consistency", "color", "content"] {allow-input: true}
content = '11.jpg'#@param {type:"string"}
style = '03.jpg'#@param {type:"string"}
# 学習済みパラメータの選択
select_param(checkpoint)
# スタイル転送
! python main.py eval --content content/$content --style style/$style
# 変換前後の画像連結
edit_pic('content/'+content, 'ics.jpg')
# 連結画像表示
from google.colab.patches import cv2_imshow
import cv2
img_cv=cv2.imread('disp.jpg')
cv2_imshow(img_cv)
#@title 画像のダウンロード
import os
import shutil
from google.colab import files
content_name = os.path.splitext(content)
style_name = os.path.splitext(style)
file_name = 'download/'+content_name[0]+'_'+style_name[0]+'_'+checkpoint+'.jpg'
shutil.copy('disp.jpg', file_name)
files.download(file_name)
#@title 動画のスタイル転送
#@markdown ・videoは音声付きとして下さい
# 学習済みパラメータの選択
checkpoint = 'consistency' #@param ["original", "consistency", "color", "content"] {allow-input: true}
select_param(checkpoint)
# --- 動画を静止画にバラす ---
video = 'open_house.mp4' #@param {type:"string"}
video_file = 'video/'+video
import os
import shutil
import cv2
# flamesフォルダーリセット
if os.path.isdir('images'):
shutil.rmtree('images')
os.makedirs('images', exist_ok=True)
def video_2_images(video_file= video_file, # ビデオの指定
image_dir='images/',
image_file='%s.jpg'):
# Initial setting
i = 0
interval = 1
length = 3000 # 最大フレーム数
cap = cv2.VideoCapture(video_file)
fps = cap.get(cv2.CAP_PROP_FPS) # fps取得
while(cap.isOpened()):
flag, frame = cap.read()
if flag == False:
break
if i == length*interval:
break
if i % interval == 0:
cv2.imwrite(image_dir+image_file % str(int(i/interval)).zfill(6), frame)
i += 1
cap.release()
return fps, i, interval
fps, i, interval = video_2_images()
#print('fps = ', fps)
#print('flames = ', i)
#print('interval = ', interval)
# --- 動画のスタイル転送 ---
# style_movieフォルダーリセット
if os.path.isdir('style_movie'):
shutil.rmtree('style_movie')
os.makedirs('style_movie', exist_ok=True)
# 指定したstyleをstyle_movieフォルダへコピー
import shutil
style = 'ani_004.jpg'#@param {type:"string"}
shutil.copy('style/'+style, 'style_movie/'+style)
# icsフォルダーリセット
if os.path.isdir('ics'):
shutil.rmtree('ics')
os.makedirs('ics', exist_ok=True)
#@title 動画の再生
from IPython.display import HTML
from base64 import b64encode
mp4 = open('./output.mp4', 'rb').read()
data_url = 'data:video/mp4;base64,' + b64encode(mp4).decode()
HTML(f"""
<video width="70%" height="70%" controls>
<source src="{data_url}" type="video/mp4">
</video>""")
#@title 動画のダウンロード
import os
import shutil
from google.colab import files
content_name = os.path.splitext(video)
style_name = os.path.splitext(style)
file_name = 'download/'+content_name[0]+'_'+style_name[0]+'_'+checkpoint+'.mp4'
shutil.copy('output.mp4', file_name)
files.download(file_name)