changed help message for pixel loss:

Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.

SAE:
previous SAE model will not work with this update.
Greatly decreased chance of model collapse.
Increased model accuracy.
Residual blocks now default and this option has been removed.
Improved 'learn mask'.
Added masked preview (switch by space key)

Converter:
fixed rct/lct in seamless mode
added mask mode (6) learned*FAN-prd*FAN-dst

added mask editor, its created for refining dataset for FANSeg model, and not for production, but you can spend your time and test it in regular fakes with face obstructions

converters/ConverterMasked.py

@@ -77,11 +77,11 @@ class ConverterMasked(Converter):
                 self.hist_match_threshold = np.clip ( io.input_int("Hist match threshold [0..255] (skip:255) :  ", 255), 0, 255)
 
         if face_type == FaceType.FULL:
-            self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst, (3) FAN-prd, (4) FAN-dst , (5) FAN-prd&dst (?) help. Default - %d : " % (1) , 1, help_message="If you learned mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd&dst' - using multiplied FAN prd and dst mask. "), 1, 5 )
+            self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst, (3) FAN-prd, (4) FAN-dst , (5) FAN-prd*FAN-dst (6) learned*FAN-prd*FAN-dst (?) help. Default - %d : " % (1) , 1, help_message="If you learned mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd*FAN-dst' or 'learned*FAN-prd*FAN-dst' - using multiplied masks."), 1, 6 )
         else:
             self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst . Default - %d : " % (1) , 1), 1, 2 )
 
-        if self.mask_mode == 3 or self.mask_mode == 4 or self.mask_mode == 5:
+        if self.mask_mode >= 3 or self.mask_mode <= 6:
             self.fan_seg = None
 
         if self.mode != 'raw':
@@ -117,7 +117,7 @@ class ConverterMasked(Converter):
 
     #overridable
     def on_cli_initialize(self):
-        if (self.mask_mode == 3 or self.mask_mode == 4 or self.mask_mode == 5) and self.fan_seg == None:
+        if (self.mask_mode >= 3 and self.mask_mode <= 6) and self.fan_seg == None:
             self.fan_seg = FANSegmentator(256, FaceType.toString(FaceType.FULL) )
 
     #override
@@ -167,26 +167,28 @@ class ConverterMasked(Converter):
 
         if self.mask_mode == 2: #dst
             prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC)
-        elif self.mask_mode >= 3 and self.mask_mode <= 5:
+        elif self.mask_mode >= 3 and self.mask_mode <= 6:
 
-            if self.mask_mode == 3 or self.mask_mode == 5: #FAN-prd
+            if self.mask_mode == 3 or self.mask_mode == 5 or self.mask_mode == 6: 
                 prd_face_bgr_256 = cv2.resize (prd_face_bgr, (256,256) )
                 prd_face_bgr_256_mask = self.fan_seg.extract_from_bgr( prd_face_bgr_256[np.newaxis,...] ) [0]
                 FAN_prd_face_mask_a_0 = cv2.resize (prd_face_bgr_256_mask, (output_size,output_size), cv2.INTER_CUBIC)
 
-            if self.mask_mode == 4 or self.mask_mode == 5: #FAN-dst
+            if self.mask_mode == 4 or self.mask_mode == 5 or self.mask_mode == 6: 
                 face_256_mat     = LandmarksProcessor.get_transform_mat (img_face_landmarks, 256, face_type=FaceType.FULL)
                 dst_face_256_bgr = cv2.warpAffine(img_bgr, face_256_mat, (256, 256), flags=cv2.INTER_LANCZOS4 )
                 dst_face_256_mask = self.fan_seg.extract_from_bgr( dst_face_256_bgr[np.newaxis,...] ) [0]
                 FAN_dst_face_mask_a_0 = cv2.resize (dst_face_256_mask, (output_size,output_size), cv2.INTER_CUBIC)
 
-            if self.mask_mode == 3:
+            if self.mask_mode == 3:   #FAN-prd
                 prd_face_mask_a_0 = FAN_prd_face_mask_a_0
-            elif self.mask_mode == 4:
+            elif self.mask_mode == 4: #FAN-dst
                 prd_face_mask_a_0 = FAN_dst_face_mask_a_0
             elif self.mask_mode == 5:
                 prd_face_mask_a_0 = FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0
-
+            elif self.mask_mode == 6:
+                prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0
+                
         prd_face_mask_a_0[ prd_face_mask_a_0 < 0.001 ] = 0.0
 
         prd_face_mask_a   = prd_face_mask_a_0[...,np.newaxis]
@@ -269,11 +271,12 @@ class ConverterMasked(Converter):
                         img_mask_blurry_aaa = cv2.blur(img_mask_blurry_aaa, (blur, blur) )
 
                 img_mask_blurry_aaa = np.clip( img_mask_blurry_aaa, 0, 1.0 )
+                face_mask_blurry_aaa = cv2.warpAffine( img_mask_blurry_aaa, face_mat, (output_size, output_size) )
 
                 if debug:
                     debugs += [img_mask_blurry_aaa.copy()]
 
-                if self.color_transfer_mode is not None:
+                if 'seamless' not in self.mode and self.color_transfer_mode is not None:
                     if self.color_transfer_mode == 'rct':
                         if debug:
                             debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]
@@ -309,15 +312,20 @@ class ConverterMasked(Converter):
 
                     if self.masked_hist_match:
                         hist_mask_a *= prd_face_mask_a
+                        
+                    white =  (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32)
 
-                    hist_match_1 = prd_face_bgr*hist_mask_a + (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32)
+                    hist_match_1 = prd_face_bgr*hist_mask_a + white
                     hist_match_1[ hist_match_1 > 1.0 ] = 1.0
 
-                    hist_match_2 = dst_face_bgr*hist_mask_a + (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32)
+                    hist_match_2 = dst_face_bgr*hist_mask_a + white
                     hist_match_2[ hist_match_1 > 1.0 ] = 1.0
 
                     prd_face_bgr = imagelib.color_hist_match(hist_match_1, hist_match_2, self.hist_match_threshold )
-
+                    
+                    #if self.masked_hist_match:
+                    #    prd_face_bgr -= white
+                        
                 if self.mode == 'hist-match-bw':
                     prd_face_bgr = prd_face_bgr.astype(dtype=np.float32)
 
@@ -364,6 +372,35 @@ class ConverterMasked(Converter):
 
                 out_img = np.clip( img_bgr*(1-img_mask_blurry_aaa) + (out_img*img_mask_blurry_aaa) , 0, 1.0 )
 
