import multiprocessing import os import pickle import sys import traceback from pathlib import Path import numpy as np from core import imagelib, pathex from core.cv2ex import * from core.interact import interact as io from core.joblib import Subprocessor from merger import MergeFaceAvatar, MergeMasked, MergerConfig from .MergerScreen import Screen, ScreenManager MERGER_DEBUG = False class InteractiveMergerSubprocessor(Subprocessor): class Frame(object): def __init__(self, prev_temporal_frame_infos=None, frame_info=None, next_temporal_frame_infos=None): self.prev_temporal_frame_infos = prev_temporal_frame_infos self.frame_info = frame_info self.next_temporal_frame_infos = next_temporal_frame_infos self.output_filepath = None self.output_mask_filepath = None self.idx = None self.cfg = None self.is_done = False self.is_processing = False self.is_shown = False self.image = None class ProcessingFrame(object): def __init__(self, idx=None, cfg=None, prev_temporal_frame_infos=None, frame_info=None, next_temporal_frame_infos=None, output_filepath=None, output_mask_filepath=None, need_return_image = False): self.idx = idx self.cfg = cfg self.prev_temporal_frame_infos = prev_temporal_frame_infos self.frame_info = frame_info self.next_temporal_frame_infos = next_temporal_frame_infos self.output_filepath = output_filepath self.output_mask_filepath = output_mask_filepath self.need_return_image = need_return_image if self.need_return_image: self.image = None class Cli(Subprocessor.Cli): #override def on_initialize(self, client_dict): self.log_info ('Running on %s.' % (client_dict['device_name']) ) self.device_idx = client_dict['device_idx'] self.device_name = client_dict['device_name'] self.predictor_func = client_dict['predictor_func'] self.predictor_input_shape = client_dict['predictor_input_shape'] self.face_enhancer_func = client_dict['face_enhancer_func'] self.xseg_256_extract_func = client_dict['xseg_256_extract_func'] #transfer and set stdin in order to work code.interact in debug subprocess stdin_fd = client_dict['stdin_fd'] if stdin_fd is not None: sys.stdin = os.fdopen(stdin_fd) return None #override def process_data(self, pf): #pf=ProcessingFrame cfg = pf.cfg.copy() frame_info = pf.frame_info filepath = frame_info.filepath if len(frame_info.landmarks_list) == 0: if cfg.mode == 'raw-predict': h,w,c = self.predictor_input_shape img_bgr = np.zeros( (h,w,3), dtype=np.uint8) img_mask = np.zeros( (h,w,1), dtype=np.uint8) else: self.log_info (f'no faces found for {filepath.name}, copying without faces') img_bgr = cv2_imread(filepath) imagelib.normalize_channels(img_bgr, 3) h,w,c = img_bgr.shape img_mask = np.zeros( (h,w,1), dtype=img_bgr.dtype) cv2_imwrite (pf.output_filepath, img_bgr) cv2_imwrite (pf.output_mask_filepath, img_mask) if pf.need_return_image: pf.image = np.concatenate ([img_bgr, img_mask], axis=-1) else: if cfg.type == MergerConfig.TYPE_MASKED: try: final_img = MergeMasked (self.predictor_func, self.predictor_input_shape, face_enhancer_func=self.face_enhancer_func, xseg_256_extract_func=self.xseg_256_extract_func, cfg=cfg, frame_info=frame_info) except Exception as e: e_str = traceback.format_exc() if 'MemoryError' in e_str: raise Subprocessor.SilenceException else: raise Exception( f'Error while merging file [{filepath}]: {e_str}' ) elif cfg.type == MergerConfig.TYPE_FACE_AVATAR: final_img = MergeFaceAvatar (self.predictor_func, self.predictor_input_shape, cfg, pf.prev_temporal_frame_infos, pf.frame_info, pf.next_temporal_frame_infos ) cv2_imwrite (pf.output_filepath, final_img[...,0:3] ) cv2_imwrite (pf.output_mask_filepath, final_img[...,3:4] ) if pf.need_return_image: pf.image = final_img return pf #overridable def get_data_name (self, pf): #return string identificator of your data return pf.frame_info.filepath #override def __init__(self, is_interactive, merger_session_filepath, predictor_func, predictor_input_shape, face_enhancer_func, xseg_256_extract_func, merger_config, frames, frames_root_path, output_path, output_mask_path, model_iter, subprocess_count=4): if len (frames) == 0: raise ValueError ("len (frames) == 0") super().__init__('Merger', InteractiveMergerSubprocessor.Cli, io_loop_sleep_time=0.001) self.is_interactive = is_interactive self.merger_session_filepath = Path(merger_session_filepath) self.merger_config = merger_config self.predictor_func = predictor_func self.predictor_input_shape = predictor_input_shape self.face_enhancer_func = face_enhancer_func self.xseg_256_extract_func = xseg_256_extract_func self.frames_root_path = frames_root_path self.output_path = output_path self.output_mask_path = output_mask_path self.model_iter = model_iter self.prefetch_frame_count = self.process_count = subprocess_count session_data = None if self.is_interactive and self.merger_session_filepath.exists(): io.input_skip_pending() if io.input_bool ("Use saved session?", True): try: with open( str(self.merger_session_filepath), "rb") as f: session_data = pickle.loads(f.read()) except Exception as e: pass rewind_to_frame_idx = None self.frames = frames self.frames_idxs = [ *range(len(self.frames)) ] self.frames_done_idxs = [] if self.is_interactive and session_data is not None: # Loaded session data, check it 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)) # frames count must match if frames_equal: for i in range(len(frames)): frame = frames[i] s_frame = s_frames[i] # frames filenames must match if frame.frame_info.filepath.name != s_frame.frame_info.filepath.name: frames_equal = False if not frames_equal: break if frames_equal: io.log_info ('Using saved session from ' + '/'.join (self.merger_session_filepath.parts[-2:]) ) for frame in s_frames: if frame.cfg is not None: # recreate MergerConfig class using constructor with get_config() as dict params # so if any new param will be added, old merger session will work properly frame.cfg = frame.cfg.__class__( **frame.cfg.get_config() ) 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 was more trained, recompute all frames rewind_to_frame_idx = -1 for frame in self.frames: frame.is_done = False elif len(self.frames_idxs) == 0: # all frames are done? rewind_to_frame_idx = -1 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 pathex.get_image_paths(self.output_path): #remove all images in output_path Path(filename).unlink() for filename in pathex.get_image_paths(self.output_mask_path): #remove all images in output_mask_path Path(filename).unlink() frames[0].cfg = self.merger_config.copy() for i in range( len(self.frames) ): frame = self.frames[i] frame.idx = i frame.output_filepath = self.output_path / ( frame.frame_info.filepath.stem + '.png' ) frame.output_mask_filepath = self.output_mask_path / ( frame.frame_info.filepath.stem + '.png' ) if not frame.output_filepath.exists() or \ not frame.output_mask_filepath.exists(): # if some frame does not exist, recompute and rewind frame.is_done = False frame.is_shown = False if rewind_to_frame_idx is None: rewind_to_frame_idx = i-1 else: rewind_to_frame_idx = min(rewind_to_frame_idx, i-1) if rewind_to_frame_idx is not None: while len(self.frames_done_idxs) > 0: if self.frames_done_idxs[-1] > rewind_to_frame_idx: prev_frame = self.frames[self.frames_done_idxs.pop()] self.frames_idxs.insert(0, prev_frame.idx) else: break #override def process_info_generator(self): r = [0] if MERGER_DEBUG else range(self.process_count) for i in r: yield 'CPU%d' % (i), {}, {'device_idx': i, 'device_name': 'CPU%d' % (i), 'predictor_func': self.predictor_func, 'predictor_input_shape' : self.predictor_input_shape, 'face_enhancer_func': self.face_enhancer_func, 'xseg_256_extract_func' : self.xseg_256_extract_func, 'stdin_fd': sys.stdin.fileno() if MERGER_DEBUG else None } #overridable optional def on_clients_initialized(self): io.progress_bar ("Merging", len(self.frames_idxs)+len(self.frames_done_idxs), initial=len(self.frames_done_idxs) ) self.process_remain_frames = not self.is_interactive self.is_interactive_quitting = not self.is_interactive if self.is_interactive: help_images = { MergerConfig.TYPE_MASKED : cv2_imread ( str(Path(__file__).parent / 'gfx' / 'help_merger_masked.jpg') ), MergerConfig.TYPE_FACE_AVATAR : cv2_imread ( str(Path(__file__).parent / 'gfx' / 'help_merger_face_avatar.jpg') ), } self.main_screen = Screen(initial_scale_to_width=1368, image=None, waiting_icon=True) self.help_screen = Screen(initial_scale_to_height=768, image=help_images[self.merger_config.type], waiting_icon=False) self.screen_manager = ScreenManager( "Merger", [self.main_screen, self.help_screen], capture_keys=True ) self.screen_manager.set_current (self.help_screen) self.screen_manager.show_current() self.masked_keys_funcs = { '`' : lambda cfg,shift_pressed: cfg.set_mode(0), '1' : lambda cfg,shift_pressed: cfg.set_mode(1), '2' : lambda cfg,shift_pressed: cfg.set_mode(2), '3' : lambda cfg,shift_pressed: cfg.set_mode(3), '4' : lambda cfg,shift_pressed: cfg.set_mode(4), '5' : lambda cfg,shift_pressed: cfg.set_mode(5), '6' : lambda cfg,shift_pressed: cfg.set_mode(6), 'q' : lambda cfg,shift_pressed: cfg.add_hist_match_threshold(1 if not shift_pressed else 5), 'a' : lambda cfg,shift_pressed: cfg.add_hist_match_threshold(-1 if not shift_pressed else -5), 'w' : lambda cfg,shift_pressed: cfg.add_erode_mask_modifier(1 if not shift_pressed else 5), 's' : lambda cfg,shift_pressed: cfg.add_erode_mask_modifier(-1 if not shift_pressed else -5), 'e' : lambda cfg,shift_pressed: cfg.add_blur_mask_modifier(1 if not shift_pressed else 5), 'd' : lambda cfg,shift_pressed: cfg.add_blur_mask_modifier(-1 if not shift_pressed else -5), 'r' : lambda cfg,shift_pressed: cfg.add_motion_blur_power(1 if not shift_pressed else 5), 'f' : lambda cfg,shift_pressed: cfg.add_motion_blur_power(-1 if not shift_pressed else -5), 't' : lambda cfg,shift_pressed: cfg.add_super_resolution_power(1 if not shift_pressed else 5), 'g' : lambda cfg,shift_pressed: cfg.add_super_resolution_power(-1 if not shift_pressed else -5), 'y' : lambda cfg,shift_pressed: cfg.add_blursharpen_amount(1 if not shift_pressed else 5), 'h' : lambda cfg,shift_pressed: cfg.add_blursharpen_amount(-1 if not shift_pressed else -5), 'u' : lambda cfg,shift_pressed: cfg.add_output_face_scale(1 if not shift_pressed else 5), 'j' : lambda cfg,shift_pressed: cfg.add_output_face_scale(-1 if not shift_pressed else -5), 'i' : lambda cfg,shift_pressed: cfg.add_image_denoise_power(1 if not shift_pressed else 5), 'k' : lambda cfg,shift_pressed: cfg.add_image_denoise_power(-1 if not shift_pressed else -5), 'o' : lambda cfg,shift_pressed: cfg.add_bicubic_degrade_power(1 if not shift_pressed else 5), 'l' : lambda cfg,shift_pressed: cfg.add_bicubic_degrade_power(-1 if not shift_pressed else -5), 'p' : lambda cfg,shift_pressed: cfg.add_color_degrade_power(1 if not shift_pressed else 5), ';' : lambda cfg,shift_pressed: cfg.add_color_degrade_power(-1), ':' : lambda cfg,shift_pressed: cfg.add_color_degrade_power(-5), 'z' : lambda cfg,shift_pressed: cfg.toggle_masked_hist_match(), 'x' : lambda cfg,shift_pressed: cfg.toggle_mask_mode(), 'c' : lambda cfg,shift_pressed: cfg.toggle_color_transfer_mode(), 'n' : lambda cfg,shift_pressed: cfg.toggle_sharpen_mode(), } self.masked_keys = list(self.masked_keys_funcs.keys()) #overridable optional def on_clients_finalized(self): io.progress_bar_close() if self.is_interactive: self.screen_manager.finalize() for frame in self.frames: frame.output_filepath = None frame.output_mask_filepath = 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, } self.merger_session_filepath.write_bytes( pickle.dumps(session_data) ) io.log_info ("Session is saved to " + '/'.join (self.merger_session_filepath.parts[-2:]) ) #override def on_tick(self): io.process_messages() go_prev_frame = False go_first_frame = False go_prev_frame_overriding_cfg = False go_first_frame_overriding_cfg = False go_next_frame = self.process_remain_frames go_next_frame_overriding_cfg = False go_last_frame_overriding_cfg = False cur_frame = None if len(self.frames_idxs) != 0: cur_frame = self.frames[self.frames_idxs[0]] if self.is_interactive: screen_image = None if self.process_remain_frames else \ self.main_screen.get_image() self.main_screen.set_waiting_icon( self.process_remain_frames or \ self.is_interactive_quitting ) if cur_frame is not None and not self.is_interactive_quitting: if not self.process_remain_frames: if cur_frame.is_done: if not cur_frame.is_shown: if cur_frame.image is None: image = cv2_imread (cur_frame.output_filepath, verbose=False) image_mask = cv2_imread (cur_frame.output_mask_filepath, verbose=False) if image is None or image_mask is None: # unable to read? recompute then cur_frame.is_done = False else: image = imagelib.normalize_channels(image, 3) image_mask = imagelib.normalize_channels(image_mask, 1) cur_frame.image = np.concatenate([image, image_mask], -1) if cur_frame.is_done: io.log_info (cur_frame.cfg.to_string( cur_frame.frame_info.filepath.name) ) cur_frame.is_shown = True screen_image = cur_frame.image else: self.main_screen.set_waiting_icon(True) self.main_screen.set_image(screen_image) self.screen_manager.show_current() key_events = self.screen_manager.get_key_events() key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[-1] if len(key_events) > 0 else (0,0,False,False,False) if key == 9: #tab self.screen_manager.switch_screens() else: if key == 27: #esc self.is_interactive_quitting = True elif self.screen_manager.get_current() is self.main_screen: if self.merger_config.type == MergerConfig.TYPE_MASKED and chr_key in self.masked_keys: self.process_remain_frames = False if cur_frame is not None: cfg = cur_frame.cfg prev_cfg = cfg.copy() if cfg.type == MergerConfig.TYPE_MASKED: self.masked_keys_funcs[chr_key](cfg, shift_pressed) if prev_cfg != cfg: io.log_info ( cfg.to_string(cur_frame.frame_info.filepath.name) ) cur_frame.is_done = False cur_frame.is_shown = False else: if chr_key == ',' or chr_key == 'm': self.process_remain_frames = False go_prev_frame = True if chr_key == ',': if shift_pressed: go_first_frame = True elif chr_key == 'm': if not shift_pressed: go_prev_frame_overriding_cfg = True else: go_first_frame_overriding_cfg = True elif chr_key == '.' or chr_key == '/': self.process_remain_frames = False go_next_frame = True if chr_key == '.': if shift_pressed: self.process_remain_frames = not self.process_remain_frames elif chr_key == '/': if not shift_pressed: go_next_frame_overriding_cfg = True else: go_last_frame_overriding_cfg = True elif chr_key == '-': self.screen_manager.get_current().diff_scale(-0.1) elif chr_key == '=': self.screen_manager.get_current().diff_scale(0.1) elif chr_key == 'v': self.screen_manager.get_current().toggle_show_checker_board() if go_prev_frame: if cur_frame is None or cur_frame.is_done: if cur_frame is not None: cur_frame.image = None while True: 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 or go_first_frame_overriding_cfg): if prev_frame.cfg != cur_frame.cfg: prev_frame.cfg = cur_frame.cfg.copy() prev_frame.is_done = False cur_frame = prev_frame if go_first_frame_overriding_cfg or go_first_frame: if len(self.frames_done_idxs) > 0: continue break elif go_next_frame: if cur_frame is not None and cur_frame.is_done: cur_frame.image = None cur_frame.is_shown = True self.frames_done_idxs.append(cur_frame.idx) self.frames_idxs.pop(0) io.progress_bar_inc(1) f = self.frames if len(self.frames_idxs) != 0: next_frame = f[ self.frames_idxs[0] ] next_frame.is_shown = False if go_next_frame_overriding_cfg or go_last_frame_overriding_cfg: if go_next_frame_overriding_cfg: to_frames = next_frame.idx+1 else: to_frames = len(f) for i in range( next_frame.idx, to_frames ): f[i].cfg = None for i in range( min(len(self.frames_idxs), self.prefetch_frame_count) ): frame = f[ self.frames_idxs[i] ] if frame.cfg is None: if i == 0: frame.cfg = cur_frame.cfg.copy() else: frame.cfg = f[ self.frames_idxs[i-1] ].cfg.copy() frame.is_done = False #initiate solve again frame.is_shown = False if len(self.frames_idxs) == 0: self.process_remain_frames = False return (self.is_interactive and self.is_interactive_quitting) or \ (not self.is_interactive and self.process_remain_frames == False) #override def on_data_return (self, host_dict, pf): frame = self.frames[pf.idx] frame.is_done = False frame.is_processing = False #override def on_result (self, host_dict, pf_sent, pf_result): frame = self.frames[pf_result.idx] frame.is_processing = False if frame.cfg == pf_result.cfg: frame.is_done = True frame.image = pf_result.image #override def get_data(self, host_dict): if self.is_interactive and self.is_interactive_quitting: return None for i in range ( min(len(self.frames_idxs), self.prefetch_frame_count) ): frame = self.frames[ self.frames_idxs[i] ] if not frame.is_done and not frame.is_processing and frame.cfg is not None: frame.is_processing = True return InteractiveMergerSubprocessor.ProcessingFrame(idx=frame.idx, cfg=frame.cfg.copy(), prev_temporal_frame_infos=frame.prev_temporal_frame_infos, frame_info=frame.frame_info, next_temporal_frame_infos=frame.next_temporal_frame_infos, output_filepath=frame.output_filepath, output_mask_filepath=frame.output_mask_filepath, need_return_image=True ) return None #override def get_result(self): return 0