From bc0b25cdda51942c3ad40770fc6200477ae32e61 Mon Sep 17 00:00:00 2001 From: Colombo Date: Fri, 30 Aug 2019 16:19:41 +0400 Subject: [PATCH] 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. --- interact/interact.py | 8 +-- mainscripts/Converter.py | 118 ++++++++++++++++++++++++++++++++------- 2 files changed, 102 insertions(+), 24 deletions(-) diff --git a/interact/interact.py b/interact/interact.py index b1b4e23..25d0549 100644 --- a/interact/interact.py +++ b/interact/interact.py @@ -137,9 +137,9 @@ class InteractBase(object): 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 @@ class InteractBase(object): 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() diff --git a/mainscripts/Converter.py b/mainscripts/Converter.py index d075613..d437a6f 100644 --- a/mainscripts/Converter.py +++ b/mainscripts/Converter.py @@ -2,6 +2,7 @@ import multiprocessing import operator import os +import pickle import shutil import sys import time @@ -13,7 +14,8 @@ import numpy as np 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 @@ class ConvertSubprocessor(Subprocessor): 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 @@ class ConvertSubprocessor(Subprocessor): 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 @@ class ConvertSubprocessor(Subprocessor): 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 @@ class ConvertSubprocessor(Subprocessor): 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 @@ class ConvertSubprocessor(Subprocessor): #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 @@ class ConvertSubprocessor(Subprocessor): #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 @@ class ConvertSubprocessor(Subprocessor): 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 @@ class ConvertSubprocessor(Subprocessor): 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()