+                if 'seamless' in self.mode and self.color_transfer_mode is not None:
+                    out_face_bgr = cv2.warpAffine( out_img, face_mat, (output_size, output_size) )
+
+                    if self.color_transfer_mode == 'rct':
+                        if debug:
+                            debugs += [ np.clip( cv2.warpAffine( out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]
+
+                        new_out_face_bgr = imagelib.reinhard_color_transfer ( np.clip( (out_face_bgr*255).astype(np.uint8), 0, 255),
+                                                                              np.clip( (dst_face_bgr*255).astype(np.uint8), 0, 255),
+                                                                             source_mask=face_mask_blurry_aaa, target_mask=face_mask_blurry_aaa)
+                        new_out_face_bgr = np.clip( new_out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0)
+
+                        if debug:
+                            debugs += [ np.clip( cv2.warpAffine( new_out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]
+
+
+                    elif self.color_transfer_mode == 'lct':
+                        if debug:
+                            debugs += [ np.clip( cv2.warpAffine( out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]
+
+                        new_out_face_bgr = imagelib.linear_color_transfer (out_face_bgr, dst_face_bgr)
+                        new_out_face_bgr = np.clip( new_out_face_bgr, 0.0, 1.0)
+
+                        if debug:
+                            debugs += [ np.clip( cv2.warpAffine( new_out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ]
+                    
+                    new_out = cv2.warpAffine( new_out_face_bgr, face_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT )
+                    out_img =  np.clip( img_bgr*(1-img_mask_blurry_aaa) + (new_out*img_mask_blurry_aaa) , 0, 1.0 )
+
                 if self.mode == 'seamless-hist-match':
                     out_face_bgr = cv2.warpAffine( out_img, face_mat, (output_size, output_size) )
                     new_out_face_bgr = imagelib.color_hist_match(out_face_bgr, dst_face_bgr, self.hist_match_threshold)

facelib/LandmarksProcessor.py

@@ -4,6 +4,7 @@ import numpy as np
 from enum import IntEnum
 import mathlib
 import imagelib
+from imagelib import IEPolys
 from mathlib.umeyama import umeyama
 from facelib import FaceType
 import math
@@ -153,7 +154,7 @@ def transform_points(points, mat, invert=False):
     return points
 
 
-def get_image_hull_mask (image_shape, image_landmarks):
+def get_image_hull_mask (image_shape, image_landmarks, ie_polys=None):
     if len(image_landmarks) != 68:
         raise Exception('get_image_hull_mask works only with 68 landmarks')
     int_lmrks = np.array(image_landmarks, dtype=np.int)
@@ -198,6 +199,9 @@ def get_image_hull_mask (image_shape, image_landmarks):
     #nose
     cv2.fillConvexPoly( hull_mask, cv2.convexHull(int_lmrks[27:36]), (1,) )
 
+    if ie_polys is not None:
+        ie_polys.overlay_mask(hull_mask)
+
     return hull_mask
 
 def get_image_eye_mask (image_shape, image_landmarks):
@@ -211,11 +215,6 @@ def get_image_eye_mask (image_shape, image_landmarks):
 
     return hull_mask
 
-def get_image_hull_mask_3D (image_shape, image_landmarks):
-    result = get_image_hull_mask(image_shape, image_landmarks)
-
-    return np.repeat ( result, (3,), -1 )
-
 def blur_image_hull_mask (hull_mask):
 
     maxregion = np.argwhere(hull_mask==1.0)
@@ -235,9 +234,6 @@ def blur_image_hull_mask (hull_mask):
 
     return hull_mask
 
-def get_blurred_image_hull_mask(image_shape, image_landmarks):
-    return blur_image_hull_mask ( get_image_hull_mask(image_shape, image_landmarks) )
-
 mirror_idxs = [
     [0,16],
     [1,15],
@@ -282,7 +278,7 @@ def mirror_landmarks (landmarks, val):
     result[:,0] = val - result[:,0] - 1
     return result
 
-def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=False):
+def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=False, ie_polys=None):
     if len(image_landmarks) != 68:
         raise Exception('get_image_eye_mask works only with 68 landmarks')
 
@@ -310,11 +306,11 @@ def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=Fa
         cv2.circle(image, (x, y), 2, color, lineType=cv2.LINE_AA)
 
     if transparent_mask:
-        mask = get_image_hull_mask (image.shape, image_landmarks)
+        mask = get_image_hull_mask (image.shape, image_landmarks, ie_polys)
         image[...] = ( image * (1-mask) + image * mask / 2 )[...]
 
-def draw_rect_landmarks (image, rect, image_landmarks, face_size, face_type, transparent_mask=False, landmarks_color=(0,255,0) ):
+def draw_rect_landmarks (image, rect, image_landmarks, face_size, face_type, transparent_mask=False, ie_polys=None, landmarks_color=(0,255,0) ):
-    draw_landmarks(image, image_landmarks, color=landmarks_color, transparent_mask=transparent_mask)
+    draw_landmarks(image, image_landmarks, color=landmarks_color, transparent_mask=transparent_mask, ie_polys=ie_polys)
     imagelib.draw_rect (image, rect, (255,0,0), 2 )
 
     image_to_face_mat = get_transform_mat (image_landmarks, face_size, face_type)

imagelib/IEPolys.py

@@ -0,0 +1,104 @@
+import numpy as np
+import cv2
+
+class IEPolysPoints:
+    def __init__(self, IEPolys_parent, type):
+        self.parent = IEPolys_parent
+        self.type = type
+        self.points = np.empty( (0,2), dtype=np.int32 )
+        self.n_max = self.n = 0
+
+    def add(self,x,y):
+        self.points = np.append(self.points[0:self.n], [ (x,y) ], axis=0)
+        self.n_max = self.n = self.n + 1
+        self.parent.dirty = True
+
+    def n_dec(self):
+        self.n = max(0, self.n-1)
+        self.parent.dirty = True
+        return self.n
+
+    def n_inc(self):
+        self.n = min(len(self.points), self.n+1)
+        self.parent.dirty = True
+        return self.n
+
+    def n_clip(self):
+        self.points = self.points[0:self.n]
+        self.n_max = self.n
+
+    def cur_point(self):
+        return self.points[self.n-1]
+
+    def points_to_n(self):
+        return self.points[0:self.n]
+
+    def set_points(self, points):
+        self.points = np.array(points)
+        self.n_max = self.n = len(points)
+        self.parent.dirty = True
+
+class IEPolys:
+    def __init__(self):
+        self.list = []
+        self.n_max = self.n = 0
+        self.dirty = True
+
+    def add(self, type):
+        self.list = self.list[0:self.n]
+        self.list.append ( IEPolysPoints(self, type) )
+        self.n_max = self.n = self.n + 1
+        self.dirty = True
+
+    def n_dec(self):
+        self.n = max(0, self.n-1)
+        self.dirty = True
+        return self.n
+        
+    def n_inc(self):
+        self.n = min(len(self.list), self.n+1)
+        self.dirty = True
+        return self.n
+
+    def n_list(self):
+        return self.list[self.n-1]
+
+    def n_clip(self):
+        self.list = self.list[0:self.n]
+        self.n_max = self.n
+        if self.n > 0:
+            self.list[-1].n_clip()
+
+    def __iter__(self):
+        for n in range(self.n):
+            yield self.list[n]
+
+    def switch_dirty(self):
+        d = self.dirty
+        self.dirty = False
+        return d
+
+    def overlay_mask(self, mask):
+        h,w,c = mask.shape
+        white = (1,)*c
+        black = (0,)*c
+        for n in range(self.n):
+            poly = self.list[n]
+            if poly.n > 0:
+                cv2.fillPoly(mask, [poly.points_to_n()], white if poly.type == 1 else black )
+
+    def dump(self):
+        result = []
+        for n in range(self.n):
+            l = self.list[n]
+            result += [ (l.type, l.points_to_n().tolist() ) ]
+        return result
+
+    @staticmethod
+    def load(ie_polys=None):
+        obj = IEPolys()
+        if ie_polys is not None:
+            for (type, points) in ie_polys:
+                obj.add(type)
+                obj.n_list().set_points(points)
+        return obj

imagelib/__init__.py

@@ -20,4 +20,6 @@ from .color_transfer import linear_color_transfer
 
 from .DCSCN import DCSCN
 
-from .common import normalize_channels
+from .common import normalize_channels
+
+from .IEPolys import IEPolys

imagelib/text.py

@@ -15,23 +15,24 @@ def _get_pil_font (font, size):
         return ImageFont.load_default()
 
 def get_text_image( shape, text, color=(1,1,1), border=0.2, font=None):
-    try:
+    h,w,c = shape
-        size = shape[1]
+    try:        
-        pil_font = _get_pil_font( localization.get_default_ttf_font_name() , size)
+        pil_font = _get_pil_font( localization.get_default_ttf_font_name() , h-2)
-        text_width, text_height = pil_font.getsize(text)
 
-        canvas = Image.new('RGB', shape[0:2], (0,0,0) )
+        canvas = Image.new('RGB', (w,h) , (0,0,0) )
         draw = ImageDraw.Draw(canvas)
         offset = ( 0, 0)
         draw.text(offset, text, font=pil_font, fill=tuple((np.array(color)*255).astype(np.int)) )
 
         result = np.asarray(canvas) / 255
-        if shape[2] != 3:
+        
-            result = np.concatenate ( (result, np.ones ( (shape[1],) + (shape[0],) + (shape[2]-3,)) ), axis=2 )
+        if c > 3:
-
+            result = np.concatenate ( (result, np.ones ((h,w,c-3)) ), axis=-1 )
+        elif c < 3:
+            result = result[...,0:c]
         return result
     except:
-        return np.zeros ( (shape[1], shape[0], shape[2]), dtype=np.float32 )
+        return np.zeros ( (h,w,c) )
 
 def draw_text( image, rect, text, color=(1,1,1), border=0.2, font=None):
     h,w,c = image.shape

interact/interact.py

@@ -18,8 +18,10 @@ except:
 class InteractBase(object):
     EVENT_LBUTTONDOWN = 1
     EVENT_LBUTTONUP = 2
+    EVENT_MBUTTONDOWN = 3
+    EVENT_MBUTTONUP = 4
     EVENT_RBUTTONDOWN = 5
-    EVENT_RBUTTONUP = 6
+    EVENT_RBUTTONUP = 6    
     EVENT_MOUSEWHEEL = 10
 
     def __init__(self):
@@ -263,6 +265,8 @@ class InteractDesktop(InteractBase):
             elif event == cv2.EVENT_LBUTTONUP: ev = InteractBase.EVENT_LBUTTONUP
             elif event == cv2.EVENT_RBUTTONDOWN: ev = InteractBase.EVENT_RBUTTONDOWN
             elif event == cv2.EVENT_RBUTTONUP: ev = InteractBase.EVENT_RBUTTONUP
+            elif event == cv2.EVENT_MBUTTONDOWN: ev = InteractBase.EVENT_MBUTTONDOWN
+            elif event == cv2.EVENT_MBUTTONUP: ev = InteractBase.EVENT_MBUTTONUP
             elif event == cv2.EVENT_MOUSEWHEEL: ev = InteractBase.EVENT_MOUSEWHEEL
 
             else: ev = 0

main.py

@@ -186,24 +186,23 @@ if __name__ == "__main__":
     p.add_argument('--lossless', action="store_true", dest="lossless", default=False, help="PNG codec.")
     p.set_defaults(func=process_videoed_video_from_sequence)
 
-    def process_labelingtool(arguments):
+    def process_labelingtool_edit_mask(arguments):
-        from mainscripts import LabelingTool
+        from mainscripts import MaskEditorTool
-        LabelingTool.main (arguments.input_dir, arguments.output_dir)
+        MaskEditorTool.mask_editor_main (arguments.input_dir, arguments.confirmed_dir, arguments.skipped_dir)
 
-    p = subparsers.add_parser( "labelingtool", help="Labeling tool.")
+    labeling_parser = subparsers.add_parser( "labelingtool", help="Labeling tool.").add_subparsers()
+    p = labeling_parser.add_parser ( "edit_mask", help="")
     p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.")
-    p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the labeled faces will be stored.")
+    p.add_argument('--confirmed-dir', required=True, action=fixPathAction, dest="confirmed_dir", help="This is where the labeled faces will be stored.")
-    p.set_defaults(func=process_labelingtool)
+    p.add_argument('--skipped-dir', required=True, action=fixPathAction, dest="skipped_dir", help="This is where the labeled faces will be stored.")
-
+    p.set_defaults(func=process_labelingtool_edit_mask)
+    
     def bad_args(arguments):
         parser.print_help()
         exit(0)
     parser.set_defaults(func=bad_args)
 
     arguments = parser.parse_args()
-
-    #os.environ['force_plaidML'] = '1'
-
     arguments.func(arguments)
 
     print ("Done.")

mainscripts/LabelingTool_unfinished.py

@@ -1,287 +0,0 @@
-import traceback
-import os
-import sys
-import time
-import numpy as np
-import numpy.linalg as npl
-
-import cv2
-from pathlib import Path
-from interact import interact as io
-from utils.cv2_utils import *
-from utils import Path_utils
-from utils.DFLPNG import DFLPNG
-from utils.DFLJPG import DFLJPG
-from facelib import LandmarksProcessor
-
-def main(input_dir, output_dir):
-    input_path = Path(input_dir)
-    output_path = Path(output_dir)
-
-    if not input_path.exists():
-        raise ValueError('Input directory not found. Please ensure it exists.')
-
-    if not output_path.exists():
-        output_path.mkdir(parents=True)
-
-    wnd_name = "Labeling tool"
-    io.named_window (wnd_name)
-    io.capture_mouse(wnd_name)
-    io.capture_keys(wnd_name)
-
-    #for filename in io.progress_bar_generator (Path_utils.get_image_paths(input_path), desc="Labeling"):
-    for filename in Path_utils.get_image_paths(input_path):
-        filepath = Path(filename)
-
-        if filepath.suffix == '.png':
-            dflimg = DFLPNG.load( str(filepath) )
-        elif filepath.suffix == '.jpg':
-            dflimg = DFLJPG.load ( str(filepath) )
-        else:
-            dflimg = None
-
-        if dflimg is None:
-            io.log_err ("%s is not a dfl image file" % (filepath.name) )
-            continue
-
-        lmrks = dflimg.get_landmarks()
-        lmrks_list = lmrks.tolist()
-        orig_img = cv2_imread(str(filepath))
-        h,w,c = orig_img.shape
-
-        mask_orig = LandmarksProcessor.get_image_hull_mask( orig_img.shape, lmrks).astype(np.uint8)[:,:,0]
-        ero_dil_rate = w // 8
-        mask_ero = cv2.erode (mask_orig, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero_dil_rate,ero_dil_rate)), iterations = 1 )
-        mask_dil = cv2.dilate(mask_orig, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero_dil_rate,ero_dil_rate)), iterations = 1 )
-
-
-        #mask_bg = np.zeros(orig_img.shape[:2],np.uint8)
-        mask_bg = 1-mask_dil
-        mask_bgp = np.ones(orig_img.shape[:2],np.uint8) #default - all background possible
-        mask_fg = np.zeros(orig_img.shape[:2],np.uint8)
-        mask_fgp = np.zeros(orig_img.shape[:2],np.uint8)
-
-        img = orig_img.copy()
-
-        l_thick=2
-
-        def draw_4_lines (masks_out, pts, thickness=1):
-            fgp,fg,bg,bgp = masks_out
-            h,w = fg.shape
-
-            fgp_pts = []
-            fg_pts = np.array([ pts[i:i+2]  for i in range(len(pts)-1)])
-            bg_pts = []
-            bgp_pts = []
-
-            for i in range(len(fg_pts)):
-                a, b = line = fg_pts[i]
-
-                ba = b-a
-                v = ba / npl.norm(ba)
-
-                ccpv = np.array([v[1],-v[0]])
-                cpv = np.array([-v[1],v[0]])
-                step = 1 / max(np.abs(cpv))
-
-                fgp_pts.append ( np.clip (line + ccpv * step * thickness,     0, w-1 ).astype(np.int) )
-                bg_pts.append  ( np.clip (line +  cpv * step * thickness,     0, w-1 ).astype(np.int) )
-                bgp_pts.append ( np.clip (line +  cpv * step * thickness * 2, 0, w-1 ).astype(np.int) )
-
-            fgp_pts = np.array(fgp_pts)
-            bg_pts = np.array(bg_pts)
-            bgp_pts = np.array(bgp_pts)
-
-            cv2.polylines(fgp, fgp_pts, False, (1,), thickness=thickness)
-            cv2.polylines(fg,  fg_pts,  False, (1,), thickness=thickness)
-            cv2.polylines(bg,  bg_pts,  False, (1,), thickness=thickness)
-            cv2.polylines(bgp, bgp_pts, False, (1,), thickness=thickness)
-
-        def draw_lines ( masks_steps, pts, thickness=1):
-            lines = np.array([ pts[i:i+2] for i in range(len(pts)-1)])
-
-
-            for mask, step in masks_steps:
-                h,w = mask.shape
-
-                mask_lines = []
-                for i in range(len(lines)):
-                    a, b = line = lines[i]
-                    ba = b-a
-                    ba_len = npl.norm(ba)
-                    if ba_len != 0:
-                        v = ba / ba_len
-                        pv = np.array([-v[1],v[0]])
-                        pv_inv_max = 1 / max(np.abs(pv))
-                        mask_lines.append ( np.clip (line + pv * pv_inv_max * thickness * step, 0, w-1 ).astype(np.int) )
-                    else:
-                        mask_lines.append ( np.array(line, dtype=np.int) )
-                cv2.polylines(mask, mask_lines, False, (1,), thickness=thickness)
-
-        def draw_fill_convex( mask_out, pts, scale=1.0 ):
-            hull = cv2.convexHull(np.array(pts))
-
-            if scale !=1.0:
-                pts_count = hull.shape[0]
-
-                sum_x = np.sum(hull[:, 0, 0])
-                sum_y = np.sum(hull[:, 0, 1])
-
-                hull_center = np.array([sum_x/pts_count, sum_y/pts_count])
-                hull = hull_center+(hull-hull_center)*scale
-                hull = hull.astype(pts.dtype)
-            cv2.fillConvexPoly( mask_out, hull, (1,) )
-
-        def get_gc_mask_bgr(gc_mask):
-            h, w = gc_mask.shape
-            bgr = np.zeros( (h,w,3), dtype=np.uint8 )
-
-            bgr [ gc_mask == 0 ] = (0,0,0)
-            bgr [ gc_mask == 1 ] = (255,255,255)
-            bgr [ gc_mask == 2 ] = (0,0,255) #RED
-            bgr [ gc_mask == 3 ] = (0,255,0) #GREEN
-            return bgr
-
-        def get_gc_mask_result(gc_mask):
-            return np.where((gc_mask==1) + (gc_mask==3),1,0).astype(np.int)
-
-        #convex inner of right chin to end of right eyebrow
-        #draw_fill_convex ( mask_fgp, lmrks_list[8:17]+lmrks_list[26:27] )
-
-        #convex inner of start right chin to right eyebrow
-        #draw_fill_convex ( mask_fgp, lmrks_list[8:9]+lmrks_list[22:27] )
-
-        #convex inner of nose
-        draw_fill_convex ( mask_fgp, lmrks[27:36] )
-
-        #convex inner of nose half
-        draw_fill_convex ( mask_fg, lmrks[27:36], scale=0.5 )
-
-
-        #left corner of mouth to left corner of nose
-        #draw_lines ( [ (mask_fg,0),   ], lmrks_list[49:50]+lmrks_list[32:33], l_thick)
-
-        #convex inner: right corner of nose to centers of eyebrows
-        #draw_fill_convex ( mask_fgp, lmrks_list[35:36]+lmrks_list[19:20]+lmrks_list[24:25])
-
-        #right corner of mouth to right corner of nose
-        #draw_lines ( [ (mask_fg,0),   ], lmrks_list[54:55]+lmrks_list[35:36], l_thick)
-
-        #left eye
-        #draw_fill_convex ( mask_fg, lmrks_list[36:40] )
-        #right eye
-        #draw_fill_convex ( mask_fg, lmrks_list[42:48] )
-
-        #right chin
-        draw_lines ( [ (mask_bg,0), (mask_fg,-1),   ], lmrks[8:17], l_thick)
-
-        #left eyebrow center to right eyeprow center
-        draw_lines ( [ (mask_bg,-1), (mask_fg,0),   ], lmrks_list[19:20] + lmrks_list[24:25], l_thick)
-        #        #draw_lines ( [ (mask_bg,-1), (mask_fg,0),   ], lmrks_list[24:25] + lmrks_list[19:17:-1], l_thick)
-
-        #half right eyebrow to end of right chin
-        draw_lines ( [ (mask_bg,-1), (mask_fg,0),   ], lmrks_list[24:27] + lmrks_list[16:17], l_thick)
-
-        #import code
-        #code.interact(local=dict(globals(), **locals()))
-
-        #compose mask layers
-        gc_mask = np.zeros(orig_img.shape[:2],np.uint8)
-        gc_mask [ mask_bgp==1 ] = 2
-        gc_mask [ mask_fgp==1 ] = 3
-        gc_mask [ mask_bg==1  ] = 0
-        gc_mask [ mask_fg==1  ] = 1
-
-        gc_bgr_before = get_gc_mask_bgr (gc_mask)
-
-
-
-        #io.show_image (wnd_name, gc_mask )
-
-        ##points, hierarcy = cv2.findContours(original_mask,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
-        ##gc_mask = ( (1-erode_mask)*2 + erode_mask )# * dilate_mask
-        #gc_mask = (1-erode_mask)*2 + erode_mask
-        #cv2.addWeighted(
-        #gc_mask = mask_0_27 + (1-mask_0_27)*2
-        #
-        ##import code
-        ##code.interact(local=dict(globals(), **locals()))
-        #
-        #rect = (1,1,img.shape[1]-2,img.shape[0]-2)
-        #
-        #
-        cv2.grabCut(img,gc_mask,None,np.zeros((1,65),np.float64),np.zeros((1,65),np.float64),5, cv2.GC_INIT_WITH_MASK)
-
-        gc_bgr = get_gc_mask_bgr (gc_mask)
-        gc_mask_result = get_gc_mask_result(gc_mask)
-        gc_mask_result_1 = gc_mask_result[:,:,np.newaxis]
-
-        #import code
-        #code.interact(local=dict(globals(), **locals()))
-        orig_img_gc_layers_masked = (0.5*orig_img + 0.5*gc_bgr).astype(np.uint8)
-        orig_img_gc_before_layers_masked = (0.5*orig_img + 0.5*gc_bgr_before).astype(np.uint8)
-
-
-
-        pink_bg = np.full ( orig_img.shape, (255,0,255), dtype=np.uint8 )
-
-
-        orig_img_result = orig_img * gc_mask_result_1
-        orig_img_result_pinked = orig_img_result + pink_bg * (1-gc_mask_result_1)
-
-        #io.show_image (wnd_name, blended_img)
-
-        ##gc_mask, bgdModel, fgdModel =
-        #
-        #mask2 = np.where((gc_mask==1) + (gc_mask==3),255,0).astype('uint8')[:,:,np.newaxis]
-        #mask2 = np.repeat(mask2, (3,), -1)
-        #
-        ##mask2 = np.where(gc_mask!=0,255,0).astype('uint8')
-        #blended_img = orig_img #-\
-        #              #0.3 * np.full(original_img.shape, (50,50,50)) * (1-mask_0_27)[:,:,np.newaxis]
-        #              #0.3 * np.full(original_img.shape, (50,50,50)) * (1-dilate_mask)[:,:,np.newaxis] +\
-        #              #0.3 * np.full(original_img.shape, (50,50,50)) * (erode_mask)[:,:,np.newaxis]
-        #blended_img = np.clip(blended_img, 0, 255).astype(np.uint8)
-        ##import code
-        ##code.interact(local=dict(globals(), **locals()))
-        orig_img_lmrked = orig_img.copy()
-        LandmarksProcessor.draw_landmarks(orig_img_lmrked, lmrks, transparent_mask=True)
-
-        screen = np.concatenate ([orig_img_gc_before_layers_masked,
-                                  orig_img_gc_layers_masked,
-                                  orig_img,
-                                  orig_img_lmrked,
-                                  orig_img_result_pinked,
-                                  orig_img_result,
-                                  ], axis=1)
-
-        io.show_image (wnd_name, screen.astype(np.uint8) )
-
-
-        while True:
-            io.process_messages()
-
-            for (x,y,ev,flags) in io.get_mouse_events(wnd_name):
-                pass
-                #print (x,y,ev,flags)
-
-            key_events = [ ev for ev, in io.get_key_events(wnd_name) ]
-            for key in key_events:
-                if key == ord('1'):
-                    pass
-                if key == ord('2'):
-                    pass
-                if key == ord('3'):
-                    pass
-
-            if ord(' ') in key_events:
-                break
-
-    import code
-    code.interact(local=dict(globals(), **locals()))
-
-
-
-
-#original_mask = np.ones(original_img.shape[:2],np.uint8)*2
-#cv2.drawContours(original_mask, points, -1, (1,), 1)

mainscripts/MaskEditorTool.py

@@ -0,0 +1,376 @@
+import os
+import sys
+import time
+import traceback
+from pathlib import Path
+
+import cv2
+import numpy as np
+import numpy.linalg as npl
+
+import imagelib
+from facelib import LandmarksProcessor
+from imagelib import IEPolys
+from interact import interact as io
+from utils import Path_utils
+from utils.cv2_utils import *
+from utils.DFLJPG import DFLJPG
+from utils.DFLPNG import DFLPNG
+
+class MaskEditor:
+    STATE_NONE=0
+    STATE_MASKING=1
+    
+    def __init__(self, img, mask=None, ie_polys=None, get_status_lines_func=None):
+        self.img = imagelib.normalize_channels (img,3)
+        h, w, c = img.shape
+        ph, pw = h // 4, w // 4
+        
+        if mask is not None:
+            self.mask = imagelib.normalize_channels (mask,3)
+        else:
+            self.mask = np.zeros ( (h,w,3) )
+        self.get_status_lines_func = get_status_lines_func
+        
+        self.state_prop = self.STATE_NONE
+        
+        self.w, self.h = w, h
+        self.pw, self.ph = pw, ph
+        self.pwh = np.array([self.pw, self.ph])
+        self.pwh2 = np.array([self.pw*2, self.ph*2])
+        self.sw, self.sh = w+pw*2, h+ph*2
+
+        if ie_polys is None:
+            ie_polys = IEPolys()
+        self.ie_polys = ie_polys
+        
+        self.polys_mask = None
+        
+        self.screen_status_block = None
+        self.screen_status_block_dirty = True
+    
+    def set_state(self, state):
+        self.state = state
+    
+    @property
+    def state(self):
+        return self.state_prop
+
+    @state.setter
+    def state(self, value):
+        self.state_prop = value
+        if value == self.STATE_MASKING:
+            self.ie_polys.dirty = True
+                  
+    def get_mask(self):        
+        if self.ie_polys.switch_dirty():
+            self.screen_status_block_dirty = True
+            self.ie_mask = img = self.mask.copy()
+
+            self.ie_polys.overlay_mask(img)     
+                          
+            return img
+        return self.ie_mask
+        
+    def get_screen_overlay(self):
+        img = np.zeros ( (self.sh, self.sw, 3) )
+        
+        if self.state == self.STATE_MASKING:
+            mouse_xy = self.mouse_xy.copy() + self.pwh
+
+            l = self.ie_polys.n_list()
+            if l.n > 0:            
+                p = l.cur_point().copy() + self.pwh
+                color = (0,1,0) if l.type == 1 else (0,0,1)
+                cv2.line(img, tuple(p), tuple(mouse_xy), color )
+
+        return img
+    
+    def undo_to_begin_point(self):
+        while not self.undo_point():
+            pass
+        
+    def undo_point(self):
+        if self.state == self.STATE_NONE:
+            if self.ie_polys.n > 0:                   
+                self.state = self.STATE_MASKING
+                
+        if self.state == self.STATE_MASKING:
+            if self.ie_polys.n_list().n_dec() == 0 and \
+               self.ie_polys.n_dec() == 0:
+                self.state = self.STATE_NONE
+            else:
+                return False
+                    
+        return True
+        
+    def redo_to_end_point(self):
+        while not self.redo_point():
+            pass
+                
+    def redo_point(self):
+        if self.state == self.STATE_NONE:
+            if self.ie_polys.n_max > 0:                    
+                self.state = self.STATE_MASKING
+                if self.ie_polys.n == 0:
+                    self.ie_polys.n_inc()
+                    
+        if self.state == self.STATE_MASKING:
+            while True:
+                l = self.ie_polys.n_list()
+                if l.n_inc() == l.n_max:
+                    if self.ie_polys.n == self.ie_polys.n_max:
+                        break
+                    self.ie_polys.n_inc()
+                else:
+                    return False
+                    
+        return True
+        
+    def combine_screens(self, screens):
+    
+        screens_len = len(screens)
+        
+        new_screens = []
+        for screen, padded_overlay in screens:
+            screen_img = np.zeros( (self.sh, self.sw, 3), dtype=np.float32 )
+            
+            screen = imagelib.normalize_channels (screen, 3)
+            h,w,c = screen.shape
+            
+            screen_img[self.ph:-self.ph, self.pw:-self.pw, :] = screen
+            
+            if padded_overlay is not None:
+                screen_img = screen_img + padded_overlay
+            
+            screen_img = np.clip(screen_img*255, 0, 255).astype(np.uint8)
+            new_screens.append(screen_img)
+            
+        return np.concatenate (new_screens, axis=1)
+    
+    def get_screen_status_block(self, w, c):
+        if self.screen_status_block_dirty:
+            self.screen_status_block_dirty = False
+            lines = [
+                    'Polys current/max = %d/%d' % (self.ie_polys.n, self.ie_polys.n_max),
+                    ]
+            if self.get_status_lines_func is not None:
+                lines += self.get_status_lines_func()
+                
+            lines_count = len(lines)
+    
+            
+            h_line = 21
+            h = lines_count * h_line
+            img = np.ones ( (h,w,c) ) * 0.1
+            
+            for i in range(lines_count):
+                img[ i*h_line:(i+1)*h_line, 0:w] += \
+                    imagelib.get_text_image (  (h_line,w,c), lines[i], color=[0.8]*c )
+                
+            self.screen_status_block = np.clip(img*255, 0, 255).astype(np.uint8)
+            
+        return self.screen_status_block
+        
+    def set_screen_status_block_dirty(self):
+        self.screen_status_block_dirty = True
+        
+    def make_screen(self):
+       
+        
+        
+        screen_overlay = self.get_screen_overlay()
+        final_mask = self.get_mask()
+  
+        masked_img = self.img*final_mask*0.5 + self.img*(1-final_mask)
+
+        pink = np.full ( (self.h, self.w, 3), (1,0,1) )        
+        pink_masked_img = self.img*final_mask + pink*(1-final_mask)
+        
+        screens = [ (self.img, None),
+                    (masked_img, screen_overlay),
+                    (pink_masked_img, screen_overlay),                    
+                    ]                
+        screens = self.combine_screens(screens)
+        
+        status_img = self.get_screen_status_block( screens.shape[1], screens.shape[2] )
+           
+        result = np.concatenate ( [screens, status_img], axis=0  )
+        
+        return result
+        
+    def mask_finish(self, n_clip=True):
+        if self.state == self.STATE_MASKING:
+            if self.ie_polys.n_list().n <= 2:
+                self.ie_polys.n_dec()
+            self.state = self.STATE_NONE
+            if n_clip:
+                self.ie_polys.n_clip()
+            
+    def set_mouse_pos(self,x,y):
+        mouse_x = x % (self.sw) - self.pw
+        mouse_y = y % (self.sh) - self.ph
+        self.mouse_xy = np.array( [mouse_x, mouse_y] )
+        self.mouse_x, self.mouse_y = self.mouse_xy
+
+    def mask_point(self, type):
+        if self.state == self.STATE_MASKING and \
+           self.ie_polys.n_list().type != type:
+            self.mask_finish()
+            
+        elif self.state == self.STATE_NONE:
+            self.state = self.STATE_MASKING            
+            self.ie_polys.add(type)
+        
+        if self.state == self.STATE_MASKING:
+            self.ie_polys.n_list().add (self.mouse_x, self.mouse_y)            
+
+    def get_ie_polys(self):
+        return self.ie_polys
+        
+def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None):
+    input_path = Path(input_dir)
+    
+    confirmed_path = Path(confirmed_dir)
+    skipped_path = Path(skipped_dir)
+
+    if not input_path.exists():
+        raise ValueError('Input directory not found. Please ensure it exists.')
+
+    if not confirmed_path.exists():
+        confirmed_path.mkdir(parents=True)
+
+    if not skipped_path.exists():
+        skipped_path.mkdir(parents=True)
+
+    wnd_name = "MaskEditor tool"
+    io.named_window (wnd_name)
+    io.capture_mouse(wnd_name)
+    io.capture_keys(wnd_name)
+
+    image_paths = [ Path(x) for x in Path_utils.get_image_paths(input_path)]
+    done_paths = []
+    
+    image_paths_total = len(image_paths)
+    
+    
+                    
+    
+    is_exit = False
+    while not is_exit:
+        
+        if len(image_paths) > 0:
+            filepath = image_paths.pop(0)
+        else:
+            filepath = None
+
+        if filepath is not None:
+            if filepath.suffix == '.png':
+                dflimg = DFLPNG.load( str(filepath) )
+            elif filepath.suffix == '.jpg':
+                dflimg = DFLJPG.load ( str(filepath) )
+            else:
+                dflimg = None
+                
+            if dflimg is None:
+                io.log_err ("%s is not a dfl image file" % (filepath.name) )
+                continue
+                
+            lmrks = dflimg.get_landmarks()
+            ie_polys = dflimg.get_ie_polys()
+            
+            img = cv2_imread(str(filepath)) / 255.0
+            mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks)
+        else:
+            img = np.zeros ( (256,256,3) )
+            mask = np.ones ( (256,256,3) )
+            ie_polys = None
+            
+        def get_status_lines_func():
+            return ['Progress: %d / %d . Current file: %s' % (len(done_paths), image_paths_total, str(filepath.name) if filepath is not None else "end" ), 
+                    '[Left mouse button] - mark include mask.',
+                    '[Right mouse button] - mark exclude mask.',
+                    '[Middle mouse button] - finish current poly.',
+                    '[Mouse wheel] - undo/redo poly or point. [+ctrl] - undo to begin/redo to end',
+                    '[q] - prev image. [w] - skip and move to %s. [e] - save and move to %s. ' % (skipped_path.name, confirmed_path.name),
+                    '[z] - prev image. [x] - skip. [c] - save. ',
+                    '[esc] - quit'
+                    ]
+        ed = MaskEditor(img, mask, ie_polys, get_status_lines_func)
+
+        next = False
+        while not next:
+            io.process_messages(0.005)
+
+            for (x,y,ev,flags) in io.get_mouse_events(wnd_name):
+                ed.set_mouse_pos(x, y) 
+                if filepath is not None:
+                    if ev == io.EVENT_LBUTTONDOWN:
+                        ed.mask_point(1)
+                    elif ev == io.EVENT_RBUTTONDOWN:
+                        ed.mask_point(0)
+                    elif ev == io.EVENT_MBUTTONDOWN:
+                        ed.mask_finish()
+                    elif ev == io.EVENT_MOUSEWHEEL:
+                        if flags & 0x80000000 != 0:
+                            if flags & 0x8 != 0:
+                                ed.undo_to_begin_point()
+                            else:
+                                ed.undo_point()
+                        else:
+                            if flags & 0x8 != 0:
+                                ed.redo_to_end_point()
+                            else:
+                                ed.redo_point()
+
+            key_events = [ ev for ev, in io.get_key_events(wnd_name) ]
+            for key in key_events:         
+                if key == ord('q') or key == ord('z'):
+                    if len(done_paths) > 0:
+                        image_paths.insert(0, filepath)                        
+                        filepath = done_paths.pop(-1)
+                        
+                        if filepath.parent != input_path:
+                            new_filename_path = input_path / filepath.name
+                            filepath.rename ( new_filename_path )                            
+                            image_paths.insert(0, new_filename_path)   
+                        else:
+                            image_paths.insert(0, filepath)   
+                            
+                        next = True
+                        break
+                elif filepath is not None and ( key == ord('e') or key == ord('c') ):
+                    ed.mask_finish()
+                    dflimg.embed_and_set (str(filepath), ie_polys=ed.get_ie_polys() )
+                    
+                    if key == ord('e'):
+                        new_filename_path = confirmed_path / filepath.name
+                        filepath.rename(new_filename_path)
+                        done_paths += [new_filename_path]
+                    else:
+                        done_paths += [filepath]                         
+                    
+                    next = True
+                    break
+                
+                elif filepath is not None and ( key == ord('w') or key == ord('x') ):           
+                    if key == ord('w'):
+                        new_filename_path = skipped_path / filepath.name
+                        filepath.rename(new_filename_path)
+                        done_paths += [new_filename_path]
+                    else:
+                        done_paths += [filepath]     
+                                       
+                    next = True
+                    break
+                elif key == 27: #esc
+                    is_exit = True
+                    next = True
+                    break
+            screen = ed.make_screen()
+
+            io.show_image (wnd_name, screen )
+        io.process_messages(0.005)
+        
+    io.destroy_all_windows()    
+

mainscripts/Trainer.py

@@ -253,7 +253,7 @@ def main(args, device_args):
                 for i in range(0, len(head_lines)):
                     t = i*head_line_height
                     b = (i+1)*head_line_height
-                    head[t:b, 0:w] += imagelib.get_text_image (  (w,head_line_height,c) , head_lines[i], color=[0.8]*c )
+                    head[t:b, 0:w] += imagelib.get_text_image (  (head_line_height,w,c) , head_lines[i], color=[0.8]*c )
 
                 final = head
 

mainscripts/Util.py

@@ -27,6 +27,7 @@ def convert_png_to_jpg_file (filepath):
     DFLJPG.embed_data( new_filepath,
                        face_type=dfl_dict.get('face_type', None),
                        landmarks=dfl_dict.get('landmarks', None),
+                       ie_polys=dfl_dict.get('ie_polys', None),
                        source_filename=dfl_dict.get('source_filename', None),
                        source_rect=dfl_dict.get('source_rect', None),
                        source_landmarks=dfl_dict.get('source_landmarks', None) )
@@ -63,7 +64,7 @@ def add_landmarks_debug_images(input_path):
 
         if img is not None:
             face_landmarks = dflimg.get_landmarks()
-            LandmarksProcessor.draw_landmarks(img, face_landmarks, transparent_mask=True)
+            LandmarksProcessor.draw_landmarks(img, face_landmarks, transparent_mask=True, ie_polys=dflimg.get_ie_polys() )
 
             output_file = '{}{}'.format( str(Path(str(input_path)) / filepath.stem),  '_debug.jpg')
             cv2_imwrite(output_file, img, [int(cv2.IMWRITE_JPEG_QUALITY), 50] )

models/ModelBase.py

@@ -496,5 +496,5 @@ class ModelBase(object):
 
         lh_text = 'Iter: %d' % (iter) if iter != 0 else ''
 
-        lh_img[last_line_t:last_line_b, 0:w] += imagelib.get_text_image (  (w,last_line_b-last_line_t,c), lh_text, color=[0.8]*c )
+        lh_img[last_line_t:last_line_b, 0:w] += imagelib.get_text_image (  (last_line_b-last_line_t,w,c), lh_text, color=[0.8]*c )
         return lh_img

models/Model_DF/Model.py

@@ -12,7 +12,7 @@ class Model(ModelBase):
     def onInitializeOptions(self, is_first_run, ask_override):
         if is_first_run or ask_override:
             def_pixel_loss = self.options.get('pixel_loss', False)
-            self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k iters to enhance fine details and decrease face jitter.")
+            self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.")
         else:
             self.options['pixel_loss'] = self.options.get('pixel_loss', False)
 

models/Model_H128/Model.py

@@ -20,7 +20,7 @@ class Model(ModelBase):
 
         if is_first_run or ask_override:
             def_pixel_loss = self.options.get('pixel_loss', False)
-            self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k iters to enhance fine details and decrease face jitter.")
+            self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.")
         else:
             self.options['pixel_loss'] = self.options.get('pixel_loss', False)
 

models/Model_H64/Model.py

@@ -20,7 +20,7 @@ class Model(ModelBase):
 
         if is_first_run or ask_override:
             def_pixel_loss = self.options.get('pixel_loss', False)
-            self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k iters to enhance fine details and decrease face jitter.")
+            self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.")
         else:
             self.options['pixel_loss'] = self.options.get('pixel_loss', False)
 

models/Model_LIAEF128/Model.py

@@ -12,7 +12,7 @@ class Model(ModelBase):
     def onInitializeOptions(self, is_first_run, ask_override):
         if is_first_run or ask_override:
             def_pixel_loss = self.options.get('pixel_loss', False)
-            self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k iters to enhance fine details and decrease face jitter.")
+            self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.")
         else:
             self.options['pixel_loss'] = self.options.get('pixel_loss', False)
 

models/Model_SAE/Model.py

@@ -63,13 +63,11 @@ class SAEModel(ModelBase):
             self.options['e_ch_dims'] = np.clip ( io.input_int("Encoder dims per channel (21-85 ?:help skip:%d) : " % (default_e_ch_dims) , default_e_ch_dims, help_message="More encoder dims help to recognize more facial features, but require more VRAM. You can fine-tune model size to fit your GPU." ), 21, 85 )
             default_d_ch_dims = self.options['e_ch_dims'] // 2
             self.options['d_ch_dims'] = np.clip ( io.input_int("Decoder dims per channel (10-85 ?:help skip:%d) : " % (default_d_ch_dims) , default_d_ch_dims, help_message="More decoder dims help to get better details, but require more VRAM. You can fine-tune model size to fit your GPU." ), 10, 85 )
-            self.options['d_residual_blocks'] = io.input_bool ("Add residual blocks to decoder? (y/n, ?:help skip:n) : ", False, help_message="These blocks help to get better details, but require more computing time.")
             self.options['remove_gray_border'] = io.input_bool ("Remove gray border? (y/n, ?:help skip:n) : ", False, help_message="Removes gray border of predicted face, but requires more computing resources.")
         else:
             self.options['ae_dims'] = self.options.get('ae_dims', default_ae_dims)
             self.options['e_ch_dims'] = self.options.get('e_ch_dims', default_e_ch_dims)
             self.options['d_ch_dims'] = self.options.get('d_ch_dims', default_d_ch_dims)
-            self.options['d_residual_blocks'] = self.options.get('d_residual_blocks', False)
             self.options['remove_gray_border'] = self.options.get('remove_gray_border', False)
 
         if is_first_run:
@@ -81,7 +79,7 @@ class SAEModel(ModelBase):
         default_bg_style_power = 0.0
         if is_first_run or ask_override:
             def_pixel_loss = self.options.get('pixel_loss', False)
-            self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: %s ) : " % (yn_str[def_pixel_loss]), def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 60k iters to enhance fine details. Warning: this option may cause collapse the model, make a backup of Model folder before apply it.")
+            self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: %s ) : " % (yn_str[def_pixel_loss]), def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time.")
 
             default_face_style_power = default_face_style_power if is_first_run else self.options.get('face_style_power', default_face_style_power)
             self.options['face_style_power'] = np.clip ( io.input_number("Face style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_face_style_power), default_face_style_power,
@@ -105,7 +103,7 @@ class SAEModel(ModelBase):
         ae_dims = self.options['ae_dims']
         e_ch_dims = self.options['e_ch_dims']
         d_ch_dims = self.options['d_ch_dims']
-        d_residual_blocks = self.options['d_residual_blocks']
+        d_residual_blocks = True
         bgr_shape = (resolution, resolution, 3)
         mask_shape = (resolution, resolution, 1)
 
@@ -127,7 +125,9 @@ class SAEModel(ModelBase):
         target_dstm_ar = [ Input ( ( mask_shape[0] // (2**i) ,)*2 + (mask_shape[-1],) ) for i in range(ms_count-1, -1, -1)]
 
         padding = 'reflect' if self.options['remove_gray_border'] else 'zero'
-        common_flow_kwargs = { 'padding': padding }
+        common_flow_kwargs = { 'padding': padding,
+                               'norm': 'bn',
+                               'act':'prelu' }
 
         weights_to_load = []
         if self.options['archi'] == 'liae':
@@ -302,7 +302,7 @@ class SAEModel(ModelBase):
                 self.src_dst_mask_train = K.function (feed,[src_mask_loss, dst_mask_loss], self.src_dst_mask_opt.get_updates(src_mask_loss+dst_mask_loss, src_dst_mask_loss_train_weights) )
 
             if self.options['learn_mask']:
-                self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_src_dst[-1], pred_src_dstm[-1]])
+                self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_dst_dstm[-1], pred_src_dst[-1], pred_src_dstm[-1]])
             else:
                 self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_src_dst[-1] ] )
 
