# -*- coding: utf-8 -*-
##------------------------------------------
## StyleGAN-e4e GUI program Ver 0.01
##
## 2024.08.05 Masahiro Izutsu
##------------------------------------------
## e4e_gui.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'
# インポート＆初期設定
import os
import argparse
import numpy as np
import cv2
from PIL import Image, ImageTk
import PySimpleGUI as sg
import my_logging
import my_dialog
from torch.cuda import is_available
gpu_d = is_available() # GPU 確認
# 定数定義
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'
KEY_ORGIMAGE = '-Org-'
KEY_PRSIMAGE = '-Prs-'
KEY_MIN = '-Min-'
KEY_MAX = '-Max-'
KEY_AGE = '-Age-'
KEY_POSE = '-Pose-'
KEY_SMILE = '-Smile-'
KEY_AGEPOSE = '-AgePose-'
KEY_IMAGESEL = '-Image-'
KEY_IMAGEPTOS = '-New-'
KEY_EXIT = '-Exit-'
KEY_TXTORG = '-Source-'
KEY_TXTPRS = '-Process-'
IMG_CANCAS_SIZE = 512
# タイトル
title = 'StyleGAN-e4e GUI 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='', help="path to source image")
parser.add_argument("--result_image", default=RESULT_PATH, help="path to output")
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(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 + 'cpu : ' + NOCOLOR, opt.cpu)
print(' - ' + YELLOW + 'log : ' + NOCOLOR, opt.log)
print('\n - ' + 'GPU status : ', gpu_d)
print(' - ' + '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 check_vector_image(img_path, vec_path):
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(vec_path):
return False
files_vec = [
f for f in os.listdir(vec_path) if os.path.isfile(os.path.join(vec_path, f))
]
logger.debug(f'vector = {files_vec}')
flag = True
for img_path in files_image:
name, ext= os.path.splitext(img_path)
path = vec_path + '/' + name + '.pt'
if not os.path.isfile(path):
flag = False
logger.debug(f'check_vector_image = {flag}')
return flag
# カレントディレクトリ配下のディレクトリからファイルを選択
def get_image_file(tgt_dir, msg = ''):
# カレントディレクトリ配下の指定のみ有効
import re
image_file = ''
cpath = os.getcwd()
cpath = re.sub(r'\\', '/', cpath)
while True:
image_file = my_dialog.select_image_file(msg, tgt_dir)
if len(image_file) == 0:
break
parir = os.path.split(image_file)
parir = os.path.split(parir[0])
if cpath == parir[0]:
break
return image_file
def main_process(opt, net):
direction = opt.direction
base_dir_pair = os.path.split(opt.source_image)
img_name, ext = os.path.splitext(base_dir_pair[1])
img_dir = base_dir_pair[0] # 画像ディレクトリ
latent = img_name + '.pt' # 潜在ベクトル
result_image = opt.result_image
image_path = VEC_PIC_DIR + '/' + base_dir_pair[1]
save_path = result_image + '/' + direction + '_' + img_name + '.mp4'
radio0 = True if direction == 'age' else False
radio1 = True if direction == 'pose' else False
radio2 = True if direction == 'smile' else False
radio2 = True if direction == 'age+pose' else False
min = 0 if direction == 'smile' else -50
max = 30 if direction == 'smile' else 50
# ウィンドウのテーマ
sg.theme('BlueMono')
# ウィンドウのレイアウト
col_image0 = [
[sg.Canvas(size=(IMG_CANCAS_SIZE, IMG_CANCAS_SIZE), key=KEY_ORGIMAGE)],
[sg.Text(image_path, background_color='LightSteelBlue1', size=(63, 1), key = KEY_TXTORG)]
]
col_image1 = [
[sg.Image(size=(IMG_CANCAS_SIZE, IMG_CANCAS_SIZE), key=KEY_PRSIMAGE)],
[sg.Text(save_path, background_color='LightSteelBlue1', size=(63, 1), key = KEY_TXTPRS)]
]
col_color = [
[sg.Text("min", size=(5, 1)), sg.Slider((-50, 50), min, 10, orientation='h', size=(35, 10), disabled = not gpu_d, key=KEY_MIN, enable_events=True)],
[sg.Text("max", size=(5, 1)), sg.Slider((-50, 50), max, 10, orientation='h', size=(35, 10), disabled = not gpu_d, key=KEY_MAX, enable_events=True)]
]
col_radio = [
[sg.Text("", size=(15, 1)), sg.Text("direction", size=(10, 1)), sg.Radio('age', group_id='direction',enable_events = True, default=radio0, key=KEY_AGE), sg.Radio('pose', group_id='direction',enable_events = True, default=radio1, key=KEY_POSE), sg.Radio('smile', group_id='direction',enable_events = True, default=radio2, key=KEY_SMILE), sg.Radio('age+pose', group_id='direction',enable_events = True, default=radio2, key=KEY_AGEPOSE)],
[sg.Text("", size=(40, 1)), sg.Button('Image', size=(10, 1), key=KEY_IMAGESEL), sg.Button('Process', size=(10, 1), disabled = not gpu_d, key=KEY_IMAGEPTOS), sg.Text("", size=(1, 1)), sg.Button('Exit', size=(10, 1), key=KEY_EXIT)],
]
layout = [[sg.Text(title + sub_title, size=(48, 1), justification='right', font='Helvetica 20')],
[sg.Column(col_image0), sg.Column(col_image1)],
[sg.Column(col_color), sg.Column(col_radio)],
[sg.Text("", size=(15, 1))]
]
# ウィンドウオブジェクトの作成
window = sg.Window(title, layout, finalize=True, return_keyboard_events=True)
im0_canvas = window[KEY_ORGIMAGE]
tcv_img0 = im0_canvas.TKCanvas
tcv_img0.create_rectangle(0, 0, IMG_CANCAS_SIZE, IMG_CANCAS_SIZE, fill = '#cccccc')
org_image = Image.open(image_path) # PIL型で読み込み（オリジナル画像）
p_image = org_image.resize((IMG_CANCAS_SIZE, IMG_CANCAS_SIZE))
pht0_image = ImageTk.PhotoImage(image = p_image)
tcv_img0.create_image(IMG_CANCAS_SIZE / 2, IMG_CANCAS_SIZE / 2, image = pht0_image) # Canvasの中心に表示
# 処理動画
cap = cv2.VideoCapture(save_path)
cap_open = cap.isOpened()
new_make_f = False
# イベントのループ
while True:
event, values = window.read(timeout=30)
if new_make_f:
pic_dir = e4e_demo.still_image_generation(net, logger, vec_dir = VEC_DIR, latent = latent, direction = direction, min = min, max = max)
e4e_demo.make_movie(pic_dir, latent = latent, direction = direction, out_dir = result_image, disp_f = False)
cap = cv2.VideoCapture(save_path)
cap_open = cap.isOpened()
new_make_f = False
if event == KEY_EXIT or event == sg.WIN_CLOSED or event == 'Escape:27':
break
if event == KEY_MIN or event == KEY_MAX:
min = int(values[KEY_MIN])
max = int(values[KEY_MAX])
logger.debug(f'min = {min}, max = {max}')
if event == KEY_AGE or event == KEY_POSE or event == KEY_SMILE or event == KEY_AGEPOSE:
logger.debug(f'{KEY_AGE}, {KEY_POSE}, {KEY_SMILE}, {KEY_AGEPOSE}')
if values[KEY_AGE]:
direction = 'age'
min = -50
max = 50
elif values[KEY_POSE]:
direction = 'pose'
min = -50
max = 50
elif values[KEY_SMILE]:
direction = 'smile'
min = 0
max = 30
elif values[KEY_AGEPOSE]:
direction = 'age+pose'
min = -50
max = 50
window[KEY_MIN].update(float(min))
window[KEY_MAX].update(float(max))
window[KEY_TXTPRS].update(save_path)
if cap_open:
cap.release()
save_path = result_image + '/' + direction + '_' + img_name + '.mp4'
cap = cv2.VideoCapture(save_path)
cap_open = cap.isOpened()
if gpu_d and not cap_open:
new_make_f = True
if event == KEY_IMAGESEL:
fpath = get_image_file(VEC_PIC_DIR, msg = VEC_PIC_DIR + ' ')
if len(fpath) > 0:
# オリジナル画像変更
image_path = fpath
window[KEY_TXTORG].update(image_path)
org_image = Image.open(image_path) # PIL型で読み込み（オリジナル画像）
p_image = org_image.resize((IMG_CANCAS_SIZE, IMG_CANCAS_SIZE))
pht0_image = ImageTk.PhotoImage(image = p_image)
tcv_img0.create_image(IMG_CANCAS_SIZE / 2, IMG_CANCAS_SIZE / 2, image = pht0_image) # Canvasの中心に表示
base_dir_pair = os.path.split(image_path)
img_name, ext = os.path.splitext(base_dir_pair[1])
latent = img_name + '.pt'
if cap_open:
cap.release()
frame = np.zeros((IMG_CANCAS_SIZE, IMG_CANCAS_SIZE, 3))
frame += 204
img = cv2.imencode('.png', frame)[1].tobytes()
window[KEY_PRSIMAGE].update(img)
save_path = result_image + '/' + direction + '_' + img_name + '.mp4'
cap = cv2.VideoCapture(save_path)
cap_open = cap.isOpened()
if gpu_d and not cap_open:
new_make_f = True
window[KEY_TXTORG].update(image_path)
window[KEY_TXTPRS].update(save_path)
if event == KEY_IMAGEPTOS:
if gpu_d:
if cap_open:
cap.release()
cap_open = False
new_make_f = True
logger.debug(f'New image → latent = {latent}, direction = {direction}, min = {min}, max = {max}{NOCOLOR}')
if cap_open:
ret, frame = cap.read()
if frame is None:
#最初のフレームに戻る
cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
ret, frame = cap.read()
img = cv2.imencode('.png', frame)[1].tobytes()
window[KEY_PRSIMAGE].update(img)
# ウィンドウ終了処理
if cap_open:
cap.release()
window.close()
# 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))
if len(opt.source_image) == 0:
opt.source_image = get_image_file(IMAGE_DIR, msg = IMAGE_DIR + ' ')
if len(opt.source_image) == 0:
exit(0)
base_dir_pair = os.path.split(opt.source_image)
img_dir = base_dir_pair[0] # 画像ディレクトリ
display_info(opt, title)
if opt.cpu:
gpu_d = False
sub_title = '' if gpu_d else ' <view mode>'
if gpu_d:
import e4e_demo
start_time = datetime.datetime.now() # 時間計測開始
# モデルに学習済みパラメータをロード
net = e4e_demo.load_model(model_path = e4e_demo.MODEL_PATH)
else:
net = None
if gpu_d and not check_vector_image(img_dir, e4e_demo.VEC_DIR):
# 潜在変数と推定画像の生成
e4e_demo.make_latent(net, img_dir, out_path = opt.result_image, disp_f = True)
# 経過時間
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)
main_process(opt, net)
logger.info('\nFinished.\n')
