# -*- coding: utf-8 -*-
##------------------------------------------
## BlendGAN demo program Ver 0.01
##
## 2024.08.23 Masahiro Izutsu
##------------------------------------------
## blendgan_demo.py
# 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'
import warnings
warnings.simplefilter('ignore')
from torch.cuda import is_available
gpu_d = is_available() # GPU 確認
# インポート＆初期設定
import os
import glob
import shutil
import argparse
import cv2
import numpy as np
from PIL import Image
from tqdm import tqdm
import dlib
from alignment import align_face
import matplotlib.pyplot as plt
from skimage.transform import resize
import glob
import my_logging
import my_imagetool
# 定数定義
DEF_IMAGE = './pic/66.jpg'
RESULT_PATH = './results'
ALIGN_DIR = './align'
PIC_DIR = './pic'
TEST_IMAGE_DIR = './test_imgs/original_imgs'
TEST_STYLE_DIR = './test_imgs/style_imgs'
STYLE_TRANSFER = 'style_transfer'
OUT_IMAGES = 'images'
OUT_MOVIE = './output.mp4'
# タイトル
title = 'BlendGAN demo program Ver. 0.01'
# Parses arguments for the application
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--source_image", default='', help="path to source image")
parser.add_argument("--result_path", default=RESULT_PATH, help="path to output")
parser.add_argument('--image_size', default = 1024, help = 'image size')
parser.add_argument("--align", action="store_false", help="make align image flag")
parser.add_argument("--cpu", dest="cpu", action="store_true", help="cpu mode.")
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(args, title):
print('\n' + GREEN + title + ': Starting application...' + NOCOLOR)
print('\n - ' + YELLOW + 'source_dir : ' + NOCOLOR, args.source_image)
print(' - ' + YELLOW + 'result_path : ' + NOCOLOR, args.result_path)
print(' - ' + YELLOW + 'image_size : ' + NOCOLOR, args.image_size)
print(' - ' + YELLOW + 'align : ' + NOCOLOR, args.align)
print(' - ' + YELLOW + 'cpu : ' + NOCOLOR, args.cpu)
print(' - ' + YELLOW + 'log : ' + NOCOLOR, args.log)
print(' ')
# フォルダー初期化
def reset_folder(path):
if os.path.isdir(path):
shutil.rmtree(path)
os.makedirs(path,exist_ok=True)
# ファイル初期化
def reset_file(path):
if os.path.exists(path):
os.remove(path)
# フォルダー内の画像ファイルリストの取得
def get_image_path(folder):
files = os.listdir(folder)
files.sort()
return files
# 顔画像の切り出し
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 = PIC_DIR, align_dir = ALIGN_DIR):
reset_folder(ALIGN_DIR) # align_dir を初期化する
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 folder_image(folder, save_path = '', pixel_size = (256,256), dpi = 64, xn = 10):
files = os.listdir(folder)
files.sort()
n = len(files)
yn = n // xn + 1
if n < xn:
xn = n
# ピクセル → インチ変換
x_inch = pixel_size[0] / dpi
y_inch = pixel_size[1] / dpi
fig = plt.figure(figsize = (x_inch * xn, y_inch * yn + 0.2), dpi = dpi)
fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
for i, file in enumerate(files):
img = Image.open(folder+'/'+file)
images = np.asarray(img)
# 正方形にする
img_h, img_w = images.shape[:2]
if img_h != img_w:
images = my_imagetool.frame_square(images)
images = resize(images, pixel_size)[..., :3]
ax = fig.add_subplot(yn, xn, 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 make_align(image_file, out_path = RESULT_PATH, disp_f = True):
base_dir_pair = os.path.split(image_file)
image_path = base_dir_pair[0]
path = out_path + '/' + base_dir_pair[1] + '_imagelist.jpg'
path0 = './tmp_align.jpg'
path1 = './tmp_pic.jpg'
# 画像フォルダの 顔画像を切り出す
align_images(image_dir = image_path, align_dir = ALIGN_DIR)
# フォルダ内の画像一覧作成
folder_image(ALIGN_DIR, save_path = path0)
folder_image(PIC_DIR, save_path = path1)
images = []
images.append(cv2.imread(path1))
images.append(cv2.imread(path0))
h, w = images[0].shape[:2]
ds_image = my_imagetool.make_tileimage(images, xmax = w, ymax = h * 2)
my_imagetool.image_disp(ds_image, winname = 'align - vec_pic', dispf = disp_f, save_path = path, maxsize = 1024, wait_s = 2)
# フォルダ内の画像のリサイズ
def resize_fol(trans_dir, out_dir, logger = None):
files = glob.glob(trans_dir + '*.jpg')
files.sort()
for i, file in enumerate(files):
img = cv2.imread(file)
img_resize = cv2.resize(img, dsize=(1536, 512))
o_path = out_dir + '/' + str(i).zfill(3) + '.jpg'
logger.debug(f'Image resize ... {file} → {o_path}')
cv2.imwrite(o_path, img_resize)
# 画像を動画に変換
def make_movie(source_image, result_path, logger = None, disp_f = True):
base_dir_pair = os.path.split(source_image)
s_name, _ = os.path.splitext(base_dir_pair[1])
o_path = result_path + '/' + s_name + '_' + STYLE_TRANSFER + '.mp4'
s_path = result_path + '/' + OUT_IMAGES + '/%3d.jpg'
if not os.path.exists(o_path):
command = f'ffmpeg -r 3/2 -i {s_path} -vcodec libx264 -pix_fmt yuv420p {OUT_MOVIE} -loglevel quiet -y'
logger.info(command)
os.system(command)
# out_dir フォルダへ名前を付けてコピー
logger.info(f' making movie... → {o_path}')
shutil.copy(OUT_MOVIE, o_path)
if disp_f:
my_imagetool.image2disp(o_path)
# 切り出し画像(align_image)の存在確認
def check_align_image(img_path, align_path, logger=None):
files_image = [
f for f in os.listdir(img_path) if os.path.isfile(os.path.join(img_path, f))
]
logger.debug(f'image = {files_image}')
if not os.path.isdir(align_path):
return False
files_align = [
f for f in os.listdir(align_path) if os.path.isfile(os.path.join(align_path, f))
]
logger.debug(f'align = {files_align}')
flag = True
for img_path in files_image:
name, ext= os.path.splitext(img_path)
path = align_path + '/' + name + ext
if not os.path.isfile(path):
flag = False
logger.info(f' check_align_image = {flag}')
return flag
# ** main関数 **
def main(opt):
# アプリケーション・ログ設定
module = os.path.basename(__file__)
module_name = os.path.splitext(module)[0]
logger = my_logging.get_module_logger_sel(module_name, int(opt.log))
# 切り出し画像一覧作成
flg = check_align_image(PIC_DIR, ALIGN_DIR, logger) and opt.align
opt.align = flg
if not opt.align:
make_align(opt.source_image, out_path = RESULT_PATH, disp_f = True)
base_dir_pair = os.path.split(opt.source_image)
align_file = ALIGN_DIR + '/' + base_dir_pair[1]
org_file = TEST_IMAGE_DIR + '/' + base_dir_pair[1]
if not opt.cpu:
# 画像ファイル指定
reset_folder(TEST_IMAGE_DIR) # original_images を初期化する
logger.info(f' copy to image... {align_file} → {org_file}')
shutil.copy(align_file, org_file)
# Style transfer を実行
trans_dir = opt.result_path + '/' + STYLE_TRANSFER + '/'
reset_folder(trans_dir) # style_transfer を初期化する
checkpoint = './pretrained_models/blendgan.pt'
psp_encord = './pretrained_models/psp_encoder.pt'
c_dir = TEST_IMAGE_DIR + '/'
s_dir = TEST_STYLE_DIR + '/'
command = f'python style_transfer_folder2.py --size {opt.image_size} --ckpt {checkpoint} --psp_encoder_ckpt {psp_encord} --style_img_path {s_dir} --input_img_path {c_dir} --outdir {trans_dir}'
logger.info(command)
os.system(command)
# 画像のリサイズ
# files = glob.glob('results/style_transfer/*.jpg')
outimage_dir = opt.result_path + '/' + OUT_IMAGES
reset_folder(outimage_dir) # outimage_dir を初期化する
resize_fol(trans_dir, outimage_dir, logger = logger)
# 画像を動画に変換
make_movie(opt.source_image, opt.result_path, logger = logger)
# main関数エントリーポイント(実行開始)
if __name__ == '__main__':
import datetime
import my_dialog
parser = parse_args()
opt = parser.parse_args()
if not opt.cpu:
opt.cpu = not gpu_d
if len(opt.source_image) == 0:
opt.source_image = my_dialog.select_image_file('元画像 ', PIC_DIR)
if len(opt.source_image) == 0:
exit(0)
start_time = datetime.datetime.now() # 時間計測開始
display_info(opt, title)
main(opt)
# 経過時間
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)
print('\nFinished.\n')
