Skip to content
Permalink
Browse files

interactive converter now saves the session.

if input frames are changed (amount or filenames)
then interactive converter automatically starts a new session.
if model is more trained then all frames will be recomputed again with their saved configs.
  • Loading branch information...
iperov committed Aug 30, 2019
1 parent bac9d5a commit bc0b25cdda51942c3ad40770fc6200477ae32e61
Showing with 102 additions and 24 deletions.
  1. +4 −4 interact/interact.py
  2. +98 −20 mainscripts/Converter.py
@@ -137,9 +137,9 @@ def capture_keys(self, wnd_name):
else: print("capture_keys: already set for window ", wnd_name)
else: print("capture_keys: named_window ", wnd_name, " not found.")

def progress_bar(self, desc, total, leave=True):
def progress_bar(self, desc, total, leave=True, initial=0):
if self.pg_bar is None:
self.pg_bar = tqdm( total=total, desc=desc, leave=leave, ascii=True )
self.pg_bar = tqdm( total=total, desc=desc, leave=leave, ascii=True, initial=initial )
else: print("progress_bar: already set.")

def progress_bar_inc(self, c):
@@ -154,8 +154,8 @@ def progress_bar_close(self):
self.pg_bar = None
else: print("progress_bar not set.")

def progress_bar_generator(self, data, desc, leave=True):
self.pg_bar = tqdm( data, desc=desc, leave=leave, ascii=True )
def progress_bar_generator(self, data, desc, leave=True, initial=0):
self.pg_bar = tqdm( data, desc=desc, leave=leave, ascii=True, initial=initial )
for x in self.pg_bar:
yield x
self.pg_bar.close()
@@ -2,6 +2,7 @@
import multiprocessing
import operator
import os
import pickle
import shutil
import sys
import time
@@ -13,7 +14,8 @@
import numpy.linalg as npla

import imagelib
from converters import FrameInfo, ConverterConfig, ConvertMasked, ConvertFaceAvatar
from converters import (ConverterConfig, ConvertFaceAvatar, ConvertMasked,
FrameInfo)
from facelib import FaceType, FANSegmentator, LandmarksProcessor
from interact import interact as io
from joblib import SubprocessFunctionCaller, Subprocessor
@@ -87,20 +89,20 @@ def on_initialize(self, client_dict):
def sharpen_func (img, sharpen_mode=0, kernel_size=3, amount=150):
if kernel_size % 2 == 0:
kernel_size += 1

