global refactoring and fixes,

removed support of extracted(aligned) PNG faces. Use old builds to convert from PNG to JPG.

fanseg model file in facelib/ is renamed

merger/InteractiveMergerSubprocessor.py

@@ -0,0 +1,571 @@
+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.fanseg_full_face_256_extract_func = client_dict['fanseg_full_face_256_extract_func']
+            self.skinseg_256_extract_func = client_dict['skinseg_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:
+                self.log_info (f'no faces found for {filepath.name}, copying without faces')
+
+                img_bgr = cv2_imread(filepath)
+                imagelib.normalize_channels(img_bgr, 3)
+                cv2_imwrite (pf.output_filepath, img_bgr)
+                h,w,c = img_bgr.shape
+
+                img_mask = np.zeros( (h,w,1), dtype=img_bgr.dtype)
+                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,
+                                                 fanseg_full_face_256_extract_func=self.fanseg_full_face_256_extract_func,
+                                                 skinseg_256_extract_func=self.skinseg_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, fanseg_full_face_256_extract_func, skinseg_256_extract_func, merger_config, frames, frames_root_path, output_path, output_mask_path, model_iter):
+        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.fanseg_full_face_256_extract_func = fanseg_full_face_256_extract_func
+        self.skinseg_256_extract_func = skinseg_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 = multiprocessing.cpu_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,
+                                      'fanseg_full_face_256_extract_func' : self.fanseg_full_face_256_extract_func,
+                                      'skinseg_256_extract_func' : self.skinseg_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),
+                    '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_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

merger/MergeMasked.py

@@ -8,7 +8,14 @@ from facelib import FaceType, LandmarksProcessor
 from core.interact import interact as io
 from core.cv2ex import *
 
-def MergeMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmarks):
+fanseg_input_size = 256
+skinseg_input_size = 256
+
+def MergeMaskedFace (predictor_func, predictor_input_shape, 
+                     face_enhancer_func,
+                     fanseg_full_face_256_extract_func, 
+                     skinseg_256_extract_func,
+                     cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmarks):
     img_size = img_bgr.shape[1], img_bgr.shape[0]
     img_face_mask_a = LandmarksProcessor.get_image_hull_mask (img_bgr.shape, img_face_landmarks)
 
@@ -53,7 +60,7 @@ def MergeMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, img
         predictor_masked = False
 
     if cfg.super_resolution_power != 0:
-        prd_face_bgr_enhanced = cfg.superres_func(prd_face_bgr)
+        prd_face_bgr_enhanced = face_enhancer_func(prd_face_bgr, is_tanh=True, preserve_size=False)
         mod = cfg.super_resolution_power / 100.0
         prd_face_bgr = cv2.resize(prd_face_bgr, (output_size,output_size))*(1.0-mod) + prd_face_bgr_enhanced*mod
         prd_face_bgr = np.clip(prd_face_bgr, 0, 1)
@@ -66,29 +73,29 @@ def MergeMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, img
 
     if cfg.mask_mode == 2: #dst
         prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC)
-    elif cfg.mask_mode >= 3 and cfg.mask_mode <= 8:
+    elif cfg.mask_mode >= 3 and cfg.mask_mode <= 7:
 
         if cfg.mask_mode == 3 or cfg.mask_mode == 5 or cfg.mask_mode == 6:
-            prd_face_fanseg_bgr = cv2.resize (prd_face_bgr, (cfg.fanseg_input_size,)*2 )
-            prd_face_fanseg_mask = cfg.fanseg_extract_func(FaceType.FULL, prd_face_fanseg_bgr)
+            prd_face_fanseg_bgr = cv2.resize (prd_face_bgr, (fanseg_input_size,)*2 )
+            prd_face_fanseg_mask = fanseg_full_face_256_extract_func(prd_face_fanseg_bgr)
             FAN_prd_face_mask_a_0 = cv2.resize ( prd_face_fanseg_mask, (output_size, output_size), cv2.INTER_CUBIC)
 
         if cfg.mask_mode >= 4 and cfg.mask_mode <= 7:
 
