Merge pull request #5 from MachineEditor/merger_updates

Merger updates

core/imagelib/__init__.py

@@ -18,7 +18,7 @@ from .common import random_crop, normalize_channels, cut_odd_image, overlay_alph
 
 from .SegIEPolys import *
 
-from .blursharpen import LinearMotionBlur, blursharpen
+from .blursharpen import LinearMotionBlur, blursharpen, gaussian_sharpen, unsharpen_mask
 
 from .filters import apply_random_rgb_levels, \
                      apply_random_overlay_triangle, \

core/imagelib/blursharpen.py

@@ -22,6 +22,8 @@ def blursharpen (img, sharpen_mode=0, kernel_size=3, amount=100):
             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
+        elif sharpen_mode == 3: #unsharpen_mask
+            img = unsharpen_mask(img, amount=amount)
     elif amount < 0:
         n = -amount
         while n > 0:
@@ -35,4 +37,18 @@ def blursharpen (img, sharpen_mode=0, kernel_size=3, amount=100):
             n = max(n-10,0)
 
         return img
+    return img
+    
+def gaussian_sharpen(img, amount=100, sigma=1.0):
+    img =  cv2.addWeighted(img, 1.0 + (0.05 * amount), cv2.GaussianBlur(img, (0, 0), sigma), -(0.05 * amount), 0)
+    return img
+    
+def unsharpen_mask(img, amount=100, sigma=0.0, threshold = (5.0 / 255.0)):
+    radius = max(1, round(img.shape[0] * (amount / 100)))
+    kernel_size = int((radius * 2) + 1)
+    kernel_size = (kernel_size, kernel_size)
+    blur = cv2.GaussianBlur(img, kernel_size, sigma)
+    low_contrast_mask = (abs(img - blur) < threshold).astype("float32")
+    sharpened = (img * (1.0 + (0.05 * amount))) + (blur * -(0.05 * amount))
+    img = (img * (1.0 - low_contrast_mask)) + (sharpened * low_contrast_mask)
     return img

mainscripts/Merger.py

@@ -146,17 +146,20 @@ def main (model_class_name=None,
                 io.log_info ("Use 'recover original filename' to determine the exact duplicates.")
                 io.log_info ("")
 
+
+
             # build frames maunally
             frames = []
             for p in input_path_image_paths:
-                path = Path(p) 
-                data = alignments.get(path.stem, None)
-                if data == None:      
-                    frame = InteractiveMergerSubprocessor.Frame(FrameInfo(frame_info=frame_info))
+                cur_path = Path(p)
+                data = alignments.get(cur_path.stem, None)
+                if data == None:
+                    frame_info=FrameInfo(filepath=cur_path)
+                    frame = InteractiveMergerSubprocessor.Frame(frame_info=frame_info)
                 else:
                     landmarks_list = [d[0] for d in data]
                     dfl_images_list = [d[1] for d in data]
-                    frame_info=FrameInfo(filepath=path, landmarks_list=landmarks_list, dfl_images_list=dfl_images_list)
+                    frame_info=FrameInfo(filepath=cur_path, landmarks_list=landmarks_list, dfl_images_list=dfl_images_list)
                     frame = InteractiveMergerSubprocessor.Frame(frame_info=frame_info)
 
                 frames.append(frame)

merger/InteractiveMergerSubprocessor.py

@@ -330,7 +330,13 @@ class InteractiveMergerSubprocessor(Subprocessor):
                     '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(),
+                    'n' : lambda cfg,shift_pressed: cfg.toggle_sharpen_mode_multi(shift_pressed),
+                    '9' : lambda cfg,shift_pressed: cfg.add_pre_sharpen_power(1),
+                    '8' : lambda cfg,shift_pressed: cfg.add_pre_sharpen_power(-1),
+                    '(' : lambda cfg,shift_pressed: cfg.add_morph_power(1),
+                    '*' : lambda cfg,shift_pressed: cfg.add_morph_power(-1),
+                    'b' : lambda cfg,shift_pressed: cfg.toggle_two_pass(),
+                    '7' : lambda cfg,shift_pressed: cfg.toggle_debug_mode(),
                     }
             self.masked_keys = list(self.masked_keys_funcs.keys())
 

merger/MergeMasked.py

@@ -16,7 +16,7 @@ def MergeMaskedFace (predictor_func, predictor_input_shape,
                      face_enhancer_func,
                      xseg_256_extract_func,
                      cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmarks, dfl_img):
-
+                     
     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)
 
@@ -60,13 +60,36 @@ def MergeMaskedFace (predictor_func, predictor_input_shape,
 
     dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC )
     dst_face_mask_a_0 = np.clip(dst_face_mask_a_0, 0, 1)
+    
+    
+    if cfg.pre_sharpen_mode > 0 and cfg.pre_sharpen_power != 0:
+        if cfg.pre_sharpen_mode==1:
+            dst_face_bgr = imagelib.gaussian_sharpen(dst_face_bgr, amount=cfg.pre_sharpen_power)
+        elif cfg.pre_sharpen_mode==2:
+            dst_face_bgr = imagelib.unsharpen_mask(dst_face_bgr, amount=cfg.pre_sharpen_power)
+        
+        dst_face_bgr = np.clip(dst_face_bgr, 0, 1, out=dst_face_bgr)
 
     predictor_input_bgr      = cv2.resize (dst_face_bgr, (input_size,input_size) )
 