@@ -310,7 +310,7 @@ class SAEModel(ModelBase):
         else:
             self.load_weights_safe(weights_to_load)
             if self.options['learn_mask']:
-                self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1], pred_src_dstm[-1] ])
+                self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1], pred_dst_dstm[-1], pred_src_dstm[-1] ])
             else:
                 self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1] ])
 
@@ -391,24 +391,34 @@ class SAEModel(ModelBase):
         test_B_m = sample[1][2][0:4]
 
         if self.options['learn_mask']:
-            S, D, SS, DD, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([test_A,test_B] + self.AE_view ([test_A, test_B]) ) ]
+            S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([test_A,test_B] + self.AE_view ([test_A, test_B]) ) ]
-            SDM, = [ np.repeat (x, (3,), -1) for x in [SDM] ]
+            DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ]
         else:
             S, D, SS, DD, SD, = [ np.clip(x, 0.0, 1.0) for x in ([test_A,test_B] + self.AE_view ([test_A, test_B]) ) ]
 
+        result = []
         st = []
         for i in range(0, len(test_A)):
             ar = S[i], SS[i], D[i], DD[i], SD[i]
-            #if self.options['learn_mask']:
-            #    ar += (SDM[i],)
             st.append ( np.concatenate ( ar, axis=1) )
-
+        
-        return [ ('SAE', np.concatenate (st, axis=0 )), ]
+        result += [ ('SAE', np.concatenate (st, axis=0 )), ]
+        
+        if self.options['learn_mask']:
+            st_m = []
+            for i in range(0, len(test_A)):
+                ar = S[i], SS[i], D[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i])
+                st_m.append ( np.concatenate ( ar, axis=1) )
+                
+            result += [ ('SAE masked', np.concatenate (st_m, axis=0 )), ]
+            
+        return result
 
     def predictor_func (self, face):
         if self.options['learn_mask']:
-            bgr, mask = self.AE_convert ([face[np.newaxis,...]])
+            bgr, mask_dst_dstm, mask_src_dstm = self.AE_convert ([face[np.newaxis,...]])
-            return bgr[0], mask[0][...,0]
+            mask = mask_dst_dstm[0] * mask_src_dstm[0]            
+            return bgr[0], mask[...,0]
         else:
             bgr, = self.AE_convert ([face[np.newaxis,...]])
             return bgr[0]
@@ -440,23 +450,39 @@ class SAEModel(ModelBase):
     def initialize_nn_functions():
         exec (nnlib.import_all(), locals(), globals())
 
-        def BatchNorm():
+        def NormPass(x):
-            return BatchNormalization(axis=-1)
+            return x
+            
+        def Norm(norm=''):
+            if norm == 'bn':
+                return BatchNormalization(axis=-1)
+            else:
+                return NormPass
+        
+        def Act(act='', lrelu_alpha=0.1):
+            if act == 'prelu':
+                return PReLU()
+            else:
+                return LeakyReLU(alpha=lrelu_alpha)
 
         class ResidualBlock(object):
-            def __init__(self, filters, kernel_size=3, padding='zero', use_reflection_padding=False):
+            def __init__(self, filters, kernel_size=3, padding='zero', use_reflection_padding=False, norm='', act='', **kwargs):
                 self.filters = filters
                 self.kernel_size = kernel_size
                 self.padding = padding