-            full_face_fanseg_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, cfg.fanseg_input_size, face_type=FaceType.FULL)
-            dst_face_fanseg_bgr = cv2.warpAffine(img_bgr, full_face_fanseg_mat, (cfg.fanseg_input_size,)*2, flags=cv2.INTER_CUBIC )
-            dst_face_fanseg_mask = cfg.fanseg_extract_func( FaceType.FULL, dst_face_fanseg_bgr )
+            full_face_fanseg_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, fanseg_input_size, face_type=FaceType.FULL)
+            dst_face_fanseg_bgr = cv2.warpAffine(img_bgr, full_face_fanseg_mat, (fanseg_input_size,)*2, flags=cv2.INTER_CUBIC )
+            dst_face_fanseg_mask = fanseg_full_face_256_extract_func(dst_face_fanseg_bgr )
 
             if cfg.face_type == FaceType.FULL:
                 FAN_dst_face_mask_a_0 = cv2.resize (dst_face_fanseg_mask, (output_size,output_size), cv2.INTER_CUBIC)
             else:
-                face_fanseg_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, cfg.fanseg_input_size, face_type=cfg.face_type)
+                face_fanseg_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, fanseg_input_size, face_type=cfg.face_type)
 
-                fanseg_rect_corner_pts = np.array ( [ [0,0], [cfg.fanseg_input_size-1,0], [0,cfg.fanseg_input_size-1] ], dtype=np.float32 )
+                fanseg_rect_corner_pts = np.array ( [ [0,0], [fanseg_input_size-1,0], [0,fanseg_input_size-1] ], dtype=np.float32 )
                 a = LandmarksProcessor.transform_points (fanseg_rect_corner_pts, face_fanseg_mat, invert=True )
                 b = LandmarksProcessor.transform_points (a, full_face_fanseg_mat )
                 m = cv2.getAffineTransform(b, fanseg_rect_corner_pts)
-                FAN_dst_face_mask_a_0 = cv2.warpAffine(dst_face_fanseg_mask, m, (cfg.fanseg_input_size,)*2, flags=cv2.INTER_CUBIC )
+                FAN_dst_face_mask_a_0 = cv2.warpAffine(dst_face_fanseg_mask, m, (fanseg_input_size,)*2, flags=cv2.INTER_CUBIC )
                 FAN_dst_face_mask_a_0 = cv2.resize (FAN_dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC)
 
         if cfg.mask_mode == 3:   #FAN-prd
@@ -101,7 +108,28 @@ def MergeMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, img
             prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0
         elif cfg.mask_mode == 7:
             prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_dst_face_mask_a_0
+            
+    elif cfg.mask_mode >= 8 and cfg.mask_mode <= 11:        
+        if cfg.mask_mode == 8 or cfg.mask_mode == 10 or cfg.mask_mode == 11:
+            prd_face_skinseg_bgr = cv2.resize (prd_face_bgr, (skinseg_input_size,)*2 )
+            prd_face_skinseg_mask = skinseg_256_extract_func(prd_face_skinseg_bgr)
+            X_prd_face_mask_a_0 = cv2.resize ( prd_face_skinseg_mask, (output_size, output_size), cv2.INTER_CUBIC)
 
+        if cfg.mask_mode >= 9 and cfg.mask_mode <= 11:            
+            whole_face_mat       = LandmarksProcessor.get_transform_mat (img_face_landmarks, skinseg_input_size, face_type=FaceType.WHOLE_FACE)
+            dst_face_skinseg_bgr  = cv2.warpAffine(img_bgr, whole_face_mat, (skinseg_input_size,)*2, flags=cv2.INTER_CUBIC )
+            dst_face_skinseg_mask = skinseg_256_extract_func(dst_face_skinseg_bgr)
+            X_dst_face_mask_a_0 = cv2.resize (dst_face_skinseg_mask, (output_size,output_size), cv2.INTER_CUBIC)
+
+        if cfg.mask_mode == 8:   #'SkinSeg-prd',
+            prd_face_mask_a_0 = X_prd_face_mask_a_0
+        elif cfg.mask_mode == 9: #'SkinSeg-dst',
+            prd_face_mask_a_0 = X_dst_face_mask_a_0
+        elif cfg.mask_mode == 10: #'SkinSeg-prd*SkinSeg-dst',
+            prd_face_mask_a_0 = X_prd_face_mask_a_0 * X_dst_face_mask_a_0
+        elif cfg.mask_mode == 11: #learned*SkinSeg-prd*SkinSeg-dst'
+            prd_face_mask_a_0 = prd_face_mask_a_0 * X_prd_face_mask_a_0 * X_dst_face_mask_a_0
+        
     prd_face_mask_a_0[ prd_face_mask_a_0 < (1.0/255.0) ] = 0.0 # get rid of noise
 
     # resize to mask_subres_size