-    predicted = predictor_func (predictor_input_bgr)
+    
+    predicted = predictor_func (predictor_input_bgr, func_morph_factor = cfg.morph_power/100.0) if cfg.is_morphable else predictor_func (predictor_input_bgr)
+
+    
     prd_face_bgr          = np.clip (predicted[0], 0, 1.0)
     prd_face_mask_a_0     = np.clip (predicted[1], 0, 1.0)
     prd_face_dst_mask_a_0 = np.clip (predicted[2], 0, 1.0)
+    
+    
+    
+    if cfg.two_pass:
+        predicted_2 = predictor_func (prd_face_bgr, func_morph_factor = 1) if cfg.is_morphable else predictor_func (prd_face_bgr)
+        prd_face_bgr = np.clip (predicted_2[0], 0, 1.0)
+        prd_face_mask_a_0     = np.clip (predicted_2[1], 0, 1.0)
+        prd_face_dst_mask_a_0 = np.clip (predicted_2[2], 0, 1.0)
+        
+    if cfg.debug_mode:
+        prd_face_bgr_unchanged = prd_face_bgr.copy()
 
     if cfg.super_resolution_power != 0:
         prd_face_bgr_enhanced = face_enhancer_func(prd_face_bgr, is_tanh=True, preserve_size=False)
@@ -333,6 +356,15 @@ def MergeMaskedFace (predictor_func, predictor_input_shape,
     if out_img is None:
         out_img = img_bgr.copy()
         
+        
+    if 'raw' not in cfg.mode and cfg.debug_mode:
+        ph, pw = predictor_input_bgr.shape[:2]
+        oh, ow = out_img.shape[:2]
+        out_img[oh-ph:,ow-pw:] =  predictor_input_bgr
+        ph, pw = prd_face_bgr_unchanged.shape[:2]
+        out_img[oh-ph:,0:pw] =  prd_face_bgr_unchanged
+    
+        
     return out_img, out_merging_mask_a
 
 

merger/MergerConfig.py

@@ -21,7 +21,7 @@ class MergerConfig(object):
                        ):
         self.type = type
 
-        self.sharpen_dict = {0:"None", 1:'box', 2:'gaussian'}
+        self.sharpen_dict = {0:"None", 1:'box', 2:'gaussian', 3:'unsharpen'}
 
         #default changeable params
         self.sharpen_mode = sharpen_mode