+                self.norm = norm
+                self.act = act
 
             def __call__(self, inp):
-                var_x = inp
+                x = inp
-                var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding)(var_x)
+                x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding)(x)
-                var_x = LeakyReLU(alpha=0.2)(var_x)
+                x = Act(self.act, lrelu_alpha=0.2)(x)
-                var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding)(var_x)
+                x = Norm(self.norm)(x)
-                var_x = Add()([var_x, inp])
+                x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding)(x)
-                var_x = LeakyReLU(alpha=0.2)(var_x)
+                x = Add()([x, inp])
-                return var_x
+                x = Act(self.act, lrelu_alpha=0.2)(x)
+                x = Norm(self.norm)(x)
+                return x
         SAEModel.ResidualBlock = ResidualBlock
 
         def ResidualBlock_pre (**base_kwargs):
@@ -466,9 +492,9 @@ class SAEModel(ModelBase):
             return func
         SAEModel.ResidualBlock_pre = ResidualBlock_pre
 
-        def downscale (dim, padding='zero'):
+        def downscale (dim, padding='zero', norm='', act='', **kwargs):
             def func(x):
-                return LeakyReLU(0.1)(Conv2D(dim, kernel_size=5, strides=2, padding=padding)(x))
+                return Norm(norm)( Act(act) (Conv2D(dim, kernel_size=5, strides=2, padding=padding)(x)) )
             return func
         SAEModel.downscale = downscale
 