@@ -280,7 +308,7 @@ def MergeMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, img
                     out_face_bgr = imagelib.LinearMotionBlur (out_face_bgr, k_size , frame_info.motion_deg)
 
             if cfg.blursharpen_amount != 0:
-                out_face_bgr = cfg.blursharpen_func ( out_face_bgr, cfg.sharpen_mode, 3, cfg.blursharpen_amount)
+                out_face_bgr = imagelib.blursharpen ( out_face_bgr, cfg.sharpen_mode, 3, cfg.blursharpen_amount)
 
 
             if cfg.image_denoise_power != 0:
@@ -315,14 +343,20 @@ def MergeMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, img
     return out_img, out_merging_mask_a
 
 
-def MergeMasked (predictor_func, predictor_input_shape, cfg, frame_info):
+def MergeMasked (predictor_func, 
+                 predictor_input_shape,
+                 face_enhancer_func,
+                 fanseg_full_face_256_extract_func, 
+                 skinseg_256_extract_func, 
+                 cfg, 
+                 frame_info):
     img_bgr_uint8 = cv2_imread(frame_info.filepath)
     img_bgr_uint8 = imagelib.normalize_channels (img_bgr_uint8, 3)
     img_bgr = img_bgr_uint8.astype(np.float32) / 255.0
 
     outs = []
     for face_num, img_landmarks in enumerate( frame_info.landmarks_list ):
-        out_img, out_img_merging_mask = MergeMaskedFace (predictor_func, predictor_input_shape, cfg, frame_info, img_bgr_uint8, img_bgr, img_landmarks)
+        out_img, out_img_merging_mask = MergeMaskedFace (predictor_func, predictor_input_shape, face_enhancer_func, fanseg_full_face_256_extract_func, skinseg_256_extract_func, cfg, frame_info, img_bgr_uint8, img_bgr, img_landmarks)
         outs += [ (out_img, out_img_merging_mask) ]
 
     #Combining multiple face outputs

merger/MergerConfig.py

@@ -21,11 +21,6 @@ class MergerConfig(object):
                        ):
         self.type = type
 
-        self.superres_func = None
-        self.blursharpen_func = None
-        self.fanseg_input_size = None
-        self.fanseg_extract_func = None
-
         self.sharpen_dict = {0:"None", 1:'box', 2:'gaussian'}
 
         #default changeable params
@@ -88,6 +83,19 @@ mode_str_dict = {}
 for key in mode_dict.keys():
     mode_str_dict[ mode_dict[key] ] = key
 
