# -*- coding: utf-8 -*-
##------------------------------------------
## StyleGAN e4e demo program Ver 0.01
## (encoder4editing)
## StyleGANを使った画像編集をe4eで高速化する
## http://http://cedro3.com/ai/e4e/
##
## 2024.08.05 Masahiro Izutsu
##------------------------------------------
## e4e_demo.py
import warnings
warnings.simplefilter('ignore')
# Color Escape Code ---------------------------
GREEN = '\033[1;32m'
RED = '\033[1;31m'
NOCOLOR = '\033[0m'
YELLOW = '\033[1;33m'
CYAN = '\033[1;36m'
BLUE = '\033[1;34m'
# インポート＆初期設定
from argparse import Namespace
import time
import os
import sys
import numpy as np
from PIL import Image
import torch
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import shutil
from tqdm import tqdm
import dlib
from utils.common import tensor2im
from models.psp import pSp # we use the pSp framework to load the e4e encoder.
from editings import latent_editor
from tqdm import trange
from utils.alignment import align_face
import argparse
from skimage.transform import resize
import ffmpeg
import my_imagetool
import my_logging
# 定数定義
IMAGE_DIR = './images'
RESULT_PATH = './results'
MODEL_PATH = './pretrained_models/e4e_ffhq_encode.pt'
ALIGN_DIR = './align'
VEC_PIC_DIR = './vec_pic'
VEC_DIR = './vec'
PIC_DIR = './pic'
OUT_MOVIE = './output.mp4'
# タイトル
title = 'StyleGAN e4e demo program Ver 0.01'
# Parses arguments for the application
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--direction", default='pose', type=str, choices=['age', 'pose', 'smile', 'age+pose'], help="select 'age'/'pose'/'smile'/'age+pose'")
parser.add_argument("--source_image", default=IMAGE_DIR, help="path to source image")
parser.add_argument("--result_image", default=RESULT_PATH, help="path to output")
parser.add_argument('--log', metavar = 'LOG', default = '3', help = 'Log level(-1/0/1/2/3/4/5) Default value is \'3\'')
return parser
# 基本情報の表示
def display_info(opt, title):
print('\n' + GREEN + title + ': Starting application...' + NOCOLOR)
print('\n - ' + YELLOW + 'source_image : ' + NOCOLOR, opt.source_image)
print(' - ' + YELLOW + 'result_image : ' + NOCOLOR, opt.result_image)
print(' - ' + YELLOW + 'log : ' + NOCOLOR, opt.log)
print('\n - ' + 'model path : ', MODEL_PATH)
print(' - ' + 'align directory : ', ALIGN_DIR)
print(' - ' + 'vector pict directory : ', VEC_PIC_DIR)
print(' - ' + 'latent variable direct. : ', VEC_DIR)
print(' - ' + 'still image directory : ', PIC_DIR)
print(' ')
# モデルに学習済みパラメータをロード
def load_model(model_path = MODEL_PATH):
ckpt = torch.load(model_path, map_location='cpu')
opts = ckpt['opts']
opts['checkpoint_path'] = model_path
opts= Namespace(**opts)
net = pSp(opts)
net.eval()
net.cuda()
print('Model successfully loaded!')
return net
# --- 顔画像の切り出し ---
def run_alignment(image_path):
predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
aligned_image = align_face(filepath=image_path, predictor=predictor)
return aligned_image
# 画像フォルダの 顔画像を切り出す
def align_images(image_dir = IMAGE_DIR, align_dir = ALIGN_DIR):
# align_dir を初期化する
if os.path.isdir(align_dir):
shutil.rmtree(align_dir)
os.makedirs(align_dir, exist_ok = True)
files = sorted(os.listdir(image_dir))
for i, file in enumerate(tqdm(files)):
input_path = image_dir + '/' + file
align_path = align_dir + '/' + file
input_image = run_alignment(input_path)
input_image.resize((256,256))
input_image.save(align_path)
# 画像の潜在変数の推定
def make_latent_variable(net, align_dir = ALIGN_DIR, vec_pic_dir = VEC_PIC_DIR, vec_dir = VEC_DIR):
# フォルダの初期化
if os.path.isdir(vec_pic_dir):
shutil.rmtree(vec_pic_dir)
os.makedirs(vec_pic_dir, exist_ok = True)
if os.path.isdir(vec_dir):
shutil.rmtree(vec_dir)
os.makedirs(vec_dir, exist_ok = True)
img_transforms = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
files = sorted(os.listdir(align_dir))
for i, file in enumerate(tqdm(files)):
align_path = align_dir + '/' + file
vec_pic_path = vec_pic_dir + '/' + file
vec_path = vec_dir + '/' + file[:-4] + '.pt'
input_image = Image.open(align_path)
transformed_image = img_transforms(input_image)
with torch.no_grad():
images, latents = net(transformed_image.unsqueeze(0).to('cuda').float(), randomize_noise=False, return_latents=True)
result_image, latent = images[0], latents[0]
tensor2im(result_image).save(vec_pic_path) # vec_pic 保存
torch.save(latents, vec_path) # vec 保存
# フォルダー内の一覧画像の作成
def folder_image(folder, save_path = '', pixel_size = (256,256), dpi = 64):
# ピクセル → インチ変換
x_inch = pixel_size[0] / dpi
y_inch = pixel_size[1] / dpi
fig = plt.figure(figsize = (x_inch * 10, y_inch * 1 + 1), dpi = dpi)
fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
files = os.listdir(folder)
files.sort()
for i, file in enumerate(files):
img = Image.open(folder+'/'+file)
images = np.asarray(img)
images = resize(images, pixel_size)[..., :3]
ax = fig.add_subplot(1, 10, i+1, xticks=[], yticks=[])
image_plt = np.array(images)
ax.imshow(image_plt)
ax.set_xlabel(folder+'/'+file, fontsize=15)
if len(save_path) > 0:
plt.savefig(save_path)
plt.close()
# 静止画の生成
def still_image_generation(net, vec_dir = VEC_DIR, latent = '03.pt', direction = 'pose', min = -50, max = 50):
pic_dir = PIC_DIR # 静止画フォルダ
logger.info(f'\n{CYAN} latent = {latent}, direction = {direction}, min = {min}, max = {max}{NOCOLOR}')
# フォルダの初期化
if os.path.isdir(pic_dir):
shutil.rmtree(pic_dir)
os.makedirs(pic_dir, exist_ok=True)
vec_path = vec_dir + '/' + latent
latents = torch.load(vec_path)
editor = latent_editor.LatentEditor(net.decoder, False)
interfacegan_directions = {
'age': 'editings/interfacegan_directions/age.pt',
'smile': 'editings/interfacegan_directions/smile.pt',
'pose': 'editings/interfacegan_directions/pose.pt',
'age+pose': 'editings/interfacegan_directions/age+pose.pt'
}
interfacegan_direction = torch.load(interfacegan_directions[direction]).cuda()
cnt = 0
for i in trange(0, min, -1, desc='0 -> min'):
result = editor.apply_interfacegan(latents, interfacegan_direction, factor=i).resize((512,512))
result.save(pic_dir + '/' + str(cnt).zfill(6) + '.jpg')
cnt +=1
for i in trange(min, max, desc='min -> max'):
result = editor.apply_interfacegan(latents, interfacegan_direction, factor=i).resize((512,512))
result.save(pic_dir + '/' + str(cnt).zfill(6) + '.jpg')
cnt +=1
for i in trange(max, 0, -1, desc='max -> 0'):
result = editor.apply_interfacegan(latents, interfacegan_direction, factor=i).resize((512,512))
result.save(pic_dir + '/' + str(cnt).zfill(6) + '.jpg')
cnt +=1
return pic_dir
# mp4 動画の作成
def make_movie(pic_dir, latent = '03.pt', direction = 'pose', out_dir = 'movie'):
# 既に ファイルがあれば削除する
if os.path.exists(OUT_MOVIE):
os.remove(OUT_MOVIE)
# pic フォルダーの静止画から動画を作成
ffmpeg.input(pic_dir + '/%6d.jpg', framerate = 30).output(OUT_MOVIE, vcodec = 'libx264', r = 30).run(quiet=True)
# out_dir フォルダへ名前を付けてコピー
shutil.copy(OUT_MOVIE, out_dir + '/' + direction + '_' + latent[:-3] + '.mp4')
# 動画を表示
my_imagetool.image2disp(OUT_MOVIE)
# main関数エントリーポイント(実行開始)
if __name__ == "__main__":
import datetime
parser = parse_args()
opt = parser.parse_args()
# アプリケーション・ログ設定
module = os.path.basename(__file__)
module_name = os.path.splitext(module)[0]
logger = my_logging.get_module_logger_sel(module_name, int(opt.log))
display_info(opt, title)
start_time = datetime.datetime.now() # 時間計測開始
os.makedirs(opt.result_image, exist_ok = True)
base_dir_pair = os.path.split(opt.source_image)
path0 = opt.result_image + '/' + base_dir_pair[1] + '_align.jpg'
path1 = opt.result_image + '/' + base_dir_pair[1] + '_vec_pic.jpg'
image_path = opt.source_image
out_path = opt.result_image
# モデルに学習済みパラメータをロード
net = load_model(model_path = MODEL_PATH)
# 画像フォルダの 顔画像を切り出す
align_images(image_dir = image_path, align_dir = ALIGN_DIR)
# 画像の潜在変数の推定
make_latent_variable(net, align_dir = ALIGN_DIR, vec_pic_dir = VEC_PIC_DIR, vec_dir = VEC_DIR)
# フォルダ内の画像一覧作成
folder_image(ALIGN_DIR, save_path = path0)
my_imagetool.image2disp(path0, maxsize = 1024)
folder_image(VEC_PIC_DIR, save_path = path1)
my_imagetool.image2disp(path1, maxsize = 1024)
latent = '03.pt' # type:"string"
direction = 'age' # "age", "pose", "smile", "age+pose"
min = -50 # min:-50, max:0, step:10
max = 50 # min:0, max:50, step:10
pic_dir = still_image_generation(net, vec_dir = VEC_DIR, latent = latent, direction = direction, min = min, max = max)
make_movie(pic_dir, latent = latent, direction = direction, out_dir = out_path)
latent = '03.pt' # type:"string"
direction = 'pose' # "age", "pose", "smile", "age+pose"
min = -50 # min:-50, max:0, step:10
max = 50 # min:0, max:50, step:10
pic_dir = still_image_generation(net, vec_dir = VEC_DIR, latent = latent, direction = direction, min = min, max = max)
make_movie(pic_dir, latent = latent, direction = direction, out_dir = out_path)
latent = '03.pt' # type:"string"
direction = 'smile' # "age", "pose", "smile", "age+pose"
min = 0 # min:-50, max:0, step:10
max = 30 # min:0, max:50, step:10
pic_dir = still_image_generation(net, vec_dir = VEC_DIR, latent = latent, direction = direction, min = min, max = max)
make_movie(pic_dir, latent = latent, direction = direction, out_dir = out_path)
latent = '03.pt' # type:"string"
direction = 'age+pose' # "age", "pose", "smile", "age+pose"
min = -50 # min:-50, max:0, step:10
max = 50 # min:0, max:50, step:10
pic_dir = still_image_generation(net, vec_dir = VEC_DIR, latent = latent, direction = direction, min = min, max = max)
make_movie(pic_dir, latent = latent, direction = direction, out_dir = out_path)
# 経過時間
end_time = datetime.datetime.now()
print(start_time.strftime('\nprocessing start >>\t %Y/%m/%d %H:%M:%S'))
print(end_time.strftime('processing end >>\t %Y/%m/%d %H:%M:%S'))
print('processing time >>\t', end_time - start_time)
logger.info('\nFinished.\n')