@@ -479,9 +505,9 @@ class SAEModel(ModelBase):
             return func
         SAEModel.downscale_pre = downscale_pre
 
-        def upscale (dim, padding='zero'):
+        def upscale (dim, padding='zero', norm='', act='', **kwargs):
             def func(x):
-                return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding=padding)(x)))
+                return SubpixelUpscaler()(Norm(norm)(Act(act)(Conv2D(dim * 4, kernel_size=3, strides=1, padding=padding)(x))))
             return func
         SAEModel.upscale = upscale
 
@@ -492,7 +518,7 @@ class SAEModel(ModelBase):
             return func
         SAEModel.upscale_pre = upscale_pre
 
-        def to_bgr (output_nc, padding='zero'):
+        def to_bgr (output_nc, padding='zero', **kwargs):
             def func(x):
                 return Conv2D(output_nc, kernel_size=5, padding=padding, activation='sigmoid')(x)
             return func
@@ -506,10 +532,10 @@ class SAEModel(ModelBase):
         SAEModel.to_bgr_pre = to_bgr_pre
 
     @staticmethod
-    def LIAEEncFlow(resolution, ch_dims, padding='zero', **kwargs):
+    def LIAEEncFlow(resolution, ch_dims, **kwargs):
         exec (nnlib.import_all(), locals(), globals())