+whole_face_mask_mode_dict = {1:'learned',
+                             2:'dst',
+                             3:'FAN-prd',
+                             4:'FAN-dst',
+                             5:'FAN-prd*FAN-dst',
+                             6:'learned*FAN-prd*FAN-dst'                                  
+                             }
+"""
+8:'SkinSeg-prd',
+9:'SkinSeg-dst',
+10:'SkinSeg-prd*SkinSeg-dst',
+11:'learned*SkinSeg-prd*SkinSeg-dst'
+"""                                   
 full_face_mask_mode_dict = {1:'learned',
                                     2:'dst',
                                     3:'FAN-prd',
@@ -164,7 +172,9 @@ class MergerConfigMasked(MergerConfig):
             self.hist_match_threshold = np.clip ( self.hist_match_threshold+diff , 0, 255)
 
     def toggle_mask_mode(self):
-        if self.face_type == FaceType.FULL:
+        if self.face_type == FaceType.WHOLE_FACE:
+            a = list( whole_face_mask_mode_dict.keys() )
+        elif self.face_type == FaceType.FULL:
             a = list( full_face_mask_mode_dict.keys() )
         else:
             a = list( half_face_mask_mode_dict.keys() )
@@ -213,7 +223,14 @@ class MergerConfigMasked(MergerConfig):
             if self.mode == 'hist-match' or self.mode == 'seamless-hist-match':
                 self.hist_match_threshold = np.clip ( io.input_int("Hist match threshold", 255, add_info="0..255"), 0, 255)
 
-        if self.face_type == FaceType.FULL:
+        if self.face_type == FaceType.WHOLE_FACE:
+            s = """Choose mask mode: \n"""
+            for key in whole_face_mask_mode_dict.keys():
+                s += f"""({key}) {whole_face_mask_mode_dict[key]}\n"""
+            io.log_info(s)
+
+            self.mask_mode = io.input_int ("", 1, valid_list=whole_face_mask_mode_dict.keys() )
+        elif self.face_type == FaceType.FULL:
             s = """Choose mask mode: \n"""
             for key in full_face_mask_mode_dict.keys():
                 s += f"""({key}) {full_face_mask_mode_dict[key]}\n"""
@@ -282,7 +299,9 @@ class MergerConfigMasked(MergerConfig):
         if self.mode == 'hist-match' or self.mode == 'seamless-hist-match':
             r += f"""hist_match_threshold: {self.hist_match_threshold}\n"""
 
-        if self.face_type == FaceType.FULL:
+        if self.face_type == FaceType.WHOLE_FACE:
+            r += f"""mask_mode: { whole_face_mask_mode_dict[self.mask_mode] }\n"""
+        elif self.face_type == FaceType.FULL:
             r += f"""mask_mode: { full_face_mask_mode_dict[self.mask_mode] }\n"""
         else:
             r += f"""mask_mode: { half_face_mask_mode_dict[self.mask_mode] }\n"""

merger/MergerScreen/MergerScreen.py

@@ -0,0 +1,149 @@
+import math
+from pathlib import Path
+
+import numpy as np
+
+from core import imagelib
+from core.interact import interact as io
+from core.cv2ex import *
+from core import osex
+
+
+class ScreenAssets(object):
+    waiting_icon_image = cv2_imread ( str(Path(__file__).parent / 'gfx' / 'sand_clock_64.png') )
+
+    @staticmethod
+    def build_checkerboard_a( sh, size=5):
+        h,w = sh[0], sh[1]
+        tile = np.array([[0,1],[1,0]]).repeat(size, axis=0).repeat(size, axis=1)
+        grid = np.tile(tile,(int(math.ceil((h+0.0)/(2*size))),int(math.ceil((w+0.0)/(2*size)))))
+        return grid[:h,:w,None]
+
+class Screen(object):
+    def __init__(self, initial_scale_to_width=0, initial_scale_to_height=0, image=None, waiting_icon=False, **kwargs):
+        self.initial_scale_to_width = initial_scale_to_width
+        self.initial_scale_to_height = initial_scale_to_height
+        self.image = image
+        self.waiting_icon = waiting_icon
+
+        self.state = -1
+        self.scale = 1
+        self.force_update = True
+        self.is_first_appear = True
+        self.show_checker_board = False
+
+        self.last_screen_shape = (480,640,3)
+        self.checkerboard_image = None
+        self.set_image (image)
+        self.scrn_manager = None
+
+    def set_waiting_icon(self, b):
+        self.waiting_icon = b
+
+    def toggle_show_checker_board(self):
+        self.show_checker_board = not self.show_checker_board
+        self.force_update = True
+    
+    def get_image(self):
+        return self.image
+        
+    def set_image(self, img):
+        if not img is self.image:
+            self.force_update = True
+
+        self.image = img
+
+        if self.image is not None:
+            self.last_screen_shape = self.image.shape
+
+            if self.initial_scale_to_width != 0:
+                if self.last_screen_shape[1] > self.initial_scale_to_width:
+                    self.scale = self.initial_scale_to_width / self.last_screen_shape[1]
+                    self.force_update = True
+                self.initial_scale_to_width = 0
+
+            elif self.initial_scale_to_height != 0:
+                if self.last_screen_shape[0] > self.initial_scale_to_height:
+                    self.scale = self.initial_scale_to_height / self.last_screen_shape[0]
+                    self.force_update = True
+                self.initial_scale_to_height = 0
+
+
+    def diff_scale(self, diff):
+        self.scale = np.clip (self.scale + diff, 0.1, 4.0)
+        self.force_update = True
+
+    def show(self, force=False):
+        new_state = 0 | self.waiting_icon
+
+        if self.state != new_state or self.force_update or force:
+            self.state = new_state
+            self.force_update = False
+
+            if self.image is None:
+                screen = np.zeros ( self.last_screen_shape, dtype=np.uint8 )
+            else:
+                screen = self.image.copy()
+
+            if self.waiting_icon:
+                imagelib.overlay_alpha_image (screen, ScreenAssets.waiting_icon_image, (0,0) )
+
+            h,w,c = screen.shape
+            if self.scale != 1.0:
+                screen = cv2.resize ( screen, ( int(w*self.scale), int(h*self.scale) ) )
+
+            if c == 4:
+                if not self.show_checker_board:
+                    screen = screen[...,0:3]
+                else:
+                    if self.checkerboard_image is None or self.checkerboard_image.shape[0:2] != screen.shape[0:2]:
+                        self.checkerboard_image = ScreenAssets.build_checkerboard_a(screen.shape)
+
+                    screen = screen[...,0:3]*0.75 + 64*self.checkerboard_image*(1- (screen[...,3:4].astype(np.float32)/255.0) )
+                    screen = screen.astype(np.uint8)
+
+            io.show_image(self.scrn_manager.wnd_name, screen)
+
+            if self.is_first_appear:
+                self.is_first_appear = False
+                #center window
+                desktop_w, desktop_h = osex.get_screen_size()
+                h,w,c = screen.shape
+                cv2.moveWindow(self.scrn_manager.wnd_name, max(0,(desktop_w-w) // 2), max(0, (desktop_h-h) // 2) )
+
+            io.process_messages(0.0001)
+
+class ScreenManager(object):
+    def __init__(self, window_name="ScreenManager", screens=None, capture_keys=False ):
+        self.screens = screens or []
+        self.current_screen_id = 0
+
+        if self.screens is not None:
+            for screen in self.screens:
+                screen.scrn_manager = self
+
+        self.wnd_name = window_name
+        io.named_window(self.wnd_name)
+
+
+        if capture_keys:
+            io.capture_keys(self.wnd_name)
+
+    def finalize(self):
+        io.destroy_all_windows()
+
+    def get_key_events(self):
+        return io.get_key_events(self.wnd_name)
+
+    def switch_screens(self):
+        self.current_screen_id = (self.current_screen_id + 1) % len(self.screens)
+        self.screens[self.current_screen_id].show(force=True)
+
+    def show_current(self):
+        self.screens[self.current_screen_id].show()
+
+    def get_current(self):
+        return self.screens[self.current_screen_id]
+
+    def set_current(self, screen):
+        self.current_screen_id = self.screens.index(screen)

merger/MergerScreen/__init__.py

@@ -0,0 +1 @@
+from .MergerScreen import Screen, ScreenManager

merger/__init__.py

@@ -2,3 +2,4 @@ from .FrameInfo import FrameInfo
 from .MergerConfig import MergerConfig, MergerConfigMasked, MergerConfigFaceAvatar
 from .MergeMasked import MergeMasked
 from .MergeAvatar import MergeFaceAvatar
+from .InteractiveMergerSubprocessor import InteractiveMergerSubprocessor