@@ -97,6 +97,9 @@ mask_mode_dict = {0:'full',
 ctm_dict = { 0: "None", 1:"rct", 2:"lct", 3:"mkl", 4:"mkl-m", 5:"idt", 6:"idt-m", 7:"sot-m", 8:"mix-m" }
 ctm_str_dict = {None:0, "rct":1, "lct":2, "mkl":3, "mkl-m":4, "idt":5, "idt-m":6, "sot-m":7, "mix-m":8 }
 
+pre_sharpen_dict = {0:"None", 1:'gaussian'}, # 2:'unsharpen_mask'}
+
+
 class MergerConfigMasked(MergerConfig):
 
     def __init__(self, face_type=FaceType.FULL,
@@ -114,6 +117,12 @@ class MergerConfigMasked(MergerConfig):
                        image_denoise_power = 0,
                        bicubic_degrade_power = 0,
                        color_degrade_power = 0,
+                       pre_sharpen_power = 0,
+                       pre_sharpen_mode=0,
+                       two_pass = False, 
+                       morph_power = 100,
+                       is_morphable = False,
+                       debug_mode = False,
                        **kwargs
                        ):
 
@@ -142,6 +151,12 @@ class MergerConfigMasked(MergerConfig):
         self.image_denoise_power = image_denoise_power
         self.bicubic_degrade_power = bicubic_degrade_power
         self.color_degrade_power = color_degrade_power
+        self.two_pass = two_pass
+        self.pre_sharpen_power = pre_sharpen_power
+        self.pre_sharpen_mode = pre_sharpen_mode
+        self.morph_power = morph_power
+        self.is_morphable = is_morphable
+        self.debug_mode = debug_mode
 
     def copy(self):
         return copy.copy(self)
@@ -152,7 +167,25 @@ class MergerConfigMasked(MergerConfig):
     def toggle_masked_hist_match(self):
         if self.mode == 'hist-match':
             self.masked_hist_match = not self.masked_hist_match
-
+            
+    def toggle_two_pass(self):
+        self.two_pass = not self.two_pass
+        
+    def toggle_debug_mode(self):
+        self.debug_mode = not self.debug_mode
+        
+                
+    def toggle_sharpen_mode_multi(self, pre_sharpen=False):
+        if pre_sharpen:
+            self.toggle_sharpen_mode_presharpen()
+        else:
+            self.toggle_sharpen_mode()
+            
+    def toggle_sharpen_mode_presharpen(self):
+        a = list( pre_sharpen_dict.keys() )
+        self.pre_sharpen_mode = a[ (a.index(self.pre_sharpen_mode)+1) % len(a) ]
+        
+        
     def add_hist_match_threshold(self, diff):
         if self.mode == 'hist-match' or self.mode == 'seamless-hist-match':
             self.hist_match_threshold = np.clip ( self.hist_match_threshold+diff , 0, 255)
@@ -187,6 +220,13 @@ class MergerConfigMasked(MergerConfig):
 
     def add_bicubic_degrade_power(self, diff):
         self.bicubic_degrade_power = np.clip ( self.bicubic_degrade_power+diff, 0, 100)
+        
+    def add_pre_sharpen_power(self, diff):
+        self.pre_sharpen_power = np.clip ( self.pre_sharpen_power+diff, 0, 200)
+        
+    def add_morph_power(self, diff):
+        if self.is_morphable:
+            self.morph_power = np.clip ( self.morph_power+diff , 0, 100)
 
     def ask_settings(self):
         s = """Choose mode: \n"""
@@ -214,6 +254,13 @@ class MergerConfigMasked(MergerConfig):
             self.erode_mask_modifier = np.clip ( io.input_int ("Choose erode mask modifier", 0, add_info="-400..400"), -400, 400)
             self.blur_mask_modifier =  np.clip ( io.input_int ("Choose blur mask modifier", 0, add_info="0..400"), 0, 400)
             self.motion_blur_power = np.clip ( io.input_int ("Choose motion blur power", 0, add_info="0..100"), 0, 100)
+        
+        self.two_pass = io.input_bool("Use two pass mode?", False, help_message="Can enhance results by feeding network output again to network.")
+        self.pre_sharpen_power = np.clip (io.input_int ("Choose pre_sharpen power", 0, help_message="Can enhance results by pre sharping before feeding it to the network.", add_info="0..100" ), 0, 200)
+        
+        if self.is_morphable:
+            self.morph_power = np.clip (io.input_int ("Choose morph_power for moprhable models", 100, add_info="0..100" ), 0, 100)
+
 
         self.output_face_scale = np.clip (io.input_int ("Choose output face scale modifier", 0, add_info="-50..50" ), -50, 50)
 
@@ -249,7 +296,13 @@ class MergerConfigMasked(MergerConfig):
                    self.super_resolution_power == other.super_resolution_power and \
                    self.image_denoise_power == other.image_denoise_power and \
                    self.bicubic_degrade_power == other.bicubic_degrade_power and \
-                   self.color_degrade_power == other.color_degrade_power
+                   self.color_degrade_power == other.color_degrade_power and \
+                   self.pre_sharpen_power == other.pre_sharpen_power and \
+                   self.pre_sharpen_mode == other.pre_sharpen_mode and \
+                   self.two_pass == other.two_pass and \
+                   self.morph_power == other.morph_power and \
+                   self.is_morphable == other.is_morphable and \
+                   self.debug_mode == other.debug_mode 
 
         return False
 
@@ -284,7 +337,14 @@ class MergerConfigMasked(MergerConfig):
             r += (f"""image_denoise_power: {self.image_denoise_power}\n"""
                   f"""bicubic_degrade_power: {self.bicubic_degrade_power}\n"""
                   f"""color_degrade_power: {self.color_degrade_power}\n""")
-
+        
+        r += f"""pre_sharpen_power: {self.pre_sharpen_power}\n"""
+        r += f"""pre_sharpen_mode: {pre_sharpen_dict[self.pre_sharpen_mode]}\n"""
+        r += f"""two_pass: {self.two_pass}\n"""
+        r += f"""morph_power: {self.morph_power}\n"""
+        #r += f"""is_morphable: {self.is_morphable}\n"""
+        r += f"""debug_mode: {self.debug_mode}\n"""
+        
         r += "================"
 
         return r

models/Model_AMP/Model.py

@@ -944,19 +944,16 @@ class AMPModel(ModelBase):
 
     #override
     def get_MergerConfig(self):
-        morph_factor = np.clip ( io.input_number ("Morph factor", 1.0, add_info="0.0 .. 1.0"), 0.0, 1.0 )
-
-        def predictor_morph(face):
-            return self.predictor_func(face, morph_factor)
-
+            
+        def predictor_morph(face, func_morph_factor=1.0):
+            return self.predictor_func(face, func_morph_factor)
 
         import merger
-        return predictor_morph, (self.options['resolution'], self.options['resolution'], 3), merger.MergerConfigMasked(face_type=self.face_type, default_mode = 'overlay')
-
+        return predictor_morph, (self.options['resolution'], self.options['resolution'], 3), merger.MergerConfigMasked(face_type=self.face_type, default_mode = 'overlay', is_morphable=True)
+    
     #override
     def get_config_schema_path(self):
         config_path = Path(__file__).parent.absolute() / Path("config_schema.json")
         return config_path
 
-
 Model = AMPModel