-        upscale = SAEModel.upscale_pre(padding=padding)
+        upscale = SAEModel.upscale_pre(**kwargs)
-        downscale = SAEModel.downscale_pre(padding=padding)
+        downscale = SAEModel.downscale_pre(**kwargs)
 
         def func(input):
             dims = K.int_shape(input)[-1]*ch_dims
@@ -525,9 +551,9 @@ class SAEModel(ModelBase):
         return func
 
     @staticmethod
-    def LIAEInterFlow(resolution, ae_dims=256, padding='zero', **kwargs):
+    def LIAEInterFlow(resolution, ae_dims=256, **kwargs):
         exec (nnlib.import_all(), locals(), globals())
-        upscale = SAEModel.upscale_pre(padding=padding)
+        upscale = SAEModel.upscale_pre(**kwargs)
         lowest_dense_res=resolution // 16
 
         def func(input):
@@ -540,12 +566,12 @@ class SAEModel(ModelBase):
         return func
 
     @staticmethod
-    def LIAEDecFlow(output_nc,ch_dims, multiscale_count=1, add_residual_blocks=False, padding='zero', **kwargs):
+    def LIAEDecFlow(output_nc,ch_dims, multiscale_count=1, add_residual_blocks=False, padding='zero', norm='', **kwargs):
         exec (nnlib.import_all(), locals(), globals())
-        upscale = SAEModel.upscale_pre(padding=padding)
+        upscale = SAEModel.upscale_pre(**kwargs)
-        to_bgr = SAEModel.to_bgr_pre(padding=padding)
+        to_bgr = SAEModel.to_bgr_pre(**kwargs)
         dims = output_nc * ch_dims
-        ResidualBlock = SAEModel.ResidualBlock_pre(padding=padding)
+        ResidualBlock = SAEModel.ResidualBlock_pre(**kwargs)
 
         def func(input):
             x = input[0]

samplelib/Sample.py

@@ -16,24 +16,26 @@ class SampleType(IntEnum):
     QTY = 5
 
 class Sample(object):