if sharpen_mode == 1: #box
kernel = np.zeros( (kernel_size, kernel_size), dtype=np.float32)
kernel[ kernel_size//2, kernel_size//2] = 1.0
kernel[ kernel_size//2, kernel_size//2] = 1.0
box_filter = np.ones( (kernel_size, kernel_size), dtype=np.float32) / (kernel_size**2)
kernel = kernel + (kernel - box_filter) * amount
return cv2.filter2D(img, -1, kernel)
elif sharpen_mode == 2: #gaussian
blur = cv2.GaussianBlur(img, (kernel_size, kernel_size) , 0)
elif sharpen_mode == 2: #gaussian
blur = cv2.GaussianBlur(img, (kernel_size, kernel_size) , 0)
img = cv2.addWeighted(img, 1.0 + (0.5 * amount), blur, -(0.5 * amount), 0)
return img
return img
self.sharpen_func = sharpen_func

self.fanseg_by_face_type = {}
self.fanseg_input_size = 256

@@ -122,7 +124,7 @@ def process_data(self, pf): #pf=ProcessingFrame
cfg.predictor_func = self.predictor_func
cfg.sharpen_func = self.sharpen_func
cfg.superres_func = self.superres_func

frame_info = pf.frame_info

filename = frame_info.filename
@@ -177,7 +179,7 @@ def get_data_name (self, pf):
return pf.frame_info.filename

#override
def __init__(self, is_interactive, converter_config, frames, output_path):
def __init__(self, is_interactive, converter_config, frames, output_path, model_iter):
if len (frames) == 0:
raise ValueError ("len (frames) == 0")

@@ -203,20 +205,84 @@ def superres_func(mode, *args, **kwargs):

self.dcscn_host, self.superres_func = SubprocessFunctionCaller.make_pair(superres_func)

self.frames = frames
self.output_path = output_path
self.model_iter = model_iter

self.prefetch_frame_count = self.process_count = min(6,multiprocessing.cpu_count())

session_data = None
session_dat_path = self.output_path / 'session.dat'
if session_dat_path.exists():

try:
with open( str(session_dat_path), "rb") as f:
session_data = pickle.loads(f.read())
except Exception as e:
pass

self.frames = frames
self.frames_idxs = [ *range(len(self.frames)) ]
self.frames_done_idxs = []

digits = [ str(i) for i in range(10)]
if self.is_interactive and session_data is not None:
s_frames = session_data.get('frames', None)
s_frames_idxs = session_data.get('frames_idxs', None)
s_frames_done_idxs = session_data.get('frames_done_idxs', None)
s_model_iter = session_data.get('model_iter', None)

frames_equal = (s_frames is not None) and \
(s_frames_idxs is not None) and \
(s_frames_done_idxs is not None) and \
(s_model_iter is not None) and \
(len(frames) == len(s_frames))

if frames_equal:
for i in range(len(frames)):
frame = frames[i]
s_frame = s_frames[i]
if frame.frame_info.filename != s_frame.frame_info.filename:
frames_equal = False
if not frames_equal:
break

if frames_equal:
io.log_info ("Using saved session.")
self.frames = s_frames
self.frames_idxs = s_frames_idxs
self.frames_done_idxs = s_frames_done_idxs

if self.model_iter != s_model_iter:
#model is more trained, recompute all frames
for frame in self.frames:
frame.is_done = False

if self.model_iter != s_model_iter or \
len(self.frames_idxs) == 0:
#rewind to begin if model is more trained or all frames are done

while len(self.frames_done_idxs) > 0:
prev_frame = self.frames[self.frames_done_idxs.pop()]
self.frames_idxs.insert(0, prev_frame.idx)

if len(self.frames_idxs) != 0:
cur_frame = self.frames[self.frames_idxs[0]]
cur_frame.is_shown = False

if not frames_equal:
session_data = None

if session_data is None:
for filename in Path_utils.get_image_paths(self.output_path): #remove all images in output_path
Path(filename).unlink()

frames[0].cfg = self.converter_config.copy()

for i in range( len(self.frames) ):
frame = self.frames[i]
frame.idx = i
frame.output_filename = self.output_path / ( Path(frame.frame_info.filename).stem + '.png' )

frames[0].cfg = self.converter_config.copy()


#override
def process_info_generator(self):
@@ -232,7 +298,7 @@ def process_info_generator(self):

#overridable optional
def on_clients_initialized(self):
io.progress_bar ("Converting", len (self.frames_idxs) )
io.progress_bar ("Converting", len (self.frames_idxs), initial=len(self.frames_done_idxs) )

self.process_remain_frames = not self.is_interactive
self.is_interactive_quitting = not self.is_interactive
@@ -251,11 +317,25 @@ def on_clients_initialized(self):

#overridable optional
def on_clients_finalized(self):
io.progress_bar_close()

if self.is_interactive:
self.screen_manager.finalize()

io.progress_bar_close()
for frame in self.frames:
frame.output_filename = None
frame.image = None

session_data = {
'frames': self.frames,
'frames_idxs': self.frames_idxs,
'frames_done_idxs': self.frames_done_idxs,
'model_iter' : self.model_iter,
}
save_path = self.output_path / 'session.dat'
save_path.write_bytes( pickle.dumps(session_data) )

io.log_info ("Session is saved to " + '/'.join (save_path.parts[-2:]) )

cfg_change_keys = ['`','1', '2', '3', '4', '5', '6', '7', '8', '9',
'q', 'a', 'w', 's', 'e', 'd', 'r', 'f', 't', 'g','y','h','u','j',
@@ -363,7 +443,7 @@ def on_tick(self):
cfg.toggle_export_mask_alpha()
elif chr_key == 'n':
cfg.toggle_sharpen_mode()

else:
if chr_key == 'y':
cfg.add_sharpen_amount(1 if not shift_pressed else 5)
@@ -401,13 +481,13 @@ def on_tick(self):
if cur_frame is None or cur_frame.is_done:
if cur_frame is not None:
cur_frame.image = None

if len(self.frames_done_idxs) > 0:
prev_frame = self.frames[self.frames_done_idxs.pop()]
self.frames_idxs.insert(0, prev_frame.idx)
prev_frame.is_shown = False
io.progress_bar_inc(-1)

if cur_frame is not None and go_prev_frame_overriding_cfg:
if prev_frame.cfg != cur_frame.cfg:
prev_frame.cfg = cur_frame.cfg.copy()
@@ -501,10 +581,7 @@ def main (args, device_args):
io.log_err('Input directory not found. Please ensure it exists.')
return

if output_path.exists():
for filename in Path_utils.get_image_paths(output_path):
Path(filename).unlink()
else:
if not output_path.exists():
output_path.mkdir(parents=True, exist_ok=True)

if not model_path.exists():
@@ -643,6 +720,7 @@ def main (args, device_args):
converter_config = cfg,
frames = frames,
output_path = output_path,
model_iter = model.get_iter()
).run()

model.finalize()

0 comments on commit bc0b25c

Please sign in to comment.
You can’t perform that action at this time.