-    def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, pitch=None, yaw=None, mirror=None, close_target_list=None):
+    def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, mirror=None, close_target_list=None):
         self.sample_type = sample_type if sample_type is not None else SampleType.IMAGE
         self.filename = filename
         self.face_type = face_type
         self.shape = shape
         self.landmarks = np.array(landmarks) if landmarks is not None else None
+        self.ie_polys = ie_polys
         self.pitch = pitch
         self.yaw = yaw
         self.mirror = mirror
         self.close_target_list = close_target_list
 
-    def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, pitch=None, yaw=None, mirror=None, close_target_list=None):
+    def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, mirror=None, close_target_list=None):
         return Sample(
             sample_type=sample_type if sample_type is not None else self.sample_type,
             filename=filename if filename is not None else self.filename,
             face_type=face_type if face_type is not None else self.face_type,
             shape=shape if shape is not None else self.shape,
             landmarks=landmarks if landmarks is not None else self.landmarks.copy(),
+            ie_polys=ie_polys if ie_polys is not None else self.ie_polys,
             pitch=pitch if pitch is not None else self.pitch,
             yaw=yaw if yaw is not None else self.yaw,
             mirror=mirror if mirror is not None else self.mirror,

samplelib/SampleLoader.py

@@ -75,6 +75,7 @@ class SampleLoader:
                                                    face_type=FaceType.fromString (dflimg.get_face_type()),
                                                    shape=dflimg.get_shape(),
                                                    landmarks=dflimg.get_landmarks(),
+                                                   ie_polys=dflimg.get_ie_polys(),
                                                    pitch=pitch,
                                                    yaw=yaw) )
             except:

samplelib/SampleProcessor.py

@@ -152,7 +152,7 @@ class SampleProcessor(object):
 
                     if is_face_sample:
                         if face_mask_type == 1:
-                            img = np.concatenate( (img, LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks) ), -1 )
+                            img = np.concatenate( (img, LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks, cur_sample.ie_polys) ), -1 )
                         elif face_mask_type == 2:
                             mask = LandmarksProcessor.get_image_eye_mask (img.shape, cur_sample.landmarks)
                             mask = np.expand_dims (cv2.blur (mask, ( w // 32, w // 32 ) ), -1)

utils/DFLJPG.py

@@ -2,6 +2,7 @@ import struct
 import pickle
 import numpy as np
 from facelib import FaceType
+from imagelib import IEPolys
 from utils.struct_utils import *
 
 class DFLJPG(object):
@@ -18,7 +19,7 @@ class DFLJPG(object):
             with open(filename, "rb") as f:
                 data = f.read()
         except:
-            raise FileNotFoundError(data)
+            raise FileNotFoundError(filename)
 
         try:
             inst = DFLJPG()
@@ -150,6 +151,7 @@ class DFLJPG(object):
     @staticmethod
     def embed_data(filename, face_type=None,
                              landmarks=None,
+                             ie_polys=None,
                              source_filename=None,
                              source_rect=None,
                              source_landmarks=None,
@@ -160,6 +162,7 @@ class DFLJPG(object):
         inst.setDFLDictData ({
                                 'face_type': face_type,
                                 'landmarks': landmarks,
+                                'ie_polys' : ie_polys.dump() if ie_polys is not None else None,
                                 'source_filename': source_filename,
                                 'source_rect': source_rect,
                                 'source_landmarks': source_landmarks,
@@ -172,6 +175,29 @@ class DFLJPG(object):
         except:
             raise Exception( 'cannot save %s' % (filename) )
 
+    def embed_and_set(self, filename, face_type=None,
+                                landmarks=None,
+                                ie_polys=None,
+                                source_filename=None,
+                                source_rect=None,
+                                source_landmarks=None,
+                                image_to_face_mat=None
+                    ):
+        if face_type is None: face_type = self.get_face_type()
+        if landmarks is None: landmarks = self.get_landmarks()
+        if ie_polys is None: ie_polys = self.get_ie_polys()
+        if source_filename is None: source_filename = self.get_source_filename()
+        if source_rect is None: source_rect = self.get_source_rect()
+        if source_landmarks is None: source_landmarks = self.get_source_landmarks()
+        if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat()
+        DFLJPG.embed_data (filename, face_type=face_type,
+                                     landmarks=landmarks,
+                                     ie_polys=ie_polys,
+                                     source_filename=source_filename,
+                                     source_rect=source_rect,
+                                     source_landmarks=source_landmarks,
+                                     image_to_face_mat=image_to_face_mat)
+        
     def dump(self):
         data = b""
 
@@ -222,6 +248,12 @@ class DFLJPG(object):
 
     def get_face_type(self): return self.dfl_dict['face_type']
     def get_landmarks(self): return np.array ( self.dfl_dict['landmarks'] )
+    def get_ie_polys(self): return IEPolys.load(self.dfl_dict.get('ie_polys',None))
     def get_source_filename(self): return self.dfl_dict['source_filename']
     def get_source_rect(self): return self.dfl_dict['source_rect']
     def get_source_landmarks(self): return np.array ( self.dfl_dict['source_landmarks'] )
+    def get_image_to_face_mat(self): 
+        mat = self.dfl_dict.get ('image_to_face_mat', None)
+        if mat is not None:
+            return np.array (mat)
+        return None

utils/DFLPNG.py

@@ -6,6 +6,7 @@ import zlib
 import pickle
 import numpy as np
 from facelib import FaceType
+from imagelib import IEPolys
 
 class Chunk(object):
     def __init__(self, name=None, data=None):
@@ -226,7 +227,7 @@ class DFLPNG(object):
             with open(filename, "rb") as f:
                 data = f.read()
         except:
-            raise FileNotFoundError(data)
+            raise FileNotFoundError(filename)
 
         inst = DFLPNG()
         inst.data = data
@@ -267,18 +268,22 @@ class DFLPNG(object):
     @staticmethod
     def embed_data(filename, face_type=None,
                              landmarks=None,
+                             ie_polys=None,
                              source_filename=None,
                              source_rect=None,
-                             source_landmarks=None
+                             source_landmarks=None,
+                             image_to_face_mat=None
                    ):
 
         inst = DFLPNG.load_raw (filename)
         inst.setDFLDictData ({
                                 'face_type': face_type,
                                 'landmarks': landmarks,
+                                'ie_polys' : ie_polys.dump() if ie_polys is not None else None,
                                 'source_filename': source_filename,
                                 'source_rect': source_rect,
-                                'source_landmarks': source_landmarks
+                                'source_landmarks': source_landmarks,
+                                'image_to_face_mat':image_to_face_mat
                              })
 
         try:
@@ -287,6 +292,29 @@ class DFLPNG(object):
         except:
             raise Exception( 'cannot save %s' % (filename) )
 
+    def embed_and_set(self, filename, face_type=None,
+                                landmarks=None,
+                                ie_polys=None,
+                                source_filename=None,
+                                source_rect=None,
+                                source_landmarks=None,
+                                image_to_face_mat=None
+                        ):
+        if face_type is None: face_type = self.get_face_type()
+        if landmarks is None: landmarks = self.get_landmarks()
+        if ie_polys is None: ie_polys = self.get_ie_polys()
+        if source_filename is None: source_filename = self.get_source_filename()
+        if source_rect is None: source_rect = self.get_source_rect()
+        if source_landmarks is None: source_landmarks = self.get_source_landmarks()
+        if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat()
+        DFLPNG.embed_data (filename, face_type=face_type,
+                                     landmarks=landmarks,
+                                     ie_polys=ie_polys,
+                                     source_filename=source_filename,
+                                     source_rect=source_rect,
+                                     source_landmarks=source_landmarks,
+                                     image_to_face_mat=image_to_face_mat)
+
     def dump(self):
         data = PNG_HEADER
         for chunk in self.chunks:
@@ -326,9 +354,15 @@ class DFLPNG(object):
 
     def get_face_type(self): return self.fcwp_dict['face_type']
     def get_landmarks(self): return np.array ( self.fcwp_dict['landmarks'] )
+    def get_ie_polys(self): return IEPolys.load(self.fcwp_dict.get('ie_polys',None))
     def get_source_filename(self): return self.fcwp_dict['source_filename']
     def get_source_rect(self): return self.fcwp_dict['source_rect']
     def get_source_landmarks(self): return np.array ( self.fcwp_dict['source_landmarks'] )
+    def get_image_to_face_mat(self):
+        mat = self.fcwp_dict.get ('image_to_face_mat', None)
+        if mat is not None:
+            return np.array (mat)
+        return None
 
     def __str__(self):
         return "<PNG length={length} chunks={}>".format(len(self.chunks), **self.__dict__)