removing trailing spaces

converters/Converter.py

@@ -7,25 +7,25 @@ class Converter(object):
     TYPE_FACE = 0                   #calls convert_face
     TYPE_IMAGE = 1                  #calls convert_image without landmarks
     TYPE_IMAGE_WITH_LANDMARKS = 2   #calls convert_image with landmarks
-    
+
     #overridable
     def __init__(self, predictor_func, type):
         self.predictor_func = predictor_func
         self.type = type
-        
+
     #overridable
     def convert_face (self, img_bgr, img_face_landmarks, debug):
-        #return float32 image        
+        #return float32 image
         #if debug , return tuple ( images of any size and channels, ...)
         return image
-        
+
     #overridable
     def convert_image (self, img_bgr, img_landmarks, debug):
         #img_landmarks not None, if input image is png with embedded data
-        #return float32 image        
+        #return float32 image
         #if debug , return tuple ( images of any size and channels, ...)
         return image
-        
+
     #overridable
     def dummy_predict(self):
         #do dummy predict here
@@ -33,8 +33,8 @@ class Converter(object):
 
     def copy(self):
         return copy.copy(self)
-        
+
     def copy_and_set_predictor(self, predictor_func):
         result = self.copy()
         result.predictor_func = predictor_func
-        return result
+        return result

converters/ConverterImage.py

@@ -7,7 +7,7 @@ import numpy as np
 from utils import image_utils
 
 '''
-predictor_func: 
+predictor_func:
     input:  [predictor_input_size, predictor_input_size, BGR]
     output: [predictor_input_size, predictor_input_size, BGR]
 '''
@@ -16,18 +16,18 @@ class ConverterImage(Converter):
 
     #override
     def __init__(self,  predictor_func,
-                        predictor_input_size=0, 
+                        predictor_input_size=0,
                         output_size=0):
-                        
+
         super().__init__(predictor_func, Converter.TYPE_IMAGE)
-         
+
         self.predictor_input_size = predictor_input_size
-        self.output_size = output_size   
-        
+        self.output_size = output_size
+
     #override
     def dummy_predict(self):
         self.predictor_func ( np.zeros ( (self.predictor_input_size, self.predictor_input_size,3), dtype=np.float32) )
-        
+
     #override
     def convert_image (self, img_bgr, img_landmarks, debug):
         img_size = img_bgr.shape[1], img_bgr.shape[0]

facelib/DLIBExtractor.py

@@ -4,36 +4,36 @@ import cv2
 
 from pathlib import Path
 
-class DLIBExtractor(object):   
+class DLIBExtractor(object):
     def __init__(self, dlib):
-        self.scale_to = 1850 
-        #3100 eats ~1.687GB VRAM on 2GB 730 desktop card, but >4Gb on 6GB card, 
+        self.scale_to = 1850
+        #3100 eats ~1.687GB VRAM on 2GB 730 desktop card, but >4Gb on 6GB card,
         #but 3100 doesnt work on 2GB 850M notebook card, I cant understand this behaviour
         #1850 works on 2GB 850M notebook card, works faster than 3100, produces good result
         self.dlib = dlib
 
-    def __enter__(self):   
+    def __enter__(self):
         self.dlib_cnn_face_detector = self.dlib.cnn_face_detection_model_v1( str(Path(__file__).parent / "mmod_human_face_detector.dat") )
-        self.dlib_cnn_face_detector ( np.zeros ( (self.scale_to, self.scale_to, 3), dtype=np.uint8), 0 ) 
+        self.dlib_cnn_face_detector ( np.zeros ( (self.scale_to, self.scale_to, 3), dtype=np.uint8), 0 )
         return self
-        
+
     def __exit__(self, exc_type=None, exc_value=None, traceback=None):
         del self.dlib_cnn_face_detector
         return False #pass exception between __enter__ and __exit__ to outter level
-        
+
     def extract_from_bgr (self, input_image):
         input_image = input_image[:,:,::-1].copy()
         (h, w, ch) = input_image.shape
 
-        detected_faces = []    
+        detected_faces = []
         input_scale = self.scale_to / (w if w > h else h)
         input_image = cv2.resize (input_image, ( int(w*input_scale), int(h*input_scale) ), interpolation=cv2.INTER_LINEAR)
         detected_faces = self.dlib_cnn_face_detector(input_image, 0)
 
-        result = []        
+        result = []
         for d_rect in detected_faces:
             if type(d_rect) == self.dlib.mmod_rectangle:
-                d_rect = d_rect.rect            
+                d_rect = d_rect.rect
             left, top, right, bottom = d_rect.left(), d_rect.top(), d_rect.right(), d_rect.bottom()
             result.append ( (int(left/input_scale), int(top/input_scale), int(right/input_scale), int(bottom/input_scale)) )
 

facelib/FANSegmentator.py

@@ -8,16 +8,16 @@ from interact import interact as io
 class FANSegmentator(object):
     def __init__ (self, resolution, face_type_str, load_weights=True, weights_file_root=None):
         exec( nnlib.import_all(), locals(), globals() )
-        
+
         self.model = FANSegmentator.BuildModel(resolution, ngf=32)
-        
+
         if weights_file_root:
             weights_file_root = Path(weights_file_root)
         else:
             weights_file_root = Path(__file__).parent
-            
+
         self.weights_path = weights_file_root / ('FANSeg_%d_%s.h5' % (resolution, face_type_str) )
-        
+
         if load_weights:
             self.model.load_weights (str(self.weights_path))
         else:
@@ -31,19 +31,19 @@ class FANSegmentator(object):
 
     def __enter__(self):
         return self
-        
+
     def __exit__(self, exc_type=None, exc_value=None, traceback=None):
         return False #pass exception between __enter__ and __exit__ to outter level
-        
+
     def save_weights(self):
         self.model.save_weights (str(self.weights_path))
-        
+
     def train_on_batch(self, inp, outp):
         return self.model.train_on_batch(inp, outp)
-        
+
     def extract_from_bgr (self, input_image):
         return np.clip ( (self.model.predict(input_image) + 1) / 2.0, 0, 1.0 )
-        
+
     @staticmethod
     def BuildModel ( resolution, ngf=64):
         exec( nnlib.import_all(), locals(), globals() )
@@ -53,7 +53,7 @@ class FANSegmentator(object):
         x = FANSegmentator.DecFlow(ngf=ngf)(x)
         model = Model(inp,x)
         return model
-        
+
     @staticmethod
     def EncFlow(ngf=64, num_downs=4):
         exec( nnlib.import_all(), locals(), globals() )
@@ -65,19 +65,19 @@ class FANSegmentator(object):
         def downscale (dim):
             def func(x):
                 return LeakyReLU(0.1)(XNormalization(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x)))
-            return func 
-            
-        def func(input):     
+            return func
+
+        def func(input):
             x = input
-            
+
             result = []
             for i in range(num_downs):
                 x = downscale ( min(ngf*(2**i), ngf*8) )(x)
-                result += [x]                
-                
+                result += [x]
+
             return result
         return func
-        
+
     @staticmethod
     def DecFlow(output_nc=1, ngf=64, activation='tanh'):
         exec (nnlib.import_all(), locals(), globals())
@@ -85,23 +85,23 @@ class FANSegmentator(object):
         use_bias = True
         def XNormalization(x):
             return InstanceNormalization (axis=3, gamma_initializer=RandomNormal(1., 0.02))(x)
-            
+
         def Conv2D (filters, kernel_size, strides=(1, 1), padding='valid', data_format=None, dilation_rate=(1, 1), activation=None, use_bias=use_bias, kernel_initializer=RandomNormal(0, 0.02), bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None):
             return keras.layers.Conv2D( filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, data_format=data_format, dilation_rate=dilation_rate, activation=activation, use_bias=use_bias, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer, activity_regularizer=activity_regularizer, kernel_constraint=kernel_constraint, bias_constraint=bias_constraint )
-        
+
         def upscale (dim):
             def func(x):
                 return SubpixelUpscaler()( LeakyReLU(0.1)(XNormalization(Conv2D(dim, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x))))
-            return func 
-            
+            return func
+
         def func(input):
             input_len = len(input)
-            
+
             x = input[input_len-1]
-            for i in range(input_len-1, -1, -1):          
+            for i in range(input_len-1, -1, -1):
                 x = upscale( min(ngf* (2**i) *4, ngf*8 *4 ) )(x)
                 if i != 0:
                     x = Concatenate(axis=3)([ input[i-1] , x])
-        
+
             return Conv2D(output_nc, 3, 1, 'same', activation=activation)(x)
-        return func
+        return func

facelib/FaceType.py

@@ -3,7 +3,7 @@ from enum import IntEnum
 class FaceType(IntEnum):
     HALF = 0,
     FULL = 1,
-    HEAD = 2,    
+    HEAD = 2,
     AVATAR = 3,    #centered nose only
     MARK_ONLY = 4, #no align at all, just embedded faceinfo
     QTY = 5
@@ -13,12 +13,12 @@ class FaceType(IntEnum):
         r = from_string_dict.get (s.lower())
         if r is None:
             raise Exception ('FaceType.fromString value error')
-        return r            
-            
-    @staticmethod        
+        return r
+
+    @staticmethod
     def toString (face_type):
         return to_string_list[face_type]
-        
+
 from_string_dict = {'half_face': FaceType.HALF,
                     'full_face': FaceType.FULL,
                     'head' : FaceType.HEAD,
@@ -29,6 +29,5 @@ to_string_list = [ 'half_face',
                    'full_face',
                    'head',
                    'avatar',
-                   'mark_only'   
+                   'mark_only'
                     ]
-                        

facelib/LandmarksExtractor.py

@@ -10,38 +10,38 @@ class LandmarksExtractor(object):
     def __init__ (self, keras):
         self.keras = keras
         K = self.keras.backend
-  
-    def __enter__(self):        
+
+    def __enter__(self):
         keras_model_path = Path(__file__).parent / "2DFAN-4.h5"
         if not keras_model_path.exists():
             return None
 
-        self.keras_model = self.keras.models.load_model (str(keras_model_path)) 
-    
+        self.keras_model = self.keras.models.load_model (str(keras_model_path))
+
         return self
-        
+
     def __exit__(self, exc_type=None, exc_value=None, traceback=None):
         del self.keras_model
         return False #pass exception between __enter__ and __exit__ to outter level
-        
+
     def extract_from_bgr (self, input_image, rects, second_pass_extractor=None):
         input_image = input_image[:,:,::-1].copy()
         (h, w, ch) = input_image.shape
-        
+
         landmarks = []
         for (left, top, right, bottom) in rects:
             try:
                 center = np.array( [ (left + right) / 2.0, (top + bottom) / 2.0] )
                 #center[1] -= (bottom - top) * 0.12
                 scale = (right - left + bottom - top) / 195.0
-            
+
                 image = self.crop(input_image, center, scale).astype(np.float32)
                 image = np.expand_dims(image, 0)
-                
+
                 predicted = self.keras_model.predict (image).transpose (0,3,1,2)
-                
+
                 pts_img = self.get_pts_from_predict ( predicted[-1], center, scale)
-                pts_img = [ ( int(pt[0]), int(pt[1]) ) for pt in pts_img ]             
+                pts_img = [ ( int(pt[0]), int(pt[1]) ) for pt in pts_img ]
                 landmarks.append ( ( (left, top, right, bottom),pts_img ) )
             except Exception as e:
                 landmarks.append ( ( (left, top, right, bottom), None ) )
@@ -52,26 +52,26 @@ class LandmarksExtractor(object):
                     rect, lmrks = landmarks[i]
                     if lmrks is None:
                         continue
-                        
+
                     image_to_face_mat = LandmarksProcessor.get_transform_mat (lmrks, 256, FaceType.FULL)
                     face_image = cv2.warpAffine(input_image, image_to_face_mat, (256, 256), cv2.INTER_CUBIC)
-                    
+
                     rects2 = second_pass_extractor.extract_from_bgr(face_image)
                     if len(rects2) != 1: #dont do second pass if more than 1 or zero faces detected in cropped image
                         continue
-                        
+
                     rect2 = rects2[0]
-                    
+
                     lmrks2 = self.extract_from_bgr (face_image, [rect2] )[0][1]
-                    source_lmrks2 = LandmarksProcessor.transform_points (lmrks2, image_to_face_mat, True)                            
+                    source_lmrks2 = LandmarksProcessor.transform_points (lmrks2, image_to_face_mat, True)
                     landmarks[i] = (rect, source_lmrks2)
                 except:
                     continue
-                
+
         return landmarks
 
     def transform(self, point, center, scale, resolution):
-        pt = np.array ( [point[0], point[1], 1.0] )            
+        pt = np.array ( [point[0], point[1], 1.0] )
         h = 200.0 * scale
         m = np.eye(3)
         m[0,0] = resolution / h
@@ -80,11 +80,11 @@ class LandmarksExtractor(object):
         m[1,2] = resolution * ( -center[1] / h + 0.5 )
         m = np.linalg.inv(m)
         return np.matmul (m, pt)[0:2]
-        
+
     def crop(self, image, center, scale, resolution=256.0):
         ul = self.transform([1, 1], center, scale, resolution).astype( np.int )
         br = self.transform([resolution, resolution], center, scale, resolution).astype( np.int )
-        
+
         if image.ndim > 2:
             newDim = np.array([br[1] - ul[1], br[0] - ul[0], image.shape[2]], dtype=np.int32)
             newImg = np.zeros(newDim, dtype=np.uint8)
@@ -98,14 +98,14 @@ class LandmarksExtractor(object):
         oldX = np.array([max(1, ul[0] + 1), min(br[0], wd)], dtype=np.int32)
         oldY = np.array([max(1, ul[1] + 1), min(br[1], ht)], dtype=np.int32)
         newImg[newY[0] - 1:newY[1], newX[0] - 1:newX[1] ] = image[oldY[0] - 1:oldY[1], oldX[0] - 1:oldX[1], :]
-        
+
         newImg = cv2.resize(newImg, dsize=(int(resolution), int(resolution)), interpolation=cv2.INTER_LINEAR)
         return newImg
-               
+
     def get_pts_from_predict(self, a, center, scale):
-        b = a.reshape ( (a.shape[0], a.shape[1]*a.shape[2]) )    
+        b = a.reshape ( (a.shape[0], a.shape[1]*a.shape[2]) )
         c = b.argmax(1).reshape ( (a.shape[0], 1) ).repeat(2, axis=1).astype(np.float)
-        c[:,0] %= a.shape[2]    
+        c[:,0] %= a.shape[2]
         c[:,1] = np.apply_along_axis ( lambda x: np.floor(x / a.shape[2]), 0, c[:,1] )
 
         for i in range(a.shape[0]):
@@ -113,6 +113,6 @@ class LandmarksExtractor(object):
             if pX > 0 and pX < 63 and pY > 0 and pY < 63:
                 diff = np.array ( [a[i,pY,pX+1]-a[i,pY,pX-1], a[i,pY+1,pX]-a[i,pY-1,pX]] )
                 c[i] += np.sign(diff)*0.25
-       
+
         c += 0.5
-        return [ self.transform (c[i], center, scale, a.shape[2]) for i in range(a.shape[0]) ]
+        return [ self.transform (c[i], center, scale, a.shape[2]) for i in range(a.shape[0]) ]

facelib/LandmarksProcessor.py

@@ -36,7 +36,7 @@ landmarks_68_pt = { "mouth": (48,68),
                     "left_eye": (42, 48),
                     "nose": (27, 36), # missed one point
                     "jaw": (0, 17) }
-    
+
 
 landmarks_68_3D = np.array( [
 [-73.393523  , -29.801432   , 47.667532   ],
@@ -107,20 +107,20 @@ landmarks_68_3D = np.array( [
 [8.449166    , 30.596216    , -20.671489  ],
 [0.205322    , 31.408738    , -21.903670  ],
 [-7.198266   , 30.844876    , -20.328022  ] ], dtype=np.float32)
-    
+
 def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0):
     if not isinstance(image_landmarks, np.ndarray):
-        image_landmarks = np.array (image_landmarks) 
-        
+        image_landmarks = np.array (image_landmarks)
+
     if face_type == FaceType.AVATAR:
         centroid = np.mean (image_landmarks, axis=0)
-        
+
         mat = umeyama(image_landmarks[17:], landmarks_2D, True)[0:2]
         a, c = mat[0,0], mat[1,0]
         scale = math.sqrt((a * a) + (c * c))
-        
+
         padding = (output_size / 64) * 32
-        
+
         mat = np.eye ( 2,3 )
         mat[0,2] = -centroid[0]
         mat[1,2] = -centroid[1]
@@ -135,15 +135,15 @@ def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0):
             padding = (output_size / 64) * 24
         else:
             raise ValueError ('wrong face_type: ', face_type)
-        
+
         mat = umeyama(image_landmarks[17:], landmarks_2D, True)[0:2]
         mat = mat * (output_size - 2 * padding)
-        mat[:,2] += padding        
+        mat[:,2] += padding
         mat *= (1 / scale)
         mat[:,2] += -output_size*( ( (1 / scale) - 1.0 ) / 2 )
-             
+
     return mat
-    
+
 def transform_points(points, mat, invert=False):
     if invert:
         mat = cv2.invertAffineTransform (mat)
@@ -151,68 +151,68 @@ def transform_points(points, mat, invert=False):
     points = cv2.transform(points, mat, points.shape)
     points = np.squeeze(points)
     return points
-    
-    
-def get_image_hull_mask (image_shape, image_landmarks):        
+
+
+def get_image_hull_mask (image_shape, image_landmarks):
     if len(image_landmarks) != 68:
         raise Exception('get_image_hull_mask works only with 68 landmarks')
-        
+
     hull_mask = np.zeros(image_shape[0:2]+(1,),dtype=np.float32)
 
-    cv2.fillConvexPoly( hull_mask, cv2.convexHull( 
-            np.concatenate ( (image_landmarks[0:9], 
+    cv2.fillConvexPoly( hull_mask, cv2.convexHull(
+            np.concatenate ( (image_landmarks[0:9],
                               image_landmarks[17:18]))) , (1,)  )
-                              
+
     cv2.fillConvexPoly( hull_mask, cv2.convexHull(
-            np.concatenate ( (image_landmarks[8:17], 
+            np.concatenate ( (image_landmarks[8:17],
                               image_landmarks[26:27]))) , (1,)  )
-                              
+
     cv2.fillConvexPoly( hull_mask, cv2.convexHull(
-            np.concatenate ( (image_landmarks[17:20], 
+            np.concatenate ( (image_landmarks[17:20],
                               image_landmarks[8:9]))) , (1,)  )
-                              
+
     cv2.fillConvexPoly( hull_mask, cv2.convexHull(
-            np.concatenate ( (image_landmarks[24:27], 
+            np.concatenate ( (image_landmarks[24:27],
                               image_landmarks[8:9]))) , (1,)  )
-                              
+
     cv2.fillConvexPoly( hull_mask, cv2.convexHull(
-            np.concatenate ( (image_landmarks[19:25], 
+            np.concatenate ( (image_landmarks[19:25],
                               image_landmarks[8:9],
                               ))) , (1,)  )
-                              
+
     cv2.fillConvexPoly( hull_mask, cv2.convexHull(
-            np.concatenate ( (image_landmarks[17:22], 
+            np.concatenate ( (image_landmarks[17:22],
                               image_landmarks[27:28],
                               image_landmarks[31:36],
                               image_landmarks[8:9]
-                              ))) , (1,)  )   
-                              
+                              ))) , (1,)  )
+
     cv2.fillConvexPoly( hull_mask, cv2.convexHull(
-            np.concatenate ( (image_landmarks[22:27], 
+            np.concatenate ( (image_landmarks[22:27],
                               image_landmarks[27:28],
                               image_landmarks[31:36],
                               image_landmarks[8:9]
-                              ))) , (1,)  )   
-                              
+                              ))) , (1,)  )
+
     #nose
     cv2.fillConvexPoly( hull_mask, cv2.convexHull(image_landmarks[27:36]), (1,) )
-    
+
     return hull_mask
 
-def get_image_eye_mask (image_shape, image_landmarks):        
+def get_image_eye_mask (image_shape, image_landmarks):
     if len(image_landmarks) != 68:
         raise Exception('get_image_eye_mask works only with 68 landmarks')
-        
+
     hull_mask = np.zeros(image_shape[0:2]+(1,),dtype=np.float32)
 
     cv2.fillConvexPoly( hull_mask, cv2.convexHull( image_landmarks[36:42]), (1,) )
     cv2.fillConvexPoly( hull_mask, cv2.convexHull( image_landmarks[42:48]), (1,) )
-    
+
     return hull_mask
-    
-def get_image_hull_mask_3D (image_shape, image_landmarks):    
+
+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):
@@ -224,7 +224,7 @@ def blur_image_hull_mask (hull_mask):
     leny = maxy - miny;
     masky = int(minx+(lenx//2))
     maskx = int(miny+(leny//2))
-    lowest_len = min (lenx, leny)        
+    lowest_len = min (lenx, leny)
     ero = int( lowest_len * 0.085 )
     blur = int( lowest_len * 0.10 )
 
@@ -233,10 +233,10 @@ def blur_image_hull_mask (hull_mask):
     hull_mask = np.expand_dims (hull_mask,-1)
 
     return hull_mask
-    
-def get_blurred_image_hull_mask(image_shape, image_landmarks):        
+
+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],
@@ -246,23 +246,23 @@ mirror_idxs = [
     [5,11],
     [6,10],
     [7,9],
-    
+
     [17,26],
     [18,25],
     [19,24],
     [20,23],
-    [21,22],    
-    
+    [21,22],
+
     [36,45],
     [37,44],
     [38,43],
     [39,42],
     [40,47],
-    [41,46],    
-    
+    [41,46],
+
     [31,35],
     [32,34],
-    
+
     [50,52],
     [49,53],
     [48,54],
@@ -271,28 +271,28 @@ mirror_idxs = [
     [67,65],
     [60,64],
     [61,63] ]
-    
-def mirror_landmarks (landmarks, val):    
+
+def mirror_landmarks (landmarks, val):
     result = landmarks.copy()
-    
+
     for idx in mirror_idxs:
         result [ idx ] = result [ idx[::-1] ]
 
     result[:,0] = val - result[:,0] - 1
     return result
-    
+
 def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=False):
     if len(image_landmarks) != 68:
-        raise Exception('get_image_eye_mask works only with 68 landmarks')   
-        
-    jaw = image_landmarks[slice(*landmarks_68_pt["jaw"])]        
+        raise Exception('get_image_eye_mask works only with 68 landmarks')
+
+    jaw = image_landmarks[slice(*landmarks_68_pt["jaw"])]
     right_eyebrow = image_landmarks[slice(*landmarks_68_pt["right_eyebrow"])]
     left_eyebrow = image_landmarks[slice(*landmarks_68_pt["left_eyebrow"])]
-    mouth = image_landmarks[slice(*landmarks_68_pt["mouth"])]                                    
-    right_eye = image_landmarks[slice(*landmarks_68_pt["right_eye"])]                     
-    left_eye = image_landmarks[slice(*landmarks_68_pt["left_eye"])]            
-    nose = image_landmarks[slice(*landmarks_68_pt["nose"])]            
-    
+    mouth = image_landmarks[slice(*landmarks_68_pt["mouth"])]
+    right_eye = image_landmarks[slice(*landmarks_68_pt["right_eye"])]
+    left_eye = image_landmarks[slice(*landmarks_68_pt["left_eye"])]
+    nose = image_landmarks[slice(*landmarks_68_pt["nose"])]
+
     # open shapes
     cv2.polylines(image, tuple(np.array([v]) for v in ( right_eyebrow, jaw, left_eyebrow, np.concatenate((nose, [nose[-6]])) )),
                   False, color, lineType=cv2.LINE_AA)
@@ -303,9 +303,9 @@ def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=Fa
     for x, y in np.concatenate((right_eyebrow, left_eyebrow, mouth, right_eye, left_eye, nose), axis=0):
         cv2.circle(image, (x, y), 1, color, 1, lineType=cv2.LINE_AA)
     # jaw big circles
-    for x, y in jaw: 
+    for x, y in jaw:
         cv2.circle(image, (x, y), 2, color, lineType=cv2.LINE_AA)
-    
+
     if transparent_mask:
         mask = get_image_hull_mask (image.shape, image_landmarks)
         image[...] = ( image * (1-mask) + image * mask / 2 )[...]
@@ -314,24 +314,24 @@ def draw_rect_landmarks (image, rect, image_landmarks, face_size, face_type, tra
     draw_landmarks(image, image_landmarks, color=landmarks_color, transparent_mask=transparent_mask)
     image_utils.draw_rect (image, rect, (255,0,0), 2 )
 
-    image_to_face_mat = get_transform_mat (image_landmarks, face_size, face_type)        
+    image_to_face_mat = get_transform_mat (image_landmarks, face_size, face_type)
     points = transform_points ( [ (0,0), (0,face_size-1), (face_size-1, face_size-1), (face_size-1,0) ], image_to_face_mat, True)
-    image_utils.draw_polygon (image, points, (0,0,255), 2)  
-    
+    image_utils.draw_polygon (image, points, (0,0,255), 2)
+
 def calc_face_pitch(landmarks):
     if not isinstance(landmarks, np.ndarray):
         landmarks = np.array (landmarks)
-    t = ( (landmarks[6][1]-landmarks[8][1]) + (landmarks[10][1]-landmarks[8][1]) ) / 2.0   
+    t = ( (landmarks[6][1]-landmarks[8][1]) + (landmarks[10][1]-landmarks[8][1]) ) / 2.0
     b = landmarks[8][1]
     return float(b-t)
-    
+
 def calc_face_yaw(landmarks):
     if not isinstance(landmarks, np.ndarray):
         landmarks = np.array (landmarks)
-    l = ( (landmarks[27][0]-landmarks[0][0]) + (landmarks[28][0]-landmarks[1][0]) + (landmarks[29][0]-landmarks[2][0]) ) / 3.0   
+    l = ( (landmarks[27][0]-landmarks[0][0]) + (landmarks[28][0]-landmarks[1][0]) + (landmarks[29][0]-landmarks[2][0]) ) / 3.0
     r = ( (landmarks[16][0]-landmarks[27][0]) + (landmarks[15][0]-landmarks[28][0]) + (landmarks[14][0]-landmarks[29][0]) ) / 3.0
     return float(r-l)
-  
+
 #returns pitch,yaw [-1...+1]
 def estimate_pitch_yaw(aligned_256px_landmarks):
     shape = (256,256)
@@ -347,8 +347,8 @@ def estimate_pitch_yaw(aligned_256px_landmarks):
         aligned_256px_landmarks.astype(np.float32),
         camera_matrix,
         np.zeros((4, 1)) )
- 
+
     pitch, yaw, _ = mathlib.rotationMatrixToEulerAngles( cv2.Rodrigues(rotation_vector)[0] )
     pitch = np.clip ( pitch*1.25, -1.0, 1.0 )
     yaw = np.clip ( yaw*1.25, -1.0, 1.0 )
-    return pitch, yaw
+    return pitch, yaw

facelib/MTCExtractor.py

@@ -5,10 +5,10 @@ import cv2
 from pathlib import Path
 from nnlib import nnlib
 
-class MTCExtractor(object):   
+class MTCExtractor(object):
     def __init__(self):
         self.scale_to = 1920
-        
+
         self.min_face_size = self.scale_to * 0.042
         self.thresh1 = 0.7
         self.thresh2 = 0.85
@@ -26,12 +26,12 @@ class MTCExtractor(object):
         x = Conv2D (32, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv3")(x)
         x = PReLU (shared_axes=[1,2], name="PReLU3" )(x)
         prob = Conv2D (2, kernel_size=(1,1), strides=(1,1), padding='valid', name="conv41")(x)
-        prob = Softmax()(prob)    
+        prob = Softmax()(prob)
         x = Conv2D (4, kernel_size=(1,1), strides=(1,1), padding='valid', name="conv42")(x)
 
-        PNet_model = Model(PNet_Input, [x,prob] )        
+        PNet_model = Model(PNet_Input, [x,prob] )
         PNet_model.load_weights ( (Path(__file__).parent / 'mtcnn_pnet.h5').__str__() )
-        
+
         RNet_Input = Input ( (24, 24, 3) )
         x = RNet_Input
         x = Conv2D (28, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv1")(x)
@@ -39,18 +39,18 @@ class MTCExtractor(object):
         x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='same' ) (x)
         x = Conv2D (48, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv2")(x)
         x = PReLU (shared_axes=[1,2], name="prelu2" )(x)
-        x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='valid' ) (x)    
+        x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='valid' ) (x)
         x = Conv2D (64, kernel_size=(2,2), strides=(1,1), padding='valid', name="conv3")(x)
         x = PReLU (shared_axes=[1,2], name="prelu3" )(x)
         x = Lambda ( lambda x: K.reshape (x, (-1, np.prod(K.int_shape(x)[1:]),) ), output_shape=(np.prod(K.int_shape(x)[1:]),) ) (x)
-        x = Dense (128, name='conv4')(x)    
+        x = Dense (128, name='conv4')(x)
         x = PReLU (name="prelu4" )(x)
         prob = Dense (2, name='conv51')(x)
-        prob = Softmax()(prob)  
-        x = Dense (4, name='conv52')(x)        
-        RNet_model = Model(RNet_Input, [x,prob] )        
+        prob = Softmax()(prob)
+        x = Dense (4, name='conv52')(x)
+        RNet_model = Model(RNet_Input, [x,prob] )
         RNet_model.load_weights ( (Path(__file__).parent / 'mtcnn_rnet.h5').__str__() )
-        
+
         ONet_Input = Input ( (48, 48, 3) )
         x = ONet_Input
         x = Conv2D (32, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv1")(x)
@@ -58,20 +58,20 @@ class MTCExtractor(object):
         x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='same' ) (x)
         x = Conv2D (64, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv2")(x)
         x = PReLU (shared_axes=[1,2], name="prelu2" )(x)
-        x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='valid' ) (x)    
+        x = MaxPooling2D( pool_size=(3,3), strides=(2,2), padding='valid' ) (x)
         x = Conv2D (64, kernel_size=(3,3), strides=(1,1), padding='valid', name="conv3")(x)
         x = PReLU (shared_axes=[1,2], name="prelu3" )(x)
-        x = MaxPooling2D( pool_size=(2,2), strides=(2,2), padding='same' ) (x) 
+        x = MaxPooling2D( pool_size=(2,2), strides=(2,2), padding='same' ) (x)
         x = Conv2D (128, kernel_size=(2,2), strides=(1,1), padding='valid', name="conv4")(x)
         x = PReLU (shared_axes=[1,2], name="prelu4" )(x)
-        x = Lambda ( lambda x: K.reshape (x, (-1, np.prod(K.int_shape(x)[1:]),) ), output_shape=(np.prod(K.int_shape(x)[1:]),) ) (x)    
+        x = Lambda ( lambda x: K.reshape (x, (-1, np.prod(K.int_shape(x)[1:]),) ), output_shape=(np.prod(K.int_shape(x)[1:]),) ) (x)
         x = Dense (256, name='conv5')(x)
         x = PReLU (name="prelu5" )(x)
         prob = Dense (2, name='conv61')(x)
-        prob = Softmax()(prob)    
+        prob = Softmax()(prob)
         x1 = Dense (4, name='conv62')(x)
-        x2 = Dense (10, name='conv63')(x)        
-        ONet_model = Model(ONet_Input, [x1,x2,prob] )        
+        x2 = Dense (10, name='conv63')(x)
+        ONet_model = Model(ONet_Input, [x1,x2,prob] )
         ONet_model.load_weights ( (Path(__file__).parent / 'mtcnn_onet.h5').__str__() )
 
         self.pnet_fun = K.function ( PNet_model.inputs, PNet_model.outputs )
@@ -79,13 +79,13 @@ class MTCExtractor(object):
         self.onet_fun = K.function ( ONet_model.inputs, ONet_model.outputs )
 
     def __enter__(self):
-        faces, pnts = detect_face ( np.zeros ( (self.scale_to, self.scale_to, 3)), self.min_face_size, self.pnet_fun, self.rnet_fun, self.onet_fun, [ self.thresh1, self.thresh2, self.thresh3 ], self.scale_factor ) 
-        
+        faces, pnts = detect_face ( np.zeros ( (self.scale_to, self.scale_to, 3)), self.min_face_size, self.pnet_fun, self.rnet_fun, self.onet_fun, [ self.thresh1, self.thresh2, self.thresh3 ], self.scale_factor )
+
         return self
-        
+
     def __exit__(self, exc_type=None, exc_value=None, traceback=None):
         return False #pass exception between __enter__ and __exit__ to outter level
-        
+
     def extract_from_bgr (self, input_image):
         input_image = input_image[:,:,::-1].copy()
         (h, w, ch) = input_image.shape
@@ -95,7 +95,7 @@ class MTCExtractor(object):
 
         detected_faces, pnts = detect_face ( input_image, self.min_face_size, self.pnet_fun, self.rnet_fun, self.onet_fun, [ self.thresh1, self.thresh2, self.thresh3 ], self.scale_factor )
         detected_faces = [ ( int(face[0]/input_scale), int(face[1]/input_scale), int(face[2]/input_scale), int(face[3]/input_scale)) for face in detected_faces ]
-          
+
         return detected_faces
 
 def detect_face(img, minsize, pnet, rnet, onet, threshold, factor):
@@ -132,9 +132,9 @@ def detect_face(img, minsize, pnet, rnet, onet, threshold, factor):
         out = pnet([img_y])
         out0 = np.transpose(out[0], (0,2,1,3))
         out1 = np.transpose(out[1], (0,2,1,3))
-        
+
         boxes, _ = generateBoundingBox(out1[0,:,:,1].copy(), out0[0,:,:,:].copy(), scale, threshold[0])
-        
+
         # inter-scale nms
         pick = nms(boxes.copy(), 0.5, 'Union')
         if boxes.size>0 and pick.size>0:
@@ -217,7 +217,7 @@ def detect_face(img, minsize, pnet, rnet, onet, threshold, factor):
             pick = nms(total_boxes.copy(), 0.7, 'Min')
             total_boxes = total_boxes[pick,:]
             points = points[:,pick]
-                
+
     return total_boxes, points
 
 
@@ -235,7 +235,7 @@ def bbreg(boundingbox,reg):
     b4 = boundingbox[:,3]+reg[:,3]*h
     boundingbox[:,0:4] = np.transpose(np.vstack([b1, b2, b3, b4 ]))
     return boundingbox
- 
+
 def generateBoundingBox(imap, reg, scale, t):
     """Use heatmap to generate bounding boxes"""
     stride=2
@@ -261,7 +261,7 @@ def generateBoundingBox(imap, reg, scale, t):
     q2 = np.fix((stride*bb+cellsize-1+1)/scale)
     boundingbox = np.hstack([q1, q2, np.expand_dims(score,1), reg])
     return boundingbox, reg
- 
+
 # function pick = nms(boxes,threshold,type)
 def nms(boxes, threshold, method):
     if boxes.size==0:
@@ -315,7 +315,7 @@ def pad(total_boxes, w, h):
     tmp = np.where(ex>w)
     edx.flat[tmp] = np.expand_dims(-ex[tmp]+w+tmpw[tmp],1)
     ex[tmp] = w
-    
+
     tmp = np.where(ey>h)
     edy.flat[tmp] = np.expand_dims(-ey[tmp]+h+tmph[tmp],1)
     ey[tmp] = h
@@ -327,7 +327,7 @@ def pad(total_boxes, w, h):
     tmp = np.where(y<1)
     dy.flat[tmp] = np.expand_dims(2-y[tmp],1)
     y[tmp] = 1
-    
+
     return dy, edy, dx, edx, y, ey, x, ex, tmpw, tmph
 
 # function [bboxA] = rerec(bboxA)

facelib/S3FDExtractor.py

@@ -3,35 +3,35 @@ from pathlib import Path
 import cv2
 from nnlib import nnlib
 
-class S3FDExtractor(object):   
+class S3FDExtractor(object):
     def __init__(self):
         exec( nnlib.import_all(), locals(), globals() )
-        
+
         model_path = Path(__file__).parent / "S3FD.h5"
         if not model_path.exists():
             return None
 
-        self.model = nnlib.keras.models.load_model ( str(model_path) ) 
+        self.model = nnlib.keras.models.load_model ( str(model_path) )
 
     def __enter__(self):
         return self
-        
+
     def __exit__(self, exc_type=None, exc_value=None, traceback=None):
         return False #pass exception between __enter__ and __exit__ to outter level
-        
+
     def extract_from_bgr (self, input_image):
         input_image = input_image[:,:,::-1].copy()
         (h, w, ch) = input_image.shape
-        
+
         d = max(w, h)
         scale_to = 640 if d >= 1280 else d / 2
         scale_to = max(64, scale_to)
-            
+
         input_scale = d / scale_to
         input_image = cv2.resize (input_image, ( int(w/input_scale), int(h/input_scale) ), interpolation=cv2.INTER_LINEAR)
 
         olist = self.model.predict( np.expand_dims(input_image,0) )
-        
+
         detected_faces = []
         for ltrb in self.refine (olist):
             l,t,r,b = [ x*input_scale for x in ltrb]
@@ -42,7 +42,7 @@ class S3FDExtractor(object):
             detected_faces.append ( [int(x) for x in (l,t,r,b) ] )
 
         return detected_faces
-        
+
     def refine(self, olist):
         bboxlist = []
         for i, ((ocls,), (oreg,)) in enumerate ( zip ( olist[::2], olist[1::2] ) ):
@@ -51,7 +51,7 @@ class S3FDExtractor(object):
             s_m4 = stride * 4
 
             for hindex, windex in zip(*np.where(ocls > 0.05)):
-                score = ocls[hindex, windex]                
+                score = ocls[hindex, windex]
                 loc   = oreg[hindex, windex, :]
                 priors = np.array([windex * stride + s_d2, hindex * stride + s_d2, s_m4, s_m4])
                 priors_2p = priors[2:]
@@ -61,15 +61,15 @@ class S3FDExtractor(object):
                 box[2:] += box[:2]
 
                 bboxlist.append([*box, score])
- 
+
         bboxlist = np.array(bboxlist)
         if len(bboxlist) == 0:
             bboxlist = np.zeros((1, 5))
-        
+
         bboxlist = bboxlist[self.refine_nms(bboxlist, 0.3), :]
         bboxlist = [ x[:-1].astype(np.int) for x in bboxlist if x[-1] >= 0.5]
         return bboxlist
-        
+
     def refine_nms(self, dets, thresh):
         keep = list()
         if len(dets) == 0:
@@ -91,4 +91,4 @@ class S3FDExtractor(object):
 
             inds = np.where(ovr <= thresh)[0]
             order = order[inds + 1]
-        return keep
+        return keep

imagelib/estimate_sharpness.py

@@ -1,6 +1,6 @@
 """
 Copyright (c) 2009-2010 Arizona Board of Regents.  All Rights Reserved.
- Contact: Lina Karam (karam@asu.edu) and Niranjan Narvekar (nnarveka@asu.edu) 
+ Contact: Lina Karam (karam@asu.edu) and Niranjan Narvekar (nnarveka@asu.edu)
  Image, Video, and Usabilty (IVU) Lab, http://ivulab.asu.edu , Arizona State University
  This copyright statement may not be removed from any file containing it or from modifications to these files.
  This copyright notice must also be included in any file or product that is derived from the source files.
@@ -267,11 +267,11 @@ def get_block_contrast(block):
     # type: (numpy.ndarray) -> int
     return int(np.max(block) - np.min(block))
 
-    
-def estimate_sharpness(image):       
+
+def estimate_sharpness(image):
     height, width = image.shape[:2]
-    
+
     if image.ndim == 3:
         image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
 
-    return compute(image)
+    return compute(image)

interact/__init__.py

@@ -1 +1 @@
-from .interact import interact
+from .interact import interact

interact/interact.py

@@ -12,28 +12,28 @@ class Interact(object):
     EVENT_RBUTTONDOWN = 5
     EVENT_RBUTTONUP = 6
     EVENT_MOUSEWHEEL = 10
-    
+
     def __init__(self):
         self.named_windows = {}
         self.capture_mouse_windows = {}
-        self.capture_keys_windows = {}        
+        self.capture_keys_windows = {}
         self.mouse_events = {}
         self.key_events = {}
         self.pg_bar = None
-    
+
     def log_info(self, msg, end='\n'):
         print (msg, end=end)
-        
+
     def log_err(self, msg, end='\n'):
-        print (msg, end=end)        
-    
+        print (msg, end=end)
+
     def named_window(self, wnd_name):
         if wnd_name not in self.named_windows:
             #we will show window only on first show_image
             self.named_windows[wnd_name] = 0
-            
+
         else: print("named_window: ", wnd_name, " already created.")
-        
+
     def destroy_all_windows(self):
         if len( self.named_windows ) != 0:
             cv2.destroyAllWindows()
@@ -42,32 +42,32 @@ class Interact(object):
             self.capture_keys_windows = {}
             self.mouse_events = {}
             self.key_events = {}
-        
+
     def show_image(self, wnd_name, img):
         if wnd_name in self.named_windows:
             if self.named_windows[wnd_name] == 0:
                 self.named_windows[wnd_name] = 1
-                cv2.namedWindow(wnd_name)                
+                cv2.namedWindow(wnd_name)
                 if wnd_name in self.capture_mouse_windows:
                     self.capture_mouse(wnd_name)
-                
+
             cv2.imshow (wnd_name, img)
         else: print("show_image: named_window ", wnd_name, " not found.")
 
     def capture_mouse(self, wnd_name):
         def onMouse(event, x, y, flags, param):
-            (inst, wnd_name) = param            
+            (inst, wnd_name) = param
             if event == cv2.EVENT_LBUTTONDOWN: ev = Interact.EVENT_LBUTTONDOWN
             elif event == cv2.EVENT_LBUTTONUP: ev = Interact.EVENT_LBUTTONUP
             elif event == cv2.EVENT_RBUTTONDOWN: ev = Interact.EVENT_RBUTTONDOWN
             elif event == cv2.EVENT_RBUTTONUP: ev = Interact.EVENT_RBUTTONUP
             elif event == cv2.EVENT_MOUSEWHEEL: ev = Interact.EVENT_MOUSEWHEEL
-            
+
             else: ev = 0
-            inst.add_mouse_event (wnd_name, x, y, ev, flags) 
+            inst.add_mouse_event (wnd_name, x, y, ev, flags)
 
         if wnd_name in self.named_windows:
-            self.capture_mouse_windows[wnd_name] = True            
+            self.capture_mouse_windows[wnd_name] = True
             if self.named_windows[wnd_name] == 1:
                 cv2.setMouseCallback(wnd_name, onMouse, (self,wnd_name) )
         else: print("capture_mouse: named_window ", wnd_name, " not found.")
@@ -95,66 +95,66 @@ class Interact(object):
             self.pg_bar.close()
             self.pg_bar = None
         else: print("progress_bar not set.")
-    
+
     def progress_bar_generator(self, data, desc, leave=True):
         for x in tqdm( data, desc=desc, leave=leave, ascii=True ):
             yield x
-    
+
     def process_messages(self, sleep_time=0):
         has_windows = False
         has_capture_keys = False
 
         if len(self.named_windows) != 0:
             has_windows = True
-            
+
         if len(self.capture_keys_windows) != 0:
             has_capture_keys = True
-                
+
         if has_windows or has_capture_keys:
             wait_key_time = max(1, int(sleep_time*1000) )
             key = cv2.waitKey(wait_key_time) & 0xFF
         else:
             if sleep_time != 0:
                 time.sleep(sleep_time)
-        
+
         if has_capture_keys and key != 255:
             for wnd_name in self.capture_keys_windows:
                 self.add_key_event (wnd_name, key)
-    
+
     def wait_any_key(self):
         cv2.waitKey(0)
-    
+
     def add_mouse_event(self, wnd_name, x, y, ev, flags):
-        if wnd_name not in self.mouse_events: 
+        if wnd_name not in self.mouse_events:
             self.mouse_events[wnd_name] = []
         self.mouse_events[wnd_name] += [ (x, y, ev, flags) ]
-        
+
     def add_key_event(self, wnd_name, key):
-        if wnd_name not in self.key_events: 
+        if wnd_name not in self.key_events:
             self.key_events[wnd_name] = []
         self.key_events[wnd_name] += [ (key,) ]
 
     def get_mouse_events(self, wnd_name):
         ar = self.mouse_events.get(wnd_name, [])
-        self.mouse_events[wnd_name] = []        
+        self.mouse_events[wnd_name] = []
         return ar
-        
+
     def get_key_events(self, wnd_name):
         ar = self.key_events.get(wnd_name, [])
-        self.key_events[wnd_name] = []        
+        self.key_events[wnd_name] = []
         return ar
-        
+
     def input_number(self, s, default_value, valid_list=None, help_message=None):
         while True:
             try:
                 inp = input(s)
                 if len(inp) == 0:
                     raise ValueError("")
-                    
+
                 if help_message is not None and inp == '?':
                     print (help_message)
                     continue
-                
+
                 i = float(inp)
                 if (valid_list is not None) and (i not in valid_list):
                     return default_value
@@ -162,18 +162,18 @@ class Interact(object):
             except:
                 print (default_value)
                 return default_value
-                
+
     def input_int(self,s, default_value, valid_list=None, help_message=None):
         while True:
             try:
                 inp = input(s)
                 if len(inp) == 0:
                     raise ValueError("")
-                    
+
                 if help_message is not None and inp == '?':
                     print (help_message)
                     continue
-                
+
                 i = int(inp)
                 if (valid_list is not None) and (i not in valid_list):
                     return default_value
@@ -181,41 +181,41 @@ class Interact(object):
             except:
                 print (default_value)
                 return default_value
-            
+
     def input_bool(self, s, default_value, help_message=None):
         while True:
             try:
                 inp = input(s)
                 if len(inp) == 0:
                     raise ValueError("")
-                    
+
                 if help_message is not None and inp == '?':
                     print (help_message)
                     continue
-                    
+
                 return bool ( {"y":True,"n":False,"1":True,"0":False}.get(inp.lower(), default_value) )
             except:
                 print ( "y" if default_value else "n" )
                 return default_value
-            
+
     def input_str(self, s, default_value, valid_list=None, help_message=None):
-        while True:        
+        while True:
             try:
                 inp = input(s)
                 if len(inp) == 0:
                     raise ValueError("")
-                    
+
                 if help_message is not None and inp == '?':
                     print (help_message)
                     continue
-                    
+
                 if (valid_list is not None) and (inp.lower() not in valid_list):
                     return default_value
                 return inp
             except:
                 print (default_value)
                 return default_value
-                
+
     def input_process(self, stdin_fd, sq, str):
         sys.stdin = os.fdopen(stdin_fd)
         try:
@@ -223,7 +223,7 @@ class Interact(object):
             sq.put (True)
         except:
             sq.put (False)
-            
+
     def input_in_time (self, str, max_time_sec):
         sq = multiprocessing.Queue()
         p = multiprocessing.Process(target=self.input_process, args=( sys.stdin.fileno(), sq, str))
@@ -240,4 +240,4 @@ class Interact(object):
         sys.stdin = os.fdopen( sys.stdin.fileno() )
         return inp
 
-interact = Interact()
+interact = Interact()

joblib/SubprocessFunctionCaller.py

@@ -7,7 +7,7 @@ class SubprocessFunctionCaller(object):
             self.s2c = s2c
             self.c2s = c2s
             self.lock = lock
-            
+
         def __call__(self, value):
             self.lock.acquire()
             self.c2s.put (value)
@@ -17,26 +17,26 @@ class SubprocessFunctionCaller(object):
                     self.lock.release()
                     return obj
                 time.sleep(0.005)
-    
+
     class HostProcessor(object):
         def __init__(self, s2c, c2s, func):
             self.s2c = s2c
             self.c2s = c2s
             self.func = func
-            
+
         def process_messages(self):
             while not self.c2s.empty():
                 obj = self.c2s.get()
                 result = self.func (obj)
                 self.s2c.put (result)
-                
+
     @staticmethod
     def make_pair( func ):
         s2c = multiprocessing.Queue()
         c2s = multiprocessing.Queue()
         lock = multiprocessing.Lock()
-        
+
         host_processor = SubprocessFunctionCaller.HostProcessor (s2c, c2s, func)
         cli_func = SubprocessFunctionCaller.CliFunction (s2c, c2s, lock)
-        
-        return host_processor, cli_func
+
+        return host_processor, cli_func

joblib/SubprocessorBase.py

@@ -3,12 +3,12 @@ import multiprocessing
 import time
 import sys
 from interact import interact as io
-                
+
 class Subprocessor(object):
 
     class SilenceException(Exception):
         pass
-        
+
     class Cli(object):
         def __init__ ( self, client_dict ):
             self.s2c = multiprocessing.Queue()
@@ -16,41 +16,41 @@ class Subprocessor(object):
             self.p = multiprocessing.Process(target=self._subprocess_run, args=(client_dict,) )
             self.p.daemon = True
             self.p.start()
-            
+
             self.state = None
             self.sent_time = None
             self.sent_data = None
             self.name = None
             self.host_dict = None
-        
+
         def kill(self):
             self.p.terminate()
             self.p.join()
-                
+
         #overridable optional
         def on_initialize(self, client_dict):
             #initialize your subprocess here using client_dict
             pass
-            
+
         #overridable optional
         def on_finalize(self):
             #finalize your subprocess here
             pass
-            
+
         #overridable
         def process_data(self, data):
             #process 'data' given from host and return result
             raise NotImplementedError
-     
+
         #overridable optional
         def get_data_name (self, data):
             #return string identificator of your 'data'
             return "undefined"
-            
+
         def log_info(self, msg): self.c2s.put ( {'op': 'log_info', 'msg':msg } )
         def log_err(self, msg): self.c2s.put ( {'op': 'log_err' , 'msg':msg } )
         def progress_bar_inc(self, c): self.c2s.put ( {'op': 'progress_bar_inc' , 'c':c } )
-        
+
         def _subprocess_run(self, client_dict):
             data = None
             s2c, c2s = self.s2c, self.c2s
@@ -65,20 +65,20 @@ class Subprocessor(object):
                     if op == 'data':
                         data = msg['data']
                         result = self.process_data (data)
-                        c2s.put ( {'op': 'success', 'data' : data, 'result' : result} )                        
+                        c2s.put ( {'op': 'success', 'data' : data, 'result' : result} )
                         data = None
                     elif op == 'close':
                         break
 
                     time.sleep(0.001)
-                        
+
                 self.on_finalize()
                 c2s.put ( {'op': 'finalized'} )
                 return
             except Subprocessor.SilenceException as e:
                 pass
             except Exception as e:
-                if data is not None:        
+                if data is not None:
                     print ('Exception while process data [%s]: %s' % (self.get_data_name(data), traceback.format_exc()) )
                 else:
                     print ('Exception: %s' % (traceback.format_exc()) )
@@ -91,10 +91,10 @@ class Subprocessor(object):
             raise ValueError("SubprocessorCli_class must be subclass of Subprocessor.Cli")
 
         self.name = name
-        self.SubprocessorCli_class = SubprocessorCli_class        
+        self.SubprocessorCli_class = SubprocessorCli_class
         self.no_response_time_sec = no_response_time_sec
 
-    #overridable    
+    #overridable
     def process_info_generator(self):
         #yield per process (name, host_dict, client_dict)
         raise NotImplementedError
@@ -103,42 +103,42 @@ class Subprocessor(object):
     def on_clients_initialized(self):
         #logic when all subprocesses initialized and ready
         pass
-        
+
     #overridable optional
     def on_clients_finalized(self):
         #logic when all subprocess finalized
         pass
-        
-    #overridable        
+
+    #overridable
     def get_data(self, host_dict):
         #return data for processing here
         raise NotImplementedError
-    
+
     #overridable
     def on_data_return (self, host_dict, data):
-        #you have to place returned 'data' back to your queue  
+        #you have to place returned 'data' back to your queue
         raise NotImplementedError
-        
+
     #overridable
     def on_result (self, host_dict, data, result):
         #your logic what to do with 'result' of 'data'
         raise NotImplementedError
-    
+
     #overridable
     def get_result(self):
         #return result that will be returned in func run()
         raise NotImplementedError
- 
+
     #overridable
     def on_tick(self):
         #tick in main loop
         pass
- 
+
     def run(self):
         self.clis = []
-            
+
         #getting info about name of subprocesses, host and client dicts, and spawning them
-        for name, host_dict, client_dict in self.process_info_generator():            
+        for name, host_dict, client_dict in self.process_info_generator():
             try:
                 cli = self.SubprocessorCli_class(client_dict)
                 cli.state = 1
@@ -146,21 +146,21 @@ class Subprocessor(object):
                 cli.sent_data = None
                 cli.name = name
                 cli.host_dict = host_dict
-                
+
                 self.clis.append (cli)
-                
+
             except:
                 raise Exception ("Unable to start subprocess %s" % (name))
 
         if len(self.clis) == 0:
             raise Exception ("Unable to start Subprocessor '%s' " % (self.name))
-            
+
         #waiting subprocesses their success(or not) initialization
         while True:
             for cli in self.clis[:]:
                 while not cli.c2s.empty():
                     obj = cli.c2s.get()
-                    op = obj.get('op','')                        
+                    op = obj.get('op','')
                     if op == 'init_ok':
                         cli.state = 0
                     elif op == 'log_info':
@@ -172,16 +172,16 @@ class Subprocessor(object):
                         self.clis.remove(cli)
                         break
             if all ([cli.state == 0 for cli in self.clis]):
-                break                
+                break
             io.process_messages(0.005)
-            
+
         if len(self.clis) == 0:
             raise Exception ( "Unable to start subprocesses." )
-        
-        #ok some processes survived, initialize host logic        
-        
+
+        #ok some processes survived, initialize host logic
+
         self.on_clients_initialized()
- 
+
         #main loop of data processing
         while True:
             for cli in self.clis[:]:
@@ -206,10 +206,10 @@ class Subprocessor(object):
                         io.log_err(obj['msg'])
                     elif op == 'progress_bar_inc':
                         io.progress_bar_inc(obj['c'])
-                        
+
             for cli in self.clis[:]:
                 if cli.state == 0:
-                    #free state of subprocess, get some data from get_data 
+                    #free state of subprocess, get some data from get_data
                     data = self.get_data(cli.host_dict)
                     if data is not None:
                         #and send it to subprocess
@@ -217,7 +217,7 @@ class Subprocessor(object):
                         cli.sent_time = time.time()
                         cli.sent_data = data
                         cli.state = 1
-                    
+
                 elif cli.state == 1:
                     if self.no_response_time_sec != 0 and (time.time() - cli.sent_time) > self.no_response_time_sec:
                         #subprocess busy too long
@@ -225,39 +225,39 @@ class Subprocessor(object):
                         self.on_data_return (cli.host_dict, cli.sent_data )
                         cli.kill()
                         self.clis.remove(cli)
-                        
+
             if all ([cli.state == 0 for cli in self.clis]):
                 #all subprocesses free and no more data available to process, ending loop
-                break                
+                break
             io.process_messages(0.005)
             self.on_tick()
-        
+
         #gracefully terminating subprocesses
         for cli in self.clis[:]:
             cli.s2c.put ( {'op': 'close'} )
             cli.sent_time = time.time()
-            
+
         while True:
             for cli in self.clis[:]:
                 terminate_it = False
                 while not cli.c2s.empty():
                     obj = cli.c2s.get()
-                    obj_op = obj['op']                    
+                    obj_op = obj['op']
                     if obj_op == 'finalized':
                         terminate_it = True
                         break
-                        
-                if self.no_response_time_sec != 0 and (time.time() - cli.sent_time) > self.no_response_time_sec:                
+
+                if self.no_response_time_sec != 0 and (time.time() - cli.sent_time) > self.no_response_time_sec:
                     terminate_it = True
-                    
+
                 if terminate_it:
                     cli.state = 2
                     cli.kill()
-                    
+
             if all ([cli.state == 2 for cli in self.clis]):
                 break
-        
+
         #finalizing host logic and return result
         self.on_clients_finalized()
-         
+
         return self.get_result()

joblib/__init__.py

@@ -1,2 +1,2 @@
 from .SubprocessorBase import Subprocessor
-from .SubprocessFunctionCaller import SubprocessFunctionCaller
+from .SubprocessFunctionCaller import SubprocessFunctionCaller

main.py

@@ -14,110 +14,110 @@ class fixPathAction(argparse.Action):
     def __call__(self, parser, namespace, values, option_string=None):
         setattr(namespace, self.dest, os.path.abspath(os.path.expanduser(values)))
 
-if __name__ == "__main__":    
+if __name__ == "__main__":
     multiprocessing.set_start_method("spawn")
-    
-    os_utils.set_process_lowest_prio()   
-    parser = argparse.ArgumentParser()    
+
+    os_utils.set_process_lowest_prio()
+    parser = argparse.ArgumentParser()
 
     subparsers = parser.add_subparsers()
-    
+
     def process_extract(arguments):
-        from mainscripts import Extractor      
-        Extractor.main( arguments.input_dir, 
-                        arguments.output_dir, 
+        from mainscripts import Extractor
+        Extractor.main( arguments.input_dir,
+                        arguments.output_dir,
                         arguments.debug_dir,
-                        arguments.detector,   
+                        arguments.detector,
                         arguments.manual_fix,
                         arguments.manual_output_debug_fix,
-                        arguments.manual_window_size, 
+                        arguments.manual_window_size,
                         face_type=arguments.face_type,
                         device_args={'cpu_only'  : arguments.cpu_only,
                                      'multi_gpu' : arguments.multi_gpu,
                                     }
                       )
-        
+
     p = subparsers.add_parser( "extract", help="Extract the faces from a pictures.")
     p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
     p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted files will be stored.")
-    p.add_argument('--debug-dir', action=fixPathAction, dest="debug_dir", help="Writes debug images to this directory.")    
-    p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'avatar', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'")    
+    p.add_argument('--debug-dir', action=fixPathAction, dest="debug_dir", help="Writes debug images to this directory.")
+    p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'avatar', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'")
     p.add_argument('--detector', dest="detector", choices=['dlib','mt','s3fd','manual'], default='dlib', help="Type of detector. Default 'dlib'. 'mt' (MTCNNv1) - faster, better, almost no jitter, perfect for gathering thousands faces for src-set. It is also good for dst-set, but can generate false faces in frames where main face not recognized! In this case for dst-set use either 'dlib' with '--manual-fix' or '--detector manual'. Manual detector suitable only for dst-set.")
     p.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.")
     p.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.")
     p.add_argument('--manual-output-debug-fix', action="store_true", dest="manual_output_debug_fix", default=False, help="Performs manual reextract input-dir frames which were deleted from [output_dir]_debug\ dir.")
-    p.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.") 
-    p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU. Forces to use MT extractor.")    
+    p.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.")
+    p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU. Forces to use MT extractor.")
     p.set_defaults (func=process_extract)
-    
-    def process_sort(arguments):        
+
+    def process_sort(arguments):
         from mainscripts import Sorter
         Sorter.main (input_path=arguments.input_dir, sort_by_method=arguments.sort_by_method)
-        
-    p = subparsers.add_parser( "sort", help="Sort faces in a directory.")     
+
+    p = subparsers.add_parser( "sort", help="Sort faces in a directory.")
     p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
     p.add_argument('--by', required=True, dest="sort_by_method", choices=("blur", "face", "face-dissim", "face-yaw", "face-pitch", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final", "final-no-blur", "test"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
     p.set_defaults (func=process_sort)
-    
-    def process_util(arguments):        
+
+    def process_util(arguments):
         from mainscripts import Util
-        
+
         if arguments.convert_png_to_jpg:
             Util.convert_png_to_jpg_folder (input_path=arguments.input_dir)
-        
+
         if arguments.add_landmarks_debug_images:
             Util.add_landmarks_debug_images (input_path=arguments.input_dir)
 
         if arguments.recover_original_aligned_filename:
             Util.recover_original_aligned_filename (input_path=arguments.input_dir)
-            
-    p = subparsers.add_parser( "util", help="Utilities.")     
+
+    p = subparsers.add_parser( "util", help="Utilities.")
     p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
     p.add_argument('--convert-png-to-jpg', action="store_true", dest="convert_png_to_jpg", default=False, help="Convert DeepFaceLAB PNG files to JPEG.")
     p.add_argument('--add-landmarks-debug-images', action="store_true", dest="add_landmarks_debug_images", default=False, help="Add landmarks debug image for aligned faces.")
     p.add_argument('--recover-original-aligned-filename', action="store_true", dest="recover_original_aligned_filename", default=False, help="Recover original aligned filename.")
     p.set_defaults (func=process_util)
-    
+
     def process_train(arguments):
-        args = {'training_data_src_dir'  : arguments.training_data_src_dir, 
-                'training_data_dst_dir'  : arguments.training_data_dst_dir, 
+        args = {'training_data_src_dir'  : arguments.training_data_src_dir,
+                'training_data_dst_dir'  : arguments.training_data_dst_dir,
                 'model_path'             : arguments.model_dir,
                 'model_name'             : arguments.model_name,
                 'no_preview'             : arguments.no_preview,
-                'debug'                  : arguments.debug,                
-                }                
+                'debug'                  : arguments.debug,
+                }
         device_args = {'cpu_only'  : arguments.cpu_only,
                        'force_gpu_idx' : arguments.force_gpu_idx,
                        }
-        from mainscripts import Trainer           
+        from mainscripts import Trainer
         Trainer.main(args, device_args)
 
-    p = subparsers.add_parser( "train", help="Trainer") 
+    p = subparsers.add_parser( "train", help="Trainer")
     p.add_argument('--training-data-src-dir', required=True, action=fixPathAction, dest="training_data_src_dir", help="Dir of src-set.")
     p.add_argument('--training-data-dst-dir', required=True, action=fixPathAction, dest="training_data_dst_dir", help="Dir of dst-set.")
     p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.")
     p.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model")
     p.add_argument('--no-preview', action="store_true", dest="no_preview", default=False, help="Disable preview window.")
-    p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug samples.")  
+    p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug samples.")
     p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.")
     p.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.")
     p.set_defaults (func=process_train)
-    
+
     def process_convert(arguments):
-        args = {'input_dir'   : arguments.input_dir, 
-                'output_dir'  : arguments.output_dir, 
+        args = {'input_dir'   : arguments.input_dir,
+                'output_dir'  : arguments.output_dir,
                 'aligned_dir' : arguments.aligned_dir,
                 'model_dir'   : arguments.model_dir,
                 'model_name'  : arguments.model_name,
-                'debug'       : arguments.debug,                
-                }                
+                'debug'       : arguments.debug,
+                }
         device_args = {'cpu_only'  : arguments.cpu_only,
                        'force_gpu_idx' : arguments.force_gpu_idx,
                        }
         from mainscripts import Converter
         Converter.main (args, device_args)
-        
-    p = subparsers.add_parser( "convert", help="Converter") 
+
+    p = subparsers.add_parser( "convert", help="Converter")
     p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
     p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the converted files will be stored.")
     p.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", help="Aligned directory. This is where the extracted of dst faces stored. Not used in AVATAR model.")
@@ -127,10 +127,10 @@ if __name__ == "__main__":
     p.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.")
     p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Convert on CPU.")
     p.set_defaults(func=process_convert)
-    
+
     videoed_parser = subparsers.add_parser( "videoed", help="Video processing.").add_subparsers()
-    
-    def process_videoed_extract_video(arguments):        
+
+    def process_videoed_extract_video(arguments):
         from mainscripts import VideoEd
         VideoEd.extract_video (arguments.input_file, arguments.output_dir, arguments.output_ext, arguments.fps)
     p = videoed_parser.add_parser( "extract-video", help="Extract images from video file.")
@@ -139,23 +139,23 @@ if __name__ == "__main__":
     p.add_argument('--ouptut-ext', dest="output_ext", default='png', help="Image format (extension) of output files.")
     p.add_argument('--fps', type=int, dest="fps", default=None, help="How many frames of every second of the video will be extracted. 0 - full fps.")
     p.set_defaults(func=process_videoed_extract_video)
-    
-    def process_videoed_cut_video(arguments):        
+
+    def process_videoed_cut_video(arguments):
         from mainscripts import VideoEd
-        VideoEd.cut_video (arguments.input_file, 
-                           arguments.from_time, 
-                           arguments.to_time, 
-                           arguments.audio_track_id, 
+        VideoEd.cut_video (arguments.input_file,
+                           arguments.from_time,
+                           arguments.to_time,
+                           arguments.audio_track_id,
                            arguments.bitrate)
     p = videoed_parser.add_parser( "cut-video", help="Cut video file.")
     p.add_argument('--input-file', required=True, action=fixPathAction, dest="input_file", help="Input file to be processed. Specify .*-extension to find first file.")
     p.add_argument('--from-time', dest="from_time", default=None, help="From time, for example 00:00:00.000")
     p.add_argument('--to-time', dest="to_time", default=None, help="To time, for example 00:00:00.000")
     p.add_argument('--audio-track-id', type=int, dest="audio_track_id", default=None, help="Specify audio track id.")
-    p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.")    
+    p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.")
     p.set_defaults(func=process_videoed_cut_video)
-    
-    def process_videoed_denoise_image_sequence(arguments):        
+
+    def process_videoed_denoise_image_sequence(arguments):
         from mainscripts import VideoEd
         VideoEd.denoise_image_sequence (arguments.input_dir, arguments.ext, arguments.factor)
     p = videoed_parser.add_parser( "denoise-image-sequence", help="Denoise sequence of images, keeping sharp edges. This allows you to make the final fake more believable, since the neural network is not able to make a detailed skin texture, but it makes the edges quite clear. Therefore, if the whole frame is more `blurred`, then a fake will seem more believable. Especially true for scenes of the film, which are usually very clear.")
@@ -163,65 +163,65 @@ if __name__ == "__main__":
     p.add_argument('--ext', dest="ext", default='png', help="Image format (extension) of input files.")
     p.add_argument('--factor', type=int, dest="factor", default=None, help="Denoise factor (1-20).")
     p.set_defaults(func=process_videoed_denoise_image_sequence)
-    
-    def process_videoed_video_from_sequence(arguments):        
+
+    def process_videoed_video_from_sequence(arguments):
         from mainscripts import VideoEd
         VideoEd.video_from_sequence (arguments.input_dir,
-                                     arguments.output_file, 
+                                     arguments.output_file,
                                      arguments.reference_file,
                                      arguments.ext,
                                      arguments.fps,
                                      arguments.bitrate,
                                      arguments.lossless)
-                                     
+
     p = videoed_parser.add_parser( "video-from-sequence", help="Make video from image sequence.")
     p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input file to be processed. Specify .*-extension to find first file.")
     p.add_argument('--output-file', required=True, action=fixPathAction, dest="output_file", help="Input file to be processed. Specify .*-extension to find first file.")
     p.add_argument('--reference-file', action=fixPathAction, dest="reference_file", help="Reference file used to determine proper FPS and transfer audio from it. Specify .*-extension to find first file.")
     p.add_argument('--ext', dest="ext", default='png', help="Image format (extension) of input files.")
     p.add_argument('--fps', type=int, dest="fps", default=None, help="FPS of output file. Overwritten by reference-file.")
-    p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.")    
+    p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.")
     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(arguments):
         from mainscripts import LabelingTool
         LabelingTool.main (arguments.input_dir, arguments.output_dir)
-    
+
     p = subparsers.add_parser( "labelingtool", help="Labeling tool.")
     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.set_defaults(func=process_labelingtool)
-    
+
     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.")
 
     """
     Suppressing error with keras 2.2.4+ on python exit:
-    
+
         Exception ignored in: <bound method BaseSession._Callable.__del__ of <tensorflow.python.client.session.BaseSession._Callable object at 0x000000001BDEA9B0>>
         Traceback (most recent call last):
         File "D:\DeepFaceLab\_internal\bin\lib\site-packages\tensorflow\python\client\session.py", line 1413, in __del__
         AttributeError: 'NoneType' object has no attribute 'raise_exception_on_not_ok_status'
-        
+
     reproduce: https://github.com/keras-team/keras/issues/11751 ( still no solution )
     """
     outnull_file = open(os.devnull, 'w')
     os.dup2 ( outnull_file.fileno(), sys.stderr.fileno() )
     sys.stderr = outnull_file
-    
-    
+
+
 '''
 import code
 code.interact(local=dict(globals(), **locals()))
-'''
+'''

mainscripts/Converter.py

@@ -18,39 +18,39 @@ from interact import interact as io
 
 class ConvertSubprocessor(Subprocessor):
     class Cli(Subprocessor.Cli):
-        
+
         #override
         def on_initialize(self, client_dict):
             io.log_info ('Running on %s.' % (client_dict['device_name']) )
             self.device_idx  = client_dict['device_idx']
             self.device_name = client_dict['device_name']
             self.converter   = client_dict['converter']
-            self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None        
+            self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None
             self.alignments  = client_dict['alignments']
             self.debug       = client_dict['debug']
-            
+
             #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)
-                
-            from nnlib import nnlib          
+
+            from nnlib import nnlib
             #model process ate all GPU mem,
             #so we cannot use GPU for any TF operations in converter processes
             #therefore forcing active_DeviceConfig to CPU only
             nnlib.active_DeviceConfig = nnlib.DeviceConfig (cpu_only=True)
-            
+
             return None
-            
+
         #override
         def process_data(self, data):
             filename_path = Path(data)
             files_processed = 1
             faces_processed = 0
-                
+
             output_filename_path = self.output_path / (filename_path.stem + '.png')
 
-            if self.converter.type == Converter.TYPE_FACE and filename_path.stem not in self.alignments.keys():                    
+            if self.converter.type == Converter.TYPE_FACE and filename_path.stem not in self.alignments.keys():
                 if not self.debug:
                     self.log_info ( 'no faces found for %s, copying without faces' % (filename_path.name) )
                     shutil.copy ( str(filename_path), str(output_filename_path) )
@@ -72,12 +72,12 @@ class ConvertSubprocessor(Subprocessor):
                         dflimg = DFLJPG.load ( str(filename_path) )
                     else:
                         dflimg = None
-                        
+
                     if dflimg is not None:
                         image_landmarks = dflimg.get_landmarks()
-                                
+
                         image = self.converter.convert_image(image, image_landmarks, self.debug)
-                        
+
                         if self.debug:
                             raise NotImplementedError
                             #for img in image:
@@ -85,14 +85,14 @@ class ConvertSubprocessor(Subprocessor):
                             #    cv2.waitKey(0)
                         faces_processed = 1
                     else:
-                        self.log_err ("%s is not a dfl image file" % (filename_path.name) ) 
-                        
+                        self.log_err ("%s is not a dfl image file" % (filename_path.name) )
+
                 elif self.converter.type == Converter.TYPE_FACE:
                     faces = self.alignments[filename_path.stem]
-                    
+
                     if self.debug:
                         debug_images = []
-                        
+
                     for face_num, image_landmarks in enumerate(faces):
                         try:
                             if self.debug:
@@ -101,56 +101,56 @@ class ConvertSubprocessor(Subprocessor):
                             if self.debug:
                                 debug_images += self.converter.convert_face(image, image_landmarks, self.debug)
                             else:
-                                image = self.converter.convert_face(image, image_landmarks, self.debug) 
-                                
+                                image = self.converter.convert_face(image, image_landmarks, self.debug)
+
                         except Exception as e:
                             e_str = traceback.format_exc()
                             if 'MemoryError' in e_str:
                                 raise Subprocessor.SilenceException
                             else:
                                 raise Exception( 'Error while converting face_num [%d] in file [%s]: %s' % (face_num, filename_path, e_str) )
-                            
+
                     if self.debug:
                         return (1, debug_images)
-                    
+
                     faces_processed = len(faces)
-                        
+
                 if not self.debug:
                     cv2_imwrite (str(output_filename_path), (image*255).astype(np.uint8) )
-                
-                
+
+
             return (0, files_processed, faces_processed)
-            
+
         #overridable
         def get_data_name (self, data):
             #return string identificator of your data
             return data
-        
-    #override
-    def __init__(self, converter, input_path_image_paths, output_path, alignments, debug = False): 
-        super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if debug == True else 60)    
 
-        self.converter = converter        
+    #override
+    def __init__(self, converter, input_path_image_paths, output_path, alignments, debug = False):
+        super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if debug == True else 60)
+
+        self.converter = converter
         self.host_processor, self.cli_func = SubprocessFunctionCaller.make_pair ( self.converter.predictor_func )
         self.process_converter = self.converter.copy_and_set_predictor(self.cli_func)
-        
+
         self.input_data = self.input_path_image_paths = input_path_image_paths
         self.output_path = output_path
         self.alignments = alignments
         self.debug = debug
-        
+
         self.files_processed = 0
         self.faces_processed = 0
-        
+
     #override
     def process_info_generator(self):
         r = [0] if self.debug else range(multiprocessing.cpu_count())
 
         for i in r:
             yield 'CPU%d' % (i), {}, {'device_idx': i,
-                                      'device_name': 'CPU%d' % (i), 
-                                      'converter' : self.process_converter, 
-                                      'output_dir' : str(self.output_path), 
+                                      'device_name': 'CPU%d' % (i),
+                                      'converter' : self.process_converter,
+                                      'output_dir' : str(self.output_path),
                                       'alignments' : self.alignments,
                                       'debug': self.debug,
                                       'stdin_fd': sys.stdin.fileno() if self.debug else None
@@ -160,25 +160,25 @@ class ConvertSubprocessor(Subprocessor):
     def on_clients_initialized(self):
         if self.debug:
             io.named_window ("Debug convert")
-        
+
         io.progress_bar ("Converting", len (self.input_data) )
-        
+
     #overridable optional
     def on_clients_finalized(self):
         io.progress_bar_close()
-        
+
         if self.debug:
             io.destroy_all_windows()
-        
+
     #override
     def get_data(self, host_dict):
         if len (self.input_data) > 0:
-            return self.input_data.pop(0)            
+            return self.input_data.pop(0)
         return None
-    
+
     #override
     def on_data_return (self, host_dict, data):
-        self.input_data.insert(0, data)   
+        self.input_data.insert(0, data)
 
     #override
     def on_result (self, host_dict, data, result):
@@ -190,25 +190,25 @@ class ConvertSubprocessor(Subprocessor):
                 io.show_image ('Debug convert', (img*255).astype(np.uint8) )
                 io.wait_any_key()
         io.progress_bar_inc(1)
-    
+
     #override
     def on_tick(self):
         self.host_processor.process_messages()
-        
+
     #override
     def get_result(self):
         return self.files_processed, self.faces_processed
-        
+
 def main (args, device_args):
     io.log_info ("Running converter.\r\n")
-    
+
     aligned_dir = args.get('aligned_dir', None)
-    
+
     try:
         input_path = Path(args['input_dir'])
         output_path = Path(args['output_dir'])
         model_path = Path(args['model_dir'])
-        
+
         if not input_path.exists():
             io.log_err('Input directory not found. Please ensure it exists.')
             return
@@ -218,69 +218,69 @@ def main (args, device_args):
                 Path(filename).unlink()
         else:
             output_path.mkdir(parents=True, exist_ok=True)
-         
+
         if not model_path.exists():
             io.log_err('Model directory not found. Please ensure it exists.')
             return
-            
-        import models 
+
+        import models
         model = models.import_model( args['model_name'] )(model_path, device_args=device_args)
         converter = model.get_converter()
         converter.dummy_predict()
 
         alignments = None
-        
+
         if converter.type == Converter.TYPE_FACE:
             if aligned_dir is None:
                 io.log_err('Aligned directory not found. Please ensure it exists.')
-                return 
-                
+                return
+
             aligned_path = Path(aligned_dir)
             if not aligned_path.exists():
                 io.log_err('Aligned directory not found. Please ensure it exists.')
-                return 
-                
+                return
+
             alignments = {}
-            
+
             aligned_path_image_paths = Path_utils.get_image_paths(aligned_path)
             for filepath in io.progress_bar_generator(aligned_path_image_paths, "Collecting alignments"):
                 filepath = Path(filepath)
-                
+
                 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) ) 
+                    io.log_err ("%s is not a dfl image file" % (filepath.name) )
                     continue
-                    
+
                 source_filename_stem = Path( dflimg.get_source_filename() ).stem
                 if source_filename_stem not in alignments.keys():
                     alignments[ source_filename_stem ] = []
 
                 alignments[ source_filename_stem ].append (dflimg.get_source_landmarks())
 
-        files_processed, faces_processed = ConvertSubprocessor ( 
+        files_processed, faces_processed = ConvertSubprocessor (
                     converter              = converter,
-                    input_path_image_paths = Path_utils.get_image_paths(input_path), 
+                    input_path_image_paths = Path_utils.get_image_paths(input_path),
                     output_path            = output_path,
                     alignments             = alignments,
                     debug                  = args.get('debug',False)
                     ).run()
 
         model.finalize()
-        
+
     except Exception as e:
         print ( 'Error: %s' % (str(e)))
         traceback.print_exc()
-        
-'''            
+
+'''
 if model_name == 'AVATAR':
     output_path_image_paths = Path_utils.get_image_paths(output_path)
-    
+
     last_ok_frame = -1
     for filename in output_path_image_paths:
         filename_path = Path(filename)
@@ -289,15 +289,15 @@ if model_name == 'AVATAR':
             frame = int(stem)
         except:
             raise Exception ('Aligned avatars must be created from indexed sequence files.')
-            
+
         if frame-last_ok_frame > 1:
             start = last_ok_frame + 1
             end = frame - 1
-            
+
             print ("Filling gaps: [%d...%d]" % (start, end) )
-            for i in range (start, end+1):                    
+            for i in range (start, end+1):
                 shutil.copy ( str(filename), str( output_path / ('%.5d%s' % (i, filename_path.suffix ))  ) )
-            
+
         last_ok_frame = frame
 '''
 #interpolate landmarks
@@ -306,28 +306,28 @@ if model_name == 'AVATAR':
 #a = sorted(alignments.keys())
 #a_len = len(a)
 #
-#box_pts = 3            
+#box_pts = 3
 #box = np.ones(box_pts)/box_pts
 #for i in range( a_len ):
 #    if i >= box_pts and i <= a_len-box_pts-1:
 #        af0 = alignments[ a[i] ][0] ##first face
-#        m0 = LandmarksProcessor.get_transform_mat (af0, 256, face_type=FaceType.FULL)      
-#    
+#        m0 = LandmarksProcessor.get_transform_mat (af0, 256, face_type=FaceType.FULL)
+#
 #        points = []
-#        
+#
 #        for j in range(-box_pts, box_pts+1):
 #            af = alignments[ a[i+j] ][0] ##first face
-#            m = LandmarksProcessor.get_transform_mat (af, 256, face_type=FaceType.FULL)                    
+#            m = LandmarksProcessor.get_transform_mat (af, 256, face_type=FaceType.FULL)
 #            p = LandmarksProcessor.transform_points (af, m)
 #            points.append (p)
-#    
+#
 #        points = np.array(points)
 #        points_len = len(points)
 #        t_points = np.transpose(points, [1,0,2])
-#        
+#
 #        p1 = np.array ( [ int(np.convolve(x[:,0], box, mode='same')[points_len//2]) for x in t_points ] )
 #        p2 = np.array ( [ int(np.convolve(x[:,1], box, mode='same')[points_len//2]) for x in t_points ] )
-#        
+#
 #        new_points = np.concatenate( [np.expand_dims(p1,-1),np.expand_dims(p2,-1)], -1 )
-#        
+#
 #        alignments[ a[i] ][0]  = LandmarksProcessor.transform_points (new_points, m0, True).astype(np.int32)

mainscripts/Extractor.py

@@ -18,9 +18,9 @@ from facelib import LandmarksProcessor
 from nnlib import nnlib
 from joblib import Subprocessor
 from interact import interact as io
-        
+
 class ExtractSubprocessor(Subprocessor):
-        
+
     class Cli(Subprocessor.Cli):
 
         #override
@@ -32,19 +32,19 @@ class ExtractSubprocessor(Subprocessor):
             self.face_type    = client_dict['face_type']
             self.device_idx   = client_dict['device_idx']
             self.cpu_only     = client_dict['device_type'] == 'CPU'
-            self.output_path  = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None        
+            self.output_path  = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None
             self.debug_dir    = client_dict['debug_dir']
             self.detector     = client_dict['detector']
-  
+
             self.cached_image = (None, None)
-            
+
             self.e = None
             device_config = nnlib.DeviceConfig ( cpu_only=self.cpu_only, force_gpu_idx=self.device_idx, allow_growth=True)
             if self.type == 'rects':
                 if self.detector is not None:
                     if self.detector == 'mt':
                         nnlib.import_all (device_config)
-                        self.e = facelib.MTCExtractor()                            
+                        self.e = facelib.MTCExtractor()
                     elif self.detector == 'dlib':
                         nnlib.import_dlib (device_config)
                         self.e = facelib.DLIBExtractor(nnlib.dlib)
@@ -53,10 +53,10 @@ class ExtractSubprocessor(Subprocessor):
                         self.e = facelib.S3FDExtractor()
                     else:
                         raise ValueError ("Wrong detector type.")
-                        
+
                     if self.e is not None:
                         self.e.__enter__()
-                    
+
             elif self.type == 'landmarks':
                 nnlib.import_all (device_config)
                 self.e = facelib.LandmarksExtractor(nnlib.keras)
@@ -66,15 +66,15 @@ class ExtractSubprocessor(Subprocessor):
                     self.second_pass_e.__enter__()
                 else:
                     self.second_pass_e = None
-                
+
             elif self.type == 'final':
                 pass
-                
+
         #override
         def on_finalize(self):
             if self.e is not None:
                 self.e.__exit__()
-                
+
         #override
         def process_data(self, data):
             filename_path = Path( data[0] )
@@ -84,64 +84,64 @@ class ExtractSubprocessor(Subprocessor):
                 image = self.cached_image[1] #cached image for manual extractor
             else:
                 image = cv2_imread( filename_path_str )
-                
+
                 if image is None:
                     self.log_err ( 'Failed to extract %s, reason: cv2_imread() fail.' % ( str(filename_path) ) )
                     return None
-                
+
                 image_shape = image.shape
                 if len(image_shape) == 2:
                     h, w = image.shape
-                    ch = 1                    
+                    ch = 1
                 else:
                     h, w, ch = image.shape
-                
+
                 if ch == 1:
                     image = np.repeat ( image [:,:,np.newaxis], 3, -1 )
                 elif ch == 4:
                     image = image[:,:,0:3]
-                    
+
                 wm = w % 2
                 hm = h % 2
                 if wm + hm != 0: #fix odd image
                     image = image[0:h-hm,0:w-wm,:]
                 self.cached_image = ( filename_path_str, image )
-       
+
             src_dflimg = None
-            h, w, ch = image.shape                
+            h, w, ch = image.shape
             if h == w:
                 #extracting from already extracted jpg image?
                 if filename_path.suffix == '.jpg':
                     src_dflimg = DFLJPG.load ( str(filename_path) )
-       
+
             if self.type == 'rects':
                 if min(w,h) < 128:
                     self.log_err ( 'Image is too small %s : [%d, %d]' % ( str(filename_path), w, h ) )
                     rects = []
-                else:                    
+                else:
                     rects = self.e.extract_from_bgr (image)
-                    
+
                 return [str(filename_path), rects]
 
             elif self.type == 'landmarks':
                 rects = data[1]
                 if rects is None:
                     landmarks = None
-                else:                    
-                    landmarks = self.e.extract_from_bgr (image, rects, self.second_pass_e if src_dflimg is None else None)                    
-                    
+                else:
+                    landmarks = self.e.extract_from_bgr (image, rects, self.second_pass_e if src_dflimg is None else None)
+
                 return [str(filename_path), landmarks]
 
             elif self.type == 'final':
-                
-            
+
+
                 result = []
                 faces = data[1]
-                
+
                 if self.debug_dir is not None:
                     debug_output_file = str( Path(self.debug_dir) / (filename_path.stem+'.jpg') )
                     debug_image = image.copy()
-                 
+
                 if src_dflimg is not None and len(faces) != 1:
                     #if re-extracting from dflimg and more than 1 or zero faces detected - dont process and just copy it
                     print("src_dflimg is not None and len(faces) != 1", str(filename_path) )
@@ -151,26 +151,26 @@ class ExtractSubprocessor(Subprocessor):
                     result.append (output_file)
                 else:
                     face_idx = 0
-                    for face in faces:   
+                    for face in faces:
                         rect = np.array(face[0])
                         image_landmarks = face[1]
                         if image_landmarks is None:
                             continue
                         image_landmarks = np.array(image_landmarks)
 
-                        if self.face_type == FaceType.MARK_ONLY:                        
+                        if self.face_type == FaceType.MARK_ONLY:
                             face_image = image
                             face_image_landmarks = image_landmarks
                         else:
-                            image_to_face_mat = LandmarksProcessor.get_transform_mat (image_landmarks, self.image_size, self.face_type)       
+                            image_to_face_mat = LandmarksProcessor.get_transform_mat (image_landmarks, self.image_size, self.face_type)
                             face_image = cv2.warpAffine(image, image_to_face_mat, (self.image_size, self.image_size), cv2.INTER_LANCZOS4)
                             face_image_landmarks = LandmarksProcessor.transform_points (image_landmarks, image_to_face_mat)
-                            
+
                             landmarks_bbox = LandmarksProcessor.transform_points ( [ (0,0), (0,self.image_size-1), (self.image_size-1, self.image_size-1), (self.image_size-1,0) ], image_to_face_mat, True)
-                            
+
                             rect_area      = mathlib.polygon_area(np.array(rect[[0,2,2,0]]), np.array(rect[[1,1,3,3]]))
                             landmarks_area = mathlib.polygon_area(landmarks_bbox[:,0], landmarks_bbox[:,1] )
-                            
+
                             if landmarks_area > 4*rect_area: #get rid of faces which umeyama-landmark-area > 4*detector-rect-area
                                 continue
 
@@ -192,24 +192,24 @@ class ExtractSubprocessor(Subprocessor):
                                                        source_rect=rect,
                                                        source_landmarks=image_landmarks.tolist(),
                                                        image_to_face_mat=image_to_face_mat
-                                            )  
-                            
+                                            )
+
                         result.append (output_file)
                         face_idx += 1
-                        
+
                 if self.debug_dir is not None:
                     cv2_imwrite(debug_output_file, debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50] )
-                    
-                return result       
-            
-             
+
+                return result
+
+
         #overridable
         def get_data_name (self, data):
             #return string identificator of your data
             return data[0]
-            
+
     #override
-    def __init__(self, input_data, type, image_size, face_type, debug_dir, multi_gpu=False, cpu_only=False, manual=False, manual_window_size=0, detector=None, output_path=None): 
+    def __init__(self, input_data, type, image_size, face_type, debug_dir, multi_gpu=False, cpu_only=False, manual=False, manual_window_size=0, detector=None, output_path=None):
         self.input_data = input_data
         self.type = type
         self.image_size = image_size
@@ -218,8 +218,8 @@ class ExtractSubprocessor(Subprocessor):
         self.multi_gpu = multi_gpu
         self.cpu_only = cpu_only
         self.detector = detector
-        self.output_path = output_path        
-        self.manual = manual        
+        self.output_path = output_path
+        self.manual = manual
         self.manual_window_size = manual_window_size
         self.result = []
 
@@ -233,32 +233,32 @@ class ExtractSubprocessor(Subprocessor):
             io.named_window(self.wnd_name)
             io.capture_mouse(self.wnd_name)
             io.capture_keys(self.wnd_name)
-            
+
             self.cache_original_image = (None, None)
             self.cache_image = (None, None)
             self.cache_text_lines_img = (None, None)
             self.hide_help = False
-            
+
             self.landmarks = None
             self.x = 0
             self.y = 0
             self.rect_size = 100
             self.rect_locked = False
             self.extract_needed = True
-            
+
         io.progress_bar (None, len (self.input_data))
-            
+
     #override
     def on_clients_finalized(self):
         if self.manual == True:
             io.destroy_all_windows()
-            
+
         io.progress_bar_close()
-        
+
     def get_devices_for_type (self, type, multi_gpu, cpu_only):
         if 'cpu' in nnlib.device.backend:
             cpu_only = True
-            
+
         if not cpu_only and (type == 'rects' or type == 'landmarks'):
             if type == 'rects' and (self.detector == 'mt') and nnlib.device.backend == "plaidML":
                 cpu_only = True
@@ -269,11 +269,11 @@ class ExtractSubprocessor(Subprocessor):
                     devices = [nnlib.device.getBestValidDeviceIdx()]
                 if len(devices) == 0:
                     devices = [0]
-                    
+
                 for idx in devices:
                     dev_name = nnlib.device.getDeviceName(idx)
                     dev_vram = nnlib.device.getDeviceVRAMTotalGb(idx)
-                    
+
                     if not self.manual and ( self.type == 'rects' and self.detector != 's3fd' ):
                         for i in range ( int (max (1, dev_vram / 2) ) ):
                             yield (idx, 'GPU', '%s #%d' % (dev_name,i) , dev_vram)
@@ -286,21 +286,21 @@ class ExtractSubprocessor(Subprocessor):
             else:
                 for i in range( min(8, multiprocessing.cpu_count() // 2) ):
                     yield (i, 'CPU', 'CPU%d' % (i), 0 )
-                
+
         if type == 'final':
             for i in range( min(8, multiprocessing.cpu_count()) ):
-                yield (i, 'CPU', 'CPU%d' % (i), 0 ) 
-        
+                yield (i, 'CPU', 'CPU%d' % (i), 0 )
+
     #override
     def process_info_generator(self):
-        base_dict = {'type' : self.type, 
-                     'image_size': self.image_size, 
-                     'face_type': self.face_type, 
-                     'debug_dir': self.debug_dir, 
-                     'output_dir': str(self.output_path), 
+        base_dict = {'type' : self.type,
+                     'image_size': self.image_size,
+                     'face_type': self.face_type,
+                     'debug_dir': self.debug_dir,
+                     'output_dir': str(self.output_path),
                      'detector': self.detector}
-    
-        for (device_idx, device_type, device_name, device_total_vram_gb) in self.get_devices_for_type(self.type, self.multi_gpu, self.cpu_only): 
+
+        for (device_idx, device_type, device_name, device_total_vram_gb) in self.get_devices_for_type(self.type, self.multi_gpu, self.cpu_only):
             client_dict = base_dict.copy()
             client_dict['device_idx'] = device_idx
             client_dict['device_name'] = device_name
@@ -311,7 +311,7 @@ class ExtractSubprocessor(Subprocessor):
     def get_data(self, host_dict):
         if not self.manual:
             if len (self.input_data) > 0:
-                return self.input_data.pop(0)    
+                return self.input_data.pop(0)
         else:
 
             need_remark_face = False
@@ -327,7 +327,7 @@ class ExtractSubprocessor(Subprocessor):
                         self.rect, self.landmarks = faces.pop()
                         faces.clear()
                         redraw_needed = True
-                        self.rect_locked = True                        
+                        self.rect_locked = True
                         self.rect_size = ( self.rect[2] - self.rect[0] ) / 2
                         self.x = ( self.rect[0] + self.rect[2] ) / 2
                         self.y = ( self.rect[1] + self.rect[3] ) / 2
@@ -338,19 +338,19 @@ class ExtractSubprocessor(Subprocessor):
                     else:
                         self.original_image = cv2_imread( filename )
                         self.cache_original_image = (filename, self.original_image )
-                    
+
                     (h,w,c) = self.original_image.shape
                     self.view_scale = 1.0 if self.manual_window_size == 0 else self.manual_window_size / ( h * (16.0/9.0) )
-                    
+
                     if self.cache_image[0] == (h,w,c) + (self.view_scale,filename):
                         self.image = self.cache_image[1]
-                    else:                                            
-                        self.image = cv2.resize (self.original_image, ( int(w*self.view_scale), int(h*self.view_scale) ), interpolation=cv2.INTER_LINEAR)    
+                    else:
+                        self.image = cv2.resize (self.original_image, ( int(w*self.view_scale), int(h*self.view_scale) ), interpolation=cv2.INTER_LINEAR)
                         self.cache_image = ( (h,w,c) + (self.view_scale,filename), self.image )
- 
+
                     (h,w,c) = self.image.shape
-                    
-                    sh = (0,0, w, min(100, h) )                    
+
+                    sh = (0,0, w, min(100, h) )
                     if self.cache_text_lines_img[0] == sh:
                         self.text_lines_img = self.cache_text_lines_img[1]
                     else:
@@ -362,30 +362,30 @@ class ExtractSubprocessor(Subprocessor):
                                                             '[,] [.]- prev frame, next frame. [Q] - skip remaining frames',
                                                             '[h] - hide this help'
                                                         ], (1, 1, 1) )*255).astype(np.uint8)
-                                                        
+
                         self.cache_text_lines_img = (sh, self.text_lines_img)
 
                     while True:
                         io.process_messages(0.0001)
-                        
+
                         new_x = self.x
                         new_y = self.y
                         new_rect_size = self.rect_size
-                        
+
                         mouse_events = io.get_mouse_events(self.wnd_name)
                         for ev in mouse_events:
                             (x, y, ev, flags) = ev
                             if ev == io.EVENT_MOUSEWHEEL and not self.rect_locked:
-                                mod = 1 if flags > 0 else -1 
+                                mod = 1 if flags > 0 else -1
                                 diff = 1 if new_rect_size <= 40 else np.clip(new_rect_size / 10, 1, 10)
-                                new_rect_size = max (5, new_rect_size + diff*mod)                    
+                                new_rect_size = max (5, new_rect_size + diff*mod)
                             elif ev == io.EVENT_LBUTTONDOWN:
                                 self.rect_locked = not self.rect_locked
                                 self.extract_needed = True
                             elif not self.rect_locked:
                                 new_x = np.clip (x, 0, w-1) / self.view_scale
                                 new_y = np.clip (y, 0, h-1) / self.view_scale
-                                
+
                         key_events = io.get_key_events(self.wnd_name)
                         key, = key_events[-1] if len(key_events) > 0 else (0,)
 
@@ -393,48 +393,48 @@ class ExtractSubprocessor(Subprocessor):
                             #confirm frame
                             is_frame_done = True
                             faces.append ( [(self.rect), self.landmarks] )
-                            
+
                             break
                         elif key == ord(' '):
                             #confirm skip frame
                             is_frame_done = True
                             break
-                        elif key == ord(',')  and len(self.result) > 0: 
-                            #go prev frame      
-                            
+                        elif key == ord(',')  and len(self.result) > 0:
+                            #go prev frame
+
                             if self.rect_locked:
                                 # Only save the face if the rect is still locked
                                 faces.append ( [(self.rect), self.landmarks] )
-                            
+
                             self.input_data.insert(0, self.result.pop() )
                             io.progress_bar_inc(-1)
                             need_remark_face = True
-                    
+
                             break
-                        elif key == ord('.'): 
-                            #go next frame   
-                            
+                        elif key == ord('.'):
+                            #go next frame
+
                             if self.rect_locked:
                                 # Only save the face if the rect is still locked
                                 faces.append ( [(self.rect), self.landmarks] )
                             need_remark_face = True
                             is_frame_done = True
-                            break                        
+                            break
                         elif key == ord('q'):
                             #skip remaining
-                            
+
                             if self.rect_locked:
                                 faces.append ( [(self.rect), self.landmarks] )
                             while len(self.input_data) > 0:
                                 self.result.append( self.input_data.pop(0) )
                                 io.progress_bar_inc(1)
-                            
+
                             break
-                            
+
                         elif key == ord('h'):
                             self.hide_help = not self.hide_help
                             break
-                        
+
                         if self.x != new_x or \
                            self.y != new_y or \
                            self.rect_size != new_rect_size or \
@@ -443,33 +443,33 @@ class ExtractSubprocessor(Subprocessor):
                             self.x = new_x
                             self.y = new_y
                             self.rect_size = new_rect_size
-                            self.rect = ( int(self.x-self.rect_size), 
-                                          int(self.y-self.rect_size), 
-                                          int(self.x+self.rect_size), 
+                            self.rect = ( int(self.x-self.rect_size),
+                                          int(self.y-self.rect_size),
+                                          int(self.x+self.rect_size),
                                           int(self.y+self.rect_size) )
-                                          
+
                             if redraw_needed:
                                 redraw_needed = False
                                 return [filename, None]
                             else:
                                 return [filename, [self.rect]]
-       
+
                 else:
                     is_frame_done = True
-                    
+
                 if is_frame_done:
                     self.result.append ( data )
                     self.input_data.pop(0)
                     io.progress_bar_inc(1)
                     self.extract_needed = True
-                    self.rect_locked = False                    
+                    self.rect_locked = False
 
         return None
-    
+
     #override
     def on_data_return (self, host_dict, data):
         if not self.manual:
-            self.input_data.insert(0, data)   
+            self.input_data.insert(0, data)
 
     #override
     def on_result (self, host_dict, data, result):
@@ -477,33 +477,33 @@ class ExtractSubprocessor(Subprocessor):
             filename, landmarks = result
             if landmarks is not None:
                 self.landmarks = landmarks[0][1]
-                                        
+
             (h,w,c) = self.image.shape
-            
+
             if not self.hide_help:
                 image = cv2.addWeighted (self.image,1.0,self.text_lines_img,1.0,0)
             else:
                 image = self.image.copy()
-                
+
             view_rect = (np.array(self.rect) * self.view_scale).astype(np.int).tolist()
             view_landmarks  = (np.array(self.landmarks) * self.view_scale).astype(np.int).tolist()
-            
+
             if self.rect_size <= 40:
                 scaled_rect_size = h // 3 if w > h else w // 3
 
                 p1 = (self.x - self.rect_size, self.y - self.rect_size)
                 p2 = (self.x + self.rect_size, self.y - self.rect_size)
                 p3 = (self.x - self.rect_size, self.y + self.rect_size)
-                
-                wh = h if h < w else w                
+
+                wh = h if h < w else w
                 np1 = (w / 2 - wh / 4, h / 2 - wh / 4)
                 np2 = (w / 2 + wh / 4, h / 2 - wh / 4)
                 np3 = (w / 2 - wh / 4, h / 2 + wh / 4)
-                
+
                 mat = cv2.getAffineTransform( np.float32([p1,p2,p3])*self.view_scale, np.float32([np1,np2,np3]) )
-                image = cv2.warpAffine(image, mat,(w,h) )                
+                image = cv2.warpAffine(image, mat,(w,h) )
                 view_landmarks = LandmarksProcessor.transform_points (view_landmarks, mat)
-     
+
             landmarks_color = (255,255,0) if self.rect_locked else (0,255,0)
             LandmarksProcessor.draw_rect_landmarks (image, view_rect, view_landmarks, self.image_size, self.face_type, landmarks_color=landmarks_color)
             self.extract_needed = False
@@ -513,10 +513,10 @@ class ExtractSubprocessor(Subprocessor):
             if self.type == 'rects':
                 self.result.append ( result )
             elif self.type == 'landmarks':
-                self.result.append ( result )                        
+                self.result.append ( result )
             elif self.type == 'final':
                 self.result += result
-                         
+
             io.progress_bar_inc(1)
 
     #override
@@ -530,47 +530,47 @@ class DeletedFilesSearcherSubprocessor(Subprocessor):
         def on_initialize(self, client_dict):
             self.debug_paths_stems = client_dict['debug_paths_stems']
             return None
-            
+
         #override
-        def process_data(self, data):  
-            input_path_stem = Path(data[0]).stem        
+        def process_data(self, data):
+            input_path_stem = Path(data[0]).stem
             return any ( [ input_path_stem == d_stem for d_stem in self.debug_paths_stems] )
-            
+
         #override
         def get_data_name (self, data):
             #return string identificator of your data
             return data[0]
-            
+
     #override
-    def __init__(self, input_paths, debug_paths ): 
+    def __init__(self, input_paths, debug_paths ):
         self.input_paths = input_paths
-        self.debug_paths_stems = [ Path(d).stem for d in debug_paths]        
+        self.debug_paths_stems = [ Path(d).stem for d in debug_paths]
         self.result = []
-        super().__init__('DeletedFilesSearcherSubprocessor', DeletedFilesSearcherSubprocessor.Cli, 60)           
-        
+        super().__init__('DeletedFilesSearcherSubprocessor', DeletedFilesSearcherSubprocessor.Cli, 60)
+
     #override
-    def process_info_generator(self):    
+    def process_info_generator(self):
         for i in range(min(multiprocessing.cpu_count(), 8)):
             yield 'CPU%d' % (i), {}, {'debug_paths_stems' : self.debug_paths_stems}
 
     #override
     def on_clients_initialized(self):
         io.progress_bar ("Searching deleted files", len (self.input_paths))
-            
+
     #override
     def on_clients_finalized(self):
         io.progress_bar_close()
 
     #override
     def get_data(self, host_dict):
-        if len (self.input_paths) > 0:        
-            return [self.input_paths.pop(0)]        
+        if len (self.input_paths) > 0:
+            return [self.input_paths.pop(0)]
         return None
-    
+
     #override
     def on_data_return (self, host_dict, data):
-        self.input_paths.insert(0, data[0])   
-     
+        self.input_paths.insert(0, data[0])
+
     #override
     def on_result (self, host_dict, data, result):
         if result == False:
@@ -591,40 +591,40 @@ def main(input_dir,
          image_size=256,
          face_type='full_face',
          device_args={}):
-         
+
     input_path = Path(input_dir)
     output_path = Path(output_dir)
     face_type = FaceType.fromString(face_type)
-    
+
     multi_gpu = device_args.get('multi_gpu', False)
     cpu_only = device_args.get('cpu_only', False)
-       
+
     if not input_path.exists():
         raise ValueError('Input directory not found. Please ensure it exists.')
-    
+
     if output_path.exists():
         if not manual_output_debug_fix and input_path != output_path:
             for filename in Path_utils.get_image_paths(output_path):
                 Path(filename).unlink()
     else:
         output_path.mkdir(parents=True, exist_ok=True)
-        
+
     if manual_output_debug_fix:
         if debug_dir is None:
             raise ValueError('debug-dir must be specified')
         detector = 'manual'
         io.log_info('Performing re-extract frames which were deleted from _debug directory.')
-        
+
     input_path_image_paths = Path_utils.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info)
     if debug_dir is not None:
         debug_output_path = Path(debug_dir)
-        
+
         if manual_output_debug_fix:
             if not debug_output_path.exists():
                 raise ValueError("%s not found " % ( str(debug_output_path) ))
 
             input_path_image_paths = DeletedFilesSearcherSubprocessor (input_path_image_paths, Path_utils.get_image_paths(debug_output_path) ).run()
-            input_path_image_paths = sorted (input_path_image_paths)            
+            input_path_image_paths = sorted (input_path_image_paths)
         else:
             if debug_output_path.exists():
                 for filename in Path_utils.get_image_paths(debug_output_path):
@@ -634,20 +634,20 @@ def main(input_dir,
 
     images_found = len(input_path_image_paths)
     faces_detected = 0
-    if images_found != 0:    
+    if images_found != 0:
         if detector == 'manual':
             io.log_info ('Performing manual extract...')
             extracted_faces = ExtractSubprocessor ([ (filename,[]) for filename in input_path_image_paths ], 'landmarks', image_size, face_type, debug_dir, cpu_only=cpu_only, manual=True, manual_window_size=manual_window_size).run()
         else:
             io.log_info ('Performing 1st pass...')
             extracted_rects = ExtractSubprocessor ([ (x,) for x in input_path_image_paths ], 'rects', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, detector=detector).run()
-                
+
             io.log_info ('Performing 2nd pass...')
             extracted_faces = ExtractSubprocessor (extracted_rects, 'landmarks', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False).run()
-                
+
             if manual_fix:
                 io.log_info ('Performing manual fix...')
-                
+
                 if all ( np.array ( [ len(data[1]) > 0 for data in extracted_faces] ) == True ):
                     io.log_info ('All faces are detected, manual fix not needed.')
                 else:
@@ -657,8 +657,8 @@ def main(input_dir,
             io.log_info ('Performing 3rd pass...')
             final_imgs_paths = ExtractSubprocessor (extracted_faces, 'final', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, output_path=output_path).run()
             faces_detected = len(final_imgs_paths)
-            
+
     io.log_info ('-------------------------')
     io.log_info ('Images found:        %d' % (images_found) )
     io.log_info ('Faces detected:      %d' % (faces_detected) )
-    io.log_info ('-------------------------')
+    io.log_info ('-------------------------')

mainscripts/LabelingTool_unfinished.py

@@ -17,18 +17,18 @@ 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)
@@ -39,165 +39,165 @@ def main(input_dir, output_dir):
             dflimg = DFLJPG.load ( str(filepath) )
         else:
             dflimg = None
-            
+
         if dflimg is None:
-            io.log_err ("%s is not a dfl image file" % (filepath.name) ) 
+            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]                    
+                    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 = 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] ) 
-        
+        #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] ) 
-        
+        #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] ) 
-        
+        draw_fill_convex ( mask_fgp, lmrks[27:36] )
+
         #convex inner of nose half
-        draw_fill_convex ( mask_fg, lmrks[27:36], scale=0.5 ) 
-        
-        
+        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] )        
+        #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) 
-        
+        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_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
@@ -211,34 +211,34 @@ def main(input_dir, output_dir):
         #
         #
         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] 
-        
+        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 = 
+
+        ##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-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)
@@ -246,25 +246,25 @@ def main(input_dir, output_dir):
         ##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,                                  
+                                  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'):
@@ -273,15 +273,15 @@ def main(input_dir, output_dir):
                     pass
                 if key == ord('3'):
                     pass
-                   
-            if ord(' ') in key_events:                    
+
+            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)
+#cv2.drawContours(original_mask, points, -1, (1,), 1)

mainscripts/Sorter.py

@@ -15,10 +15,10 @@ from joblib import Subprocessor
 import multiprocessing
 from interact import interact as io
 from imagelib import estimate_sharpness
-    
+
 class BlurEstimatorSubprocessor(Subprocessor):
     class Cli(Subprocessor.Cli):
-    
+
         #override
         def on_initialize(self, client_dict):
             self.log_info('Running on %s.' % (client_dict['device_name']) )
@@ -26,58 +26,58 @@ class BlurEstimatorSubprocessor(Subprocessor):
         #override
         def process_data(self, data):
             filepath = Path( data[0] )
-        
+
             if filepath.suffix == '.png':
                 dflimg = DFLPNG.load( str(filepath) )
             elif filepath.suffix == '.jpg':
                 dflimg = DFLJPG.load ( str(filepath) )
             else:
                 dflimg = None
-                
+
             if dflimg is not None:
                 image = cv2_imread( str(filepath) )
                 return [ str(filepath), estimate_sharpness(image) ]
             else:
-                self.log_err ("%s is not a dfl image file" % (filepath.name) ) 
+                self.log_err ("%s is not a dfl image file" % (filepath.name) )
                 return [ str(filepath), 0 ]
 
         #override
         def get_data_name (self, data):
             #return string identificator of your data
             return data[0]
-            
+
     #override
-    def __init__(self, input_data ): 
+    def __init__(self, input_data ):
         self.input_data = input_data
         self.img_list = []
         self.trash_img_list = []
-        super().__init__('BlurEstimator', BlurEstimatorSubprocessor.Cli, 60)           
+        super().__init__('BlurEstimator', BlurEstimatorSubprocessor.Cli, 60)
 
     #override
     def on_clients_initialized(self):
         io.progress_bar ("", len (self.input_data))
-        
+
     #override
     def on_clients_finalized(self):
         io.progress_bar_close ()
-        
+
     #override
-    def process_info_generator(self):    
+    def process_info_generator(self):
         for i in range(0, multiprocessing.cpu_count() ):
             yield 'CPU%d' % (i), {}, {'device_idx': i,
-                                      'device_name': 'CPU%d' % (i), 
+                                      'device_name': 'CPU%d' % (i),
                                       }
 
     #override
     def get_data(self, host_dict):
         if len (self.input_data) > 0:
-            return self.input_data.pop(0)    
-        
+            return self.input_data.pop(0)
+
         return None
-    
+
     #override
     def on_data_return (self, host_dict, data):
-        self.input_data.insert(0, data)   
+        self.input_data.insert(0, data)
 
     #override
     def on_result (self, host_dict, data, result):
@@ -85,20 +85,20 @@ class BlurEstimatorSubprocessor(Subprocessor):
             self.trash_img_list.append ( result )
         else:
             self.img_list.append ( result )
-            
+
         io.progress_bar_inc(1)
-    
+
     #override
     def get_result(self):
         return self.img_list, self.trash_img_list
 
-    
+
 def sort_by_blur(input_path):
-    io.log_info ("Sorting by blur...")        
-    
-    img_list = [ (filename,[]) for filename in Path_utils.get_image_paths(input_path) ] 
+    io.log_info ("Sorting by blur...")
+
+    img_list = [ (filename,[]) for filename in Path_utils.get_image_paths(input_path) ]
     img_list, trash_img_list = BlurEstimatorSubprocessor (img_list).run()
-    
+
     io.log_info ("Sorting...")
     img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
 
@@ -111,21 +111,21 @@ def sort_by_face(input_path):
     trash_img_list = []
     for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"):
         filepath = Path(filepath)
-        
+
         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) )
             trash_img_list.append ( [str(filepath)] )
             continue
 
         img_list.append( [str(filepath), dflimg.get_landmarks()] )
-        
+
 
     img_list_len = len(img_list)
     for i in io.progress_bar_generator ( range(0, img_list_len-1), "Sorting"):
@@ -152,21 +152,21 @@ def sort_by_face_dissim(input_path):
     trash_img_list = []
     for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"):
         filepath = Path(filepath)
-        
+
         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) )
             trash_img_list.append ( [str(filepath)] )
-            continue 
-        
+            continue
+
         img_list.append( [str(filepath), dflimg.get_landmarks(), 0 ] )
-        
+
     img_list_len = len(img_list)
     for i in io.progress_bar_generator( range(img_list_len-1), "Sorting"):
         score_total = 0
@@ -183,79 +183,79 @@ def sort_by_face_dissim(input_path):
     img_list = sorted(img_list, key=operator.itemgetter(2), reverse=True)
 
     return img_list, trash_img_list
-    
+
 def sort_by_face_yaw(input_path):
     io.log_info ("Sorting by face yaw...")
     img_list = []
     trash_img_list = []
     for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"):
         filepath = Path(filepath)
-        
+
         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) )
             trash_img_list.append ( [str(filepath)] )
             continue
-        
+
         pitch, yaw = LandmarksProcessor.estimate_pitch_yaw ( dflimg.get_landmarks() )
-       
+
         img_list.append( [str(filepath), yaw ] )
 
     io.log_info ("Sorting...")
     img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
-    
+
     return img_list, trash_img_list
-    
+
 def sort_by_face_pitch(input_path):
     io.log_info ("Sorting by face pitch...")
     img_list = []
     trash_img_list = []
     for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"):
         filepath = Path(filepath)
-        
+
         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) )
             trash_img_list.append ( [str(filepath)] )
             continue
-        
+
         pitch, yaw = LandmarksProcessor.estimate_pitch_yaw ( dflimg.get_landmarks() )
-       
+
         img_list.append( [str(filepath), pitch ] )
 
     io.log_info ("Sorting...")
     img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
-    
+
     return img_list, trash_img_list
 
 class HistSsimSubprocessor(Subprocessor):
     class Cli(Subprocessor.Cli):
         #override
-        def on_initialize(self, client_dict):  
+        def on_initialize(self, client_dict):
             self.log_info ('Running on %s.' % (client_dict['device_name']) )
 
         #override
         def process_data(self, data):
             img_list = []
             for x in data:
-                img = cv2_imread(x)    
+                img = cv2_imread(x)
                 img_list.append ([x, cv2.calcHist([img], [0], None, [256], [0, 256]),
                                      cv2.calcHist([img], [1], None, [256], [0, 256]),
                                      cv2.calcHist([img], [2], None, [256], [0, 256])
                                  ])
-        
+
             img_list_len = len(img_list)
             for i in range(img_list_len-1):
                 min_score = float("inf")
@@ -268,23 +268,23 @@ class HistSsimSubprocessor(Subprocessor):
                         min_score = score
                         j_min_score = j
                 img_list[i+1], img_list[j_min_score] = img_list[j_min_score], img_list[i+1]
-                
+
                 self.progress_bar_inc(1)
-            
-            return img_list     
+
+            return img_list
 
         #override
         def get_data_name (self, data):
             return "Bunch of images"
-    
+
     #override
-    def __init__(self, img_list ): 
+    def __init__(self, img_list ):
         self.img_list = img_list
         self.img_list_len = len(img_list)
-        
+
         slice_count = 20000
         sliced_count = self.img_list_len // slice_count
-        
+
         if sliced_count > 12:
             sliced_count = 11.9
             slice_count = int(self.img_list_len / sliced_count)
@@ -294,10 +294,10 @@ class HistSsimSubprocessor(Subprocessor):
                                [ self.img_list[sliced_count*slice_count:] ]
 
         self.result = []
-        super().__init__('HistSsim', HistSsimSubprocessor.Cli, 0)           
+        super().__init__('HistSsim', HistSsimSubprocessor.Cli, 0)
 
     #override
-    def process_info_generator(self):    
+    def process_info_generator(self):
         for i in range( len(self.img_chunks_list) ):
             yield 'CPU%d' % (i), {'i':i}, {'device_idx': i,
                                            'device_name': 'CPU%d' % (i)
@@ -306,21 +306,21 @@ class HistSsimSubprocessor(Subprocessor):
     def on_clients_initialized(self):
         io.progress_bar ("Sorting", len(self.img_list))
         io.progress_bar_inc(len(self.img_chunks_list))
-        
+
     #override
     def on_clients_finalized(self):
         io.progress_bar_close()
-    
+
     #override
-    def get_data(self, host_dict):     
-        if len (self.img_chunks_list) > 0:        
-            return self.img_chunks_list.pop(0)        
+    def get_data(self, host_dict):
+        if len (self.img_chunks_list) > 0:
+            return self.img_chunks_list.pop(0)
         return None
-    
+
     #override
     def on_data_return (self, host_dict, data):
         raise Exception("Fail to process data. Decrease number of images and try again.")
-        
+
     #override
     def on_result (self, host_dict, data, result):
         self.result += result
@@ -329,10 +329,10 @@ class HistSsimSubprocessor(Subprocessor):
     #override
     def get_result(self):
         return self.result
-    
+
 def sort_by_hist(input_path):
     io.log_info ("Sorting by histogram similarity...")
-    img_list = HistSsimSubprocessor(Path_utils.get_image_paths(input_path)).run()    
+    img_list = HistSsimSubprocessor(Path_utils.get_image_paths(input_path)).run()
     return img_list
 
 class HistDissimSubprocessor(Subprocessor):
@@ -344,7 +344,7 @@ class HistDissimSubprocessor(Subprocessor):
             self.img_list_len = len(self.img_list)
 
         #override
-        def process_data(self, data):  
+        def process_data(self, data):
             i = data[0]
             score_total = 0
             for j in range( 0, self.img_list_len):
@@ -358,40 +358,40 @@ class HistDissimSubprocessor(Subprocessor):
         def get_data_name (self, data):
             #return string identificator of your data
             return self.img_list[data[0]][0]
-    
+
     #override
-    def __init__(self, img_list ): 
+    def __init__(self, img_list ):
         self.img_list = img_list
         self.img_list_range = [i for i in range(0, len(img_list) )]
         self.result = []
-        super().__init__('HistDissim', HistDissimSubprocessor.Cli, 60)           
+        super().__init__('HistDissim', HistDissimSubprocessor.Cli, 60)
 
     #override
     def on_clients_initialized(self):
         io.progress_bar ("Sorting", len (self.img_list) )
-        
+
     #override
     def on_clients_finalized(self):
         io.progress_bar_close()
-        
+
     #override
-    def process_info_generator(self):    
+    def process_info_generator(self):
         for i in range(0, min(multiprocessing.cpu_count(), 8) ):
             yield 'CPU%d' % (i), {}, {'device_idx': i,
-                                      'device_name': 'CPU%d' % (i), 
+                                      'device_name': 'CPU%d' % (i),
                                       'img_list' : self.img_list
                                       }
     #override
     def get_data(self, host_dict):
-        if len (self.img_list_range) > 0:        
+        if len (self.img_list_range) > 0:
             return [self.img_list_range.pop(0)]
-        
+
         return None
-    
+
     #override
     def on_data_return (self, host_dict, data):
-        self.img_list_range.insert(0, data[0])   
-    
+        self.img_list_range.insert(0, data[0])
+
     #override
     def on_result (self, host_dict, data, result):
         self.img_list[data[0]][2] = result
@@ -400,7 +400,7 @@ class HistDissimSubprocessor(Subprocessor):
     #override
     def get_result(self):
         return self.img_list
-        
+
 def sort_by_hist_dissim(input_path):
     io.log_info ("Sorting by histogram dissimilarity...")
 
@@ -408,19 +408,19 @@ def sort_by_hist_dissim(input_path):
     trash_img_list = []
     for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"):
         filepath = Path(filepath)
-        
+
         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) )
             trash_img_list.append ([str(filepath)])
             continue
-            
+
         image = cv2_imread(str(filepath))
         face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks())
         image = (image*face_mask).astype(np.uint8)
@@ -428,26 +428,26 @@ def sort_by_hist_dissim(input_path):
         img_list.append ([str(filepath), cv2.calcHist([cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)], [0], None, [256], [0, 256]), 0 ])
 
     img_list = HistDissimSubprocessor(img_list).run()
-                         
+
     io.log_info ("Sorting...")
     img_list = sorted(img_list, key=operator.itemgetter(2), reverse=True)
 
     return img_list, trash_img_list
-    
+
 def sort_by_brightness(input_path):
     io.log_info ("Sorting by brightness...")
     img_list = [ [x, np.mean ( cv2.cvtColor(cv2_imread(x), cv2.COLOR_BGR2HSV)[...,2].flatten()  )] for x in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading") ]
     io.log_info ("Sorting...")
-    img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)    
+    img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
     return img_list
-    
+
 def sort_by_hue(input_path):
     io.log_info ("Sorting by hue...")
     img_list = [ [x, np.mean ( cv2.cvtColor(cv2_imread(x), cv2.COLOR_BGR2HSV)[...,0].flatten()  )] for x in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading") ]
     io.log_info ("Sorting...")
-    img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)    
+    img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
     return img_list
-    
+
 def sort_by_black(input_path):
     io.log_info ("Sorting by amount of black pixels...")
 
@@ -460,22 +460,22 @@ def sort_by_black(input_path):
     img_list = sorted(img_list, key=operator.itemgetter(1), reverse=False)
 
     return img_list
-    
+
 def sort_by_origname(input_path):
     io.log_info ("Sort by original filename...")
-    
+
     img_list = []
     trash_img_list = []
     for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"):
         filepath = Path(filepath)
-        
+
         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) )
             trash_img_list.append( [str(filepath)] )
@@ -486,7 +486,7 @@ def sort_by_origname(input_path):
     io.log_info ("Sorting...")
     img_list = sorted(img_list, key=operator.itemgetter(1))
     return img_list, trash_img_list
-    
+
 def sort_by_oneface_in_image(input_path):
     io.log_info ("Sort by one face in images...")
     image_paths = Path_utils.get_image_paths(input_path)
@@ -503,17 +503,17 @@ def sort_by_oneface_in_image(input_path):
             trash_img_list = [ (image_paths[x],) for x in idxs ]
             return img_list, trash_img_list
     return [], []
-    
+
 class FinalLoaderSubprocessor(Subprocessor):
     class Cli(Subprocessor.Cli):
         #override
-        def on_initialize(self, client_dict):        
+        def on_initialize(self, client_dict):
             self.log_info ('Running on %s.' % (client_dict['device_name']) )
             self.include_by_blur = client_dict['include_by_blur']
-            
+
         #override
-        def process_data(self, data):        
-            filepath = Path(data[0])     
+        def process_data(self, data):
+            filepath = Path(data[0])
 
             try:
                 if filepath.suffix == '.png':
@@ -522,40 +522,40 @@ class FinalLoaderSubprocessor(Subprocessor):
                     dflimg = DFLJPG.load( str(filepath) )
                 else:
                     dflimg = None
-                    
+
                 if dflimg is None:
                     self.log_err("%s is not a dfl image file" % (filepath.name))
                     return [ 1, [str(filepath)] ]
-                
+
                 bgr = cv2_imread(str(filepath))
                 if bgr is None:
-                    raise Exception ("Unable to load %s" % (filepath.name) ) 
-                
+                    raise Exception ("Unable to load %s" % (filepath.name) )
+
                 gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
                 sharpness = estimate_sharpness(gray) if self.include_by_blur else 0
                 pitch, yaw = LandmarksProcessor.estimate_pitch_yaw ( dflimg.get_landmarks() )
-                
+
                 hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
             except Exception as e:
                 self.log_err (e)
                 return [ 1, [str(filepath)] ]
-                
+
             return [ 0, [str(filepath), sharpness, hist, yaw ] ]
 
         #override
         def get_data_name (self, data):
             #return string identificator of your data
             return data[0]
-    
+
     #override
-    def __init__(self, img_list, include_by_blur ): 
+    def __init__(self, img_list, include_by_blur ):
         self.img_list = img_list
 
         self.include_by_blur = include_by_blur
         self.result = []
         self.result_trash = []
 
-        super().__init__('FinalLoader', FinalLoaderSubprocessor.Cli, 60)           
+        super().__init__('FinalLoader', FinalLoaderSubprocessor.Cli, 60)
 
     #override
     def on_clients_initialized(self):
@@ -564,9 +564,9 @@ class FinalLoaderSubprocessor(Subprocessor):
     #override
     def on_clients_finalized(self):
         io.progress_bar_close()
-        
+
     #override
-    def process_info_generator(self):    
+    def process_info_generator(self):
         for i in range(0, min(multiprocessing.cpu_count(), 8) ):
             yield 'CPU%d' % (i), {}, {'device_idx': i,
                                       'device_name': 'CPU%d' % (i),
@@ -575,15 +575,15 @@ class FinalLoaderSubprocessor(Subprocessor):
 
     #override
     def get_data(self, host_dict):
-        if len (self.img_list) > 0:        
+        if len (self.img_list) > 0:
             return [self.img_list.pop(0)]
-        
+
         return None
-    
+
     #override
     def on_data_return (self, host_dict, data):
-        self.img_list.insert(0, data[0])   
-     
+        self.img_list.insert(0, data[0])
+
     #override
     def on_result (self, host_dict, data, result):
         if result[0] == 0:
@@ -599,7 +599,7 @@ class FinalLoaderSubprocessor(Subprocessor):
 class FinalHistDissimSubprocessor(Subprocessor):
     class Cli(Subprocessor.Cli):
         #override
-        def on_initialize(self, client_dict):  
+        def on_initialize(self, client_dict):
             self.log_info ('Running on %s.' % (client_dict['device_name']) )
 
         #override
@@ -611,25 +611,25 @@ class FinalHistDissimSubprocessor(Subprocessor):
                     if i == j:
                         continue
                     score_total += cv2.compareHist(img_list[i][2], img_list[j][2], cv2.HISTCMP_BHATTACHARYYA)
-                img_list[i][3] = score_total                
+                img_list[i][3] = score_total
             img_list = sorted(img_list, key=operator.itemgetter(3), reverse=True)
-            return idx, img_list     
+            return idx, img_list
 
         #override
         def get_data_name (self, data):
             return "Bunch of images"
-    
-    #override
-    def __init__(self, yaws_sample_list ): 
-        self.yaws_sample_list = yaws_sample_list
-        self.yaws_sample_list_len = len(yaws_sample_list)
-        
-        self.yaws_sample_list_idxs = [ i for i in range(self.yaws_sample_list_len) if self.yaws_sample_list[i] is not None ]
-        self.result = [ None for _ in range(self.yaws_sample_list_len) ]
-        super().__init__('FinalHistDissimSubprocessor', FinalHistDissimSubprocessor.Cli)           
 
     #override
-    def process_info_generator(self):    
+    def __init__(self, yaws_sample_list ):
+        self.yaws_sample_list = yaws_sample_list
+        self.yaws_sample_list_len = len(yaws_sample_list)
+
+        self.yaws_sample_list_idxs = [ i for i in range(self.yaws_sample_list_len) if self.yaws_sample_list[i] is not None ]
+        self.result = [ None for _ in range(self.yaws_sample_list_len) ]
+        super().__init__('FinalHistDissimSubprocessor', FinalHistDissimSubprocessor.Cli)
+
+    #override
+    def process_info_generator(self):
         for i in range(min(multiprocessing.cpu_count(), 8) ):
             yield 'CPU%d' % (i), {'i':i}, {'device_idx': i,
                                            'device_name': 'CPU%d' % (i)
@@ -637,38 +637,38 @@ class FinalHistDissimSubprocessor(Subprocessor):
     #override
     def on_clients_initialized(self):
         io.progress_bar ("Sort by hist-dissim", self.yaws_sample_list_len)
-        
+
     #override
     def on_clients_finalized(self):
         io.progress_bar_close()
-    
+
     #override
-    def get_data(self, host_dict):     
+    def get_data(self, host_dict):
         if len (self.yaws_sample_list_idxs) > 0:
             idx = self.yaws_sample_list_idxs.pop(0)
-            
+
             return idx, self.yaws_sample_list[idx]
         return None
-    
+
     #override
     def on_data_return (self, host_dict, data):
         self.yaws_sample_list_idxs.insert(0, data[0])
-        
+
     #override
     def on_result (self, host_dict, data, result):
-        idx, yaws_sample_list = data        
+        idx, yaws_sample_list = data
         self.result[idx] = yaws_sample_list
         io.progress_bar_inc(1)
 
     #override
     def get_result(self):
         return self.result
-        
+
 def sort_final(input_path, include_by_blur=True):
     io.log_info ("Performing final sort.")
-    
+
     target_count = io.input_int ("Target number of images? (default:2000) : ", 2000)
-    
+
     img_list, trash_img_list = FinalLoaderSubprocessor( Path_utils.get_image_paths(input_path), include_by_blur ).run()
     final_img_list = []
 
@@ -676,12 +676,12 @@ def sort_final(input_path, include_by_blur=True):
     imgs_per_grad = round (target_count / grads)
 
     grads_space = np.linspace (-1.0,1.0,grads)
-    
+
     yaws_sample_list = [None]*grads
-    for g in io.progress_bar_generator ( range(grads), "Sort by yaw"):    
+    for g in io.progress_bar_generator ( range(grads), "Sort by yaw"):
         yaw = grads_space[g]
         next_yaw = grads_space[g+1] if g < grads-1 else yaw
-        
+
         yaw_samples = []
         for img in img_list:
             s_yaw = -img[3]
@@ -691,17 +691,17 @@ def sort_final(input_path, include_by_blur=True):
                 yaw_samples += [ img ]
         if len(yaw_samples) > 0:
             yaws_sample_list[g] = yaw_samples
-    
+
     total_lack = 0
     for g in io.progress_bar_generator ( range(grads), ""):
         img_list = yaws_sample_list[g]
         img_list_len = len(img_list) if img_list is not None else 0
-        
-        lack = imgs_per_grad - img_list_len
-        total_lack += max(lack, 0)        
 
-    imgs_per_grad += total_lack // grads    
-        
+        lack = imgs_per_grad - img_list_len
+        total_lack += max(lack, 0)
+
+    imgs_per_grad += total_lack // grads
+
     if include_by_blur:
         sharpned_imgs_per_grad = imgs_per_grad*10
         for g in io.progress_bar_generator ( range (grads), "Sort by blur"):
@@ -709,47 +709,47 @@ def sort_final(input_path, include_by_blur=True):
             if img_list is None:
                 continue
 
-            img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)    
-            
+            img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
+
             if len(img_list) > sharpned_imgs_per_grad:
                 trash_img_list += img_list[sharpned_imgs_per_grad:]
                 img_list = img_list[0:sharpned_imgs_per_grad]
-                
+
             yaws_sample_list[g] = img_list
-            
+
     yaws_sample_list = FinalHistDissimSubprocessor(yaws_sample_list).run()
 
     for g in io.progress_bar_generator ( range (grads), "Fetching best"):
         img_list = yaws_sample_list[g]
         if img_list is None:
             continue
-            
+
         final_img_list += img_list[0:imgs_per_grad]
         trash_img_list += img_list[imgs_per_grad:]
 
     return final_img_list, trash_img_list
-  
+
 def final_process(input_path, img_list, trash_img_list):
     if len(trash_img_list) != 0:
         parent_input_path = input_path.parent
         trash_path = parent_input_path / (input_path.stem + '_trash')
         trash_path.mkdir (exist_ok=True)
-        
-        io.log_info ("Trashing %d items to %s" % ( len(trash_img_list), str(trash_path) ) )        
-        
+
+        io.log_info ("Trashing %d items to %s" % ( len(trash_img_list), str(trash_path) ) )
+
         for filename in Path_utils.get_image_paths(trash_path):
             Path(filename).unlink()
 
         for i in io.progress_bar_generator( range(len(trash_img_list)), "Moving trash", leave=False):
-            src = Path (trash_img_list[i][0])        
+            src = Path (trash_img_list[i][0])
             dst = trash_path / src.name
             try:
                 src.rename (dst)
             except:
                 io.log_info ('fail to trashing %s' % (src.name) )
-                
+
         io.log_info ("")
-    
+
     if len(img_list) != 0:
         for i in io.progress_bar_generator( [*range(len(img_list))], "Renaming", leave=False):
             src = Path (img_list[i][0])
@@ -758,24 +758,24 @@ def final_process(input_path, img_list, trash_img_list):
                 src.rename (dst)
             except:
                 io.log_info ('fail to rename %s' % (src.name) )
-                
+
         for i in io.progress_bar_generator( [*range(len(img_list))], "Renaming"):
-            src = Path (img_list[i][0])            
+            src = Path (img_list[i][0])
             src = input_path / ('%.5d_%s' % (i, src.name))
             dst = input_path / ('%.5d%s' % (i, src.suffix))
             try:
                 src.rename (dst)
             except:
-                io.log_info ('fail to rename %s' % (src.name) )   
+                io.log_info ('fail to rename %s' % (src.name) )
+
 
 
-    
 def main (input_path, sort_by_method):
     input_path = Path(input_path)
     sort_by_method = sort_by_method.lower()
 
     io.log_info ("Running sort tool.\r\n")
-    
+
     img_list = []
     trash_img_list = []
     if sort_by_method == 'blur':            img_list, trash_img_list = sort_by_blur (input_path)
@@ -787,10 +787,10 @@ def main (input_path, sort_by_method):
     elif sort_by_method == 'hist-dissim':   img_list, trash_img_list = sort_by_hist_dissim (input_path)
     elif sort_by_method == 'brightness':    img_list = sort_by_brightness (input_path)
     elif sort_by_method == 'hue':           img_list = sort_by_hue (input_path)
-    elif sort_by_method == 'black':         img_list = sort_by_black (input_path)    
+    elif sort_by_method == 'black':         img_list = sort_by_black (input_path)
     elif sort_by_method == 'origname':      img_list, trash_img_list = sort_by_origname (input_path)
-    elif sort_by_method == 'oneface':       img_list, trash_img_list = sort_by_oneface_in_image (input_path)      
-    elif sort_by_method == 'final':         img_list, trash_img_list = sort_final (input_path)  
+    elif sort_by_method == 'oneface':       img_list, trash_img_list = sort_by_oneface_in_image (input_path)
+    elif sort_by_method == 'final':         img_list, trash_img_list = sort_final (input_path)
     elif sort_by_method == 'final-no-blur': img_list, trash_img_list = sort_final (input_path, include_by_blur=False)
-    
+
     final_process (input_path, img_list, trash_img_list)

mainscripts/Trainer.py

@@ -7,39 +7,39 @@ import numpy as np
 import itertools
 from pathlib import Path
 from utils import Path_utils
-from utils import image_utils    
+from utils import image_utils
 import cv2
 import models
 from interact import interact as io
 
 def trainerThread (s2c, c2s, args, device_args):
     while True:
-        try: 
+        try:
             training_data_src_path = Path( args.get('training_data_src_dir', '') )
             training_data_dst_path = Path( args.get('training_data_dst_dir', '') )
             model_path = Path( args.get('model_path', '') )
             model_name = args.get('model_name', '')
-            save_interval_min = 15            
+            save_interval_min = 15
             debug = args.get('debug', '')
-            
+
             if not training_data_src_path.exists():
                 io.log_err('Training data src directory does not exist.')
                 break
-                
+
             if not training_data_dst_path.exists():
                 io.log_err('Training data dst directory does not exist.')
                 break
-                
+
             if not model_path.exists():
                 model_path.mkdir(exist_ok=True)
-   
+
             model = models.import_model(model_name)(
-                        model_path, 
-                        training_data_src_path=training_data_src_path, 
-                        training_data_dst_path=training_data_dst_path, 
+                        model_path,
+                        training_data_src_path=training_data_src_path,
+                        training_data_dst_path=training_data_dst_path,
                         debug=debug,
                         device_args=device_args)
-            
+
             is_reached_goal = model.is_reached_iter_goal()
             is_upd_save_time_after_train = False
             loss_string = ""
@@ -49,37 +49,37 @@ def trainerThread (s2c, c2s, args, device_args):
                     model.save()
                     io.log_info(loss_string)
                     is_upd_save_time_after_train = True
-            
+
             def send_preview():
-                if not debug:                        
-                    previews = model.get_previews()                
+                if not debug:
+                    previews = model.get_previews()
                     c2s.put ( {'op':'show', 'previews': previews, 'iter':model.get_iter(), 'loss_history': model.get_loss_history().copy() } )
                 else:
                     previews = [( 'debug, press update for new', model.debug_one_iter())]
                     c2s.put ( {'op':'show', 'previews': previews} )
-            
-            
+
+
             if model.is_first_run():
                 model_save()
-                
+
             if model.get_target_iter() != 0:
                 if is_reached_goal:
                     io.log_info('Model already trained to target iteration. You can use preview.')
                 else:
                     io.log_info('Starting. Target iteration: %d. Press "Enter" to stop training and save model.' % ( model.get_target_iter()  ) )
-            else: 
+            else:
                 io.log_info('Starting. Press "Enter" to stop training and save model.')
- 
+
             last_save_time = time.time()
-            
+
             for i in itertools.count(0,1):
                 if not debug:
                     if not is_reached_goal:
-                        loss_string = model.train_one_iter()     
+                        loss_string = model.train_one_iter()
                         if is_upd_save_time_after_train:
                             #save resets plaidML programs, so upd last_save_time only after plaidML rebuild them
                             last_save_time = time.time()
-                        
+
                         io.log_info (loss_string, end='\r')
                         if model.get_target_iter() != 0 and model.is_reached_iter_goal():
                             io.log_info ('Reached target iteration.')
@@ -91,77 +91,77 @@ def trainerThread (s2c, c2s, args, device_args):
                     last_save_time = time.time()
                     model_save()
                     send_preview()
-                    
+
                 if i==0:
                     if is_reached_goal:
-                        model.pass_one_iter()    
+                        model.pass_one_iter()
                     send_preview()
-                    
+
                 if debug:
                     time.sleep(0.005)
-                    
+
                 while not s2c.empty():
                     input = s2c.get()
                     op = input['op']
                     if op == 'save':
                         model_save()
-                    elif op == 'preview':                    
+                    elif op == 'preview':
                         if is_reached_goal:
-                            model.pass_one_iter()                    
+                            model.pass_one_iter()
                         send_preview()
                     elif op == 'close':
                         model_save()
                         i = -1
                         break
-                        
+
                 if i == -1:
                     break
-                    
-                
+
+
 
             model.finalize()
-                
+
         except Exception as e:
             print ('Error: %s' % (str(e)))
             traceback.print_exc()
         break
     c2s.put ( {'op':'close'} )
 
-   
-       
+
+
 def main(args, device_args):
     io.log_info ("Running trainer.\r\n")
-    
+
     no_preview = args.get('no_preview', False)
-    
+
     s2c = queue.Queue()
     c2s = queue.Queue()
-    
+
     thread = threading.Thread(target=trainerThread, args=(s2c, c2s, args, device_args) )
     thread.start()
 
     if no_preview:
-        while True:      
+        while True:
             if not c2s.empty():
                 input = c2s.get()
                 op = input.get('op','')
                 if op == 'close':
                     break
             io.process_messages(0.1)
-    else:    
+    else:
         wnd_name = "Training preview"
         io.named_window(wnd_name)
         io.capture_keys(wnd_name)
-        
+
         previews = None
         loss_history = None
         selected_preview = 0
         update_preview = False
         is_showing = False
         is_waiting_preview = False
-        show_last_history_iters_count = 0    
+        show_last_history_iters_count = 0
         iter = 0
-        while True:      
+        while True:
             if not c2s.empty():
                 input = c2s.get()
                 op = input['op']
@@ -177,7 +177,7 @@ def main(args, device_args):
                             (h, w, c) = preview_rgb.shape
                             max_h = max (max_h, h)
                             max_w = max (max_w, w)
-                        
+
                         max_size = 800
                         if max_h > max_size:
                             max_w = int( max_w / (max_h / max_size) )
@@ -194,49 +194,49 @@ def main(args, device_args):
                         update_preview = True
                 elif op == 'close':
                     break
-                    
+
             if update_preview:
                 update_preview = False
 
                 selected_preview_name = previews[selected_preview][0]
                 selected_preview_rgb = previews[selected_preview][1]
                 (h,w,c) = selected_preview_rgb.shape
-                
+
                 # HEAD
                 head_lines = [
                     '[s]:save [enter]:exit',
                     '[p]:update [space]:next preview [l]:change history range',
                     'Preview: "%s" [%d/%d]' % (selected_preview_name,selected_preview+1, len(previews) )
-                    ] 
+                    ]
                 head_line_height = 15
                 head_height = len(head_lines) * head_line_height
                 head = np.ones ( (head_height,w,c) ) * 0.1
-                  
+
                 for i in range(0, len(head_lines)):
                     t = i*head_line_height
                     b = (i+1)*head_line_height
                     head[t:b, 0:w] += image_utils.get_text_image (  (w,head_line_height,c) , head_lines[i], color=[0.8]*c )
-                    
+
                 final = head
-       
-                if loss_history is not None:                
+
+                if loss_history is not None:
                     if show_last_history_iters_count == 0:
                         loss_history_to_show = loss_history
                     else:
                         loss_history_to_show = loss_history[-show_last_history_iters_count:]
-                        
+
                     lh_img = models.ModelBase.get_loss_history_preview(loss_history_to_show, iter, w, c)
                     final = np.concatenate ( [final, lh_img], axis=0 )
 
                 final = np.concatenate ( [final, selected_preview_rgb], axis=0 )
                 final = np.clip(final, 0, 1)
-                
+
                 io.show_image( wnd_name, (final*255).astype(np.uint8) )
                 is_showing = True
-            
+
             key_events = io.get_key_events(wnd_name)
             key, = key_events[-1] if len(key_events) > 0 else (0,)
-                    
+
             if key == ord('\n') or key == ord('\r'):
                 s2c.put ( {'op': 'close'} )
             elif key == ord('s'):
@@ -253,14 +253,14 @@ def main(args, device_args):
                 elif show_last_history_iters_count == 10000:
                     show_last_history_iters_count = 50000
                 elif show_last_history_iters_count == 50000:
-                    show_last_history_iters_count = 100000    
+                    show_last_history_iters_count = 100000
                 elif show_last_history_iters_count == 100000:
-                    show_last_history_iters_count = 0                
+                    show_last_history_iters_count = 0
                 update_preview = True
             elif key == ord(' '):
                 selected_preview = (selected_preview + 1) % len(previews)
                 update_preview = True
-            
+
             io.process_messages(0.1)
-            
-        io.destroy_all_windows()
+
+        io.destroy_all_windows()

mainscripts/Util.py

@@ -9,30 +9,30 @@ from interact import interact as io
 
 def convert_png_to_jpg_file (filepath):
     filepath = Path(filepath)
-    
-    if filepath.suffix != '.png': 
+
+    if filepath.suffix != '.png':
         return
-        
+
     dflpng = DFLPNG.load (str(filepath) )
     if dflpng is None:
-        io.log_err ("%s is not a dfl image file" % (filepath.name) ) 
+        io.log_err ("%s is not a dfl image file" % (filepath.name) )
         return
-    
+
     dfl_dict = dflpng.getDFLDictData()
-    
+
     img = cv2_imread (str(filepath))
     new_filepath = str(filepath.parent / (filepath.stem + '.jpg'))
     cv2_imwrite ( new_filepath, img, [int(cv2.IMWRITE_JPEG_QUALITY), 85])
 
-    DFLJPG.embed_data( new_filepath, 
+    DFLJPG.embed_data( new_filepath,
                        face_type=dfl_dict.get('face_type', None),
                        landmarks=dfl_dict.get('landmarks', 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) )
-                       
+
     filepath.unlink()
-        
+
 def convert_png_to_jpg_folder (input_path):
     input_path = Path(input_path)
 
@@ -41,73 +41,73 @@ def convert_png_to_jpg_folder (input_path):
     for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Converting"):
         filepath = Path(filepath)
         convert_png_to_jpg_file(filepath)
-        
+
 def add_landmarks_debug_images(input_path):
     io.log_info ("Adding landmarks debug images...")
 
     for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Processing"):
         filepath = Path(filepath)
- 
+
         img = cv2_imread(str(filepath))
- 
+
         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) ) 
+            io.log_err ("%s is not a dfl image file" % (filepath.name) )
             continue
 
         if img is not None:
             face_landmarks = dflimg.get_landmarks()
             LandmarksProcessor.draw_landmarks(img, face_landmarks, transparent_mask=True)
-            
+
             output_file = '{}{}'.format( str(Path(str(input_path)) / filepath.stem),  '_debug.jpg')
             cv2_imwrite(output_file, img, [int(cv2.IMWRITE_JPEG_QUALITY), 50] )
-            
+
 def recover_original_aligned_filename(input_path):
     io.log_info ("Recovering original aligned filename...")
-    
+
     files = []
     for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Processing"):
         filepath = Path(filepath)
- 
+
         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) ) 
+            io.log_err ("%s is not a dfl image file" % (filepath.name) )
             continue
-            
+
         files += [ [filepath, None, dflimg.get_source_filename(), False] ]
-      
+
     files_len = len(files)
     for i in io.progress_bar_generator( range(files_len), "Sorting" ):
         fp, _, sf, converted = files[i]
-        
+
         if converted:
             continue
-            
+
         sf_stem = Path(sf).stem
-        
+
         files[i][1] = fp.parent / ( sf_stem + '_0' + fp.suffix )
         files[i][3] = True
         c = 1
-        
+
         for j in range(i+1, files_len):
             fp_j, _, sf_j, converted_j = files[j]
             if converted_j:
                 continue
-            
+
             if sf_j == sf:
-                files[j][1] = fp_j.parent / ( sf_stem + ('_%d' % (c)) + fp_j.suffix )                 
+                files[j][1] = fp_j.parent / ( sf_stem + ('_%d' % (c)) + fp_j.suffix )
                 files[j][3] = True
                 c += 1
 
@@ -118,11 +118,11 @@ def recover_original_aligned_filename(input_path):
             fs.rename (dst)
         except:
             io.log_err ('fail to rename %s' % (fs.name) )
-            
+
     for file in io.progress_bar_generator( files, "Renaming" ):
         fs, fd, _, _ = file
         fs = fs.parent / ( fs.stem + '_tmp' + fs.suffix )
         try:
             fs.rename (fd)
         except:
-            io.log_err ('fail to rename %s' % (fs.name) )
+            io.log_err ('fail to rename %s' % (fs.name) )

mainscripts/VideoEd.py

@@ -8,38 +8,38 @@ from interact import interact as io
 def extract_video(input_file, output_dir, output_ext=None, fps=None):
     input_file_path = Path(input_file)
     output_path = Path(output_dir)
-    
+
     if not output_path.exists():
         output_path.mkdir(exist_ok=True)
-    
-    
+
+
     if input_file_path.suffix == '.*':
         input_file_path = Path_utils.get_first_file_by_stem (input_file_path.parent, input_file_path.stem)
     else:
         if not input_file_path.exists():
             input_file_path = None
-    
+
     if input_file_path is None:
         io.log_err("input_file not found.")
         return
-    
+
     if output_ext is None:
         output_ext = io.input_str ("Output image format (extension)? ( default:png ) : ", "png")
-    
+
     if fps is None:
         fps = io.input_int ("Enter FPS ( ?:help skip:fullfps ) : ", 0, help_message="How many frames of every second of the video will be extracted.")
-    
+
     for filename in Path_utils.get_image_paths (output_path, ['.'+output_ext]):
         Path(filename).unlink()
 
     job = ffmpeg.input(str(input_file_path))
-    
-    kwargs = {}    
+
+    kwargs = {}
     if fps != 0:
         kwargs.update ({'r':str(fps)})
 
     job = job.output( str (output_path / ('%5d.'+output_ext)), **kwargs )
-    
+
     try:
         job = job.run()
     except:
@@ -50,18 +50,18 @@ def cut_video ( input_file, from_time=None, to_time=None, audio_track_id=None, b
     if input_file_path is None:
         io.log_err("input_file not found.")
         return
-    
+
     output_file_path = input_file_path.parent / (input_file_path.stem + "_cut" + input_file_path.suffix)
-    
+
     if from_time is None:
         from_time = io.input_str ("From time (skip: 00:00:00.000) : ", "00:00:00.000")
-    
+
     if to_time is None:
         to_time = io.input_str ("To time (skip: 00:00:00.000) : ", "00:00:00.000")
-        
+
     if audio_track_id is None:
         audio_track_id = io.input_int ("Specify audio track id. ( skip:0 ) : ", 0)
-        
+
     if bitrate is None:
         bitrate = max (1, io.input_int ("Bitrate of output file in MB/s ? (default:25) : ", 25) )
 
@@ -69,64 +69,64 @@ def cut_video ( input_file, from_time=None, to_time=None, audio_track_id=None, b
               "b:v": "%dM" %(bitrate),
               "pix_fmt": "yuv420p",
              }
-    
+
     job = ffmpeg.input(str(input_file_path), ss=from_time, to=to_time)
-    
+
     job_v = job['v:0']
     job_a = job['a:' + str(audio_track_id) + '?' ]
-    
+
     job = ffmpeg.output(job_v, job_a, str(output_file_path), **kwargs).overwrite_output()
-    
+
     try:
         job = job.run()
     except:
         io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) )
-    
+
 def denoise_image_sequence( input_dir, ext=None, factor=None ):
     input_path = Path(input_dir)
-    
+
     if not input_path.exists():
         io.log_err("input_dir not found.")
         return
-        
+
     if ext is None:
         ext = io.input_str ("Input image format (extension)? ( default:png ) : ", "png")
-    
+
     if factor is None:
         factor = np.clip ( io.input_int ("Denoise factor? (1-20 default:5) : ", 5), 1, 20 )
-    
+
     job = ( ffmpeg
             .input(str ( input_path / ('%5d.'+ext) ) )
             .filter("hqdn3d", factor, factor, 5,5)
             .output(str ( input_path / ('%5d.'+ext) ) )
-           )    
-           
+           )
+
     try:
         job = job.run()
     except:
         io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) )
-        
+
 def video_from_sequence( input_dir, output_file, reference_file=None, ext=None, fps=None, bitrate=None, lossless=None ):
-    input_path = Path(input_dir)    
+    input_path = Path(input_dir)
     output_file_path = Path(output_file)
     reference_file_path = Path(reference_file) if reference_file is not None else None
-    
+
     if not input_path.exists():
         io.log_err("input_dir not found.")
         return
-        
+
     if not output_file_path.parent.exists():
         output_file_path.parent.mkdir(parents=True, exist_ok=True)
         return
-    
+
     out_ext = output_file_path.suffix
-    
+
     if ext is None:
         ext = io.input_str ("Input image format (extension)? ( default:png ) : ", "png")
-    
+
     if lossless is None:
         lossless = io.input_bool ("Use lossless codec ? ( default:no ) : ", False)
-    
+
     video_id = None
     audio_id = None
     ref_in_a = None
@@ -136,7 +136,7 @@ def video_from_sequence( input_dir, output_file, reference_file=None, ext=None,
         else:
             if not reference_file_path.exists():
                 reference_file_path = None
-    
+
         if reference_file_path is None:
             io.log_err("reference_file not found.")
             return
@@ -149,32 +149,32 @@ def video_from_sequence( input_dir, output_file, reference_file=None, ext=None,
             if video_id is None and stream['codec_type'] == 'video':
                 video_id = stream['index']
                 fps = stream['r_frame_rate']
-                
+
             if audio_id is None and stream['codec_type'] == 'audio':
                 audio_id = stream['index']
 
         if audio_id is not None:
             #has audio track
             ref_in_a = ffmpeg.input (str(reference_file_path))[str(audio_id)]
-        
+
     if fps is None:
         #if fps not specified and not overwritten by reference-file
         fps = max (1, io.input_int ("FPS ? (default:25) : ", 25) )
-        
+
     if not lossless and bitrate is None:
         bitrate = max (1, io.input_int ("Bitrate of output file in MB/s ? (default:16) : ", 16) )
-    
+
     i_in = ffmpeg.input(str (input_path / ('%5d.'+ext)), r=fps)
-    
+
     output_args = [i_in]
-    
+
     if ref_in_a is not None:
         output_args += [ref_in_a]
-    
+
     output_args += [str (output_file_path)]
-    
+
     output_kwargs = {}
-    
+
     if lossless:
         output_kwargs.update ({"c:v": "png"
                               })
@@ -183,15 +183,14 @@ def video_from_sequence( input_dir, output_file, reference_file=None, ext=None,
                                "b:v": "%dM" %(bitrate),
                                "pix_fmt": "yuv420p",
                               })
-                          
+
     output_kwargs.update ({"c:a": "aac",
                            "b:a": "192k",
                            "ar" : "48000"
                           })
-    
+
     job = ( ffmpeg.output(*output_args, **output_kwargs).overwrite_output() )
     try:
         job = job.run()
     except:
         io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) )
-           

mathlib/__init__.py

@@ -7,10 +7,10 @@ def get_power_of_two(x):
     while (1 << i) < x:
         i += 1
     return i
-    
+
 def rotationMatrixToEulerAngles(R) :
-    sy = math.sqrt(R[0,0] * R[0,0] +  R[1,0] * R[1,0])     
-    singular = sy < 1e-6 
+    sy = math.sqrt(R[0,0] * R[0,0] +  R[1,0] * R[1,0])
+    singular = sy < 1e-6
     if  not singular :
         x = math.atan2(R[2,1] , R[2,2])
         y = math.atan2(-R[2,0], sy)
@@ -18,8 +18,8 @@ def rotationMatrixToEulerAngles(R) :
     else :
         x = math.atan2(-R[1,2], R[1,1])
         y = math.atan2(-R[2,0], sy)
-        z = 0 
+        z = 0
     return np.array([x, y, z])
-    
+
 def polygon_area(x,y):
-    return 0.5*np.abs(np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1)))
+    return 0.5*np.abs(np.dot(x,np.roll(y,1))-np.dot(y,np.roll(x,1)))

mathlib/umeyama.py

@@ -68,4 +68,4 @@ def umeyama(src, dst, estimate_scale):
     T[:dim, dim] = dst_mean - scale * np.dot(T[:dim, :dim], src_mean.T)
     T[:dim, :dim] *= scale
 
-    return T
+    return T

models/ModelBase.py

@@ -23,11 +23,11 @@ class ModelBase(object):
     def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None, debug = False, device_args = None,
                  ask_write_preview_history=True, ask_target_iter=True, ask_batch_size=True, ask_sort_by_yaw=True,
                  ask_random_flip=True, ask_src_scale_mod=True):
-        
+
         device_args['force_gpu_idx'] = device_args.get('force_gpu_idx',-1)
         device_args['cpu_only'] = device_args.get('cpu_only',False)
-        
-        if device_args['force_gpu_idx'] == -1 and not device_args['cpu_only']: 
+
+        if device_args['force_gpu_idx'] == -1 and not device_args['cpu_only']:
             idxs_names_list = nnlib.device.getValidDevicesIdxsWithNamesList()
             if len(idxs_names_list) > 1:
                 io.log_info ("You have multi GPUs in a system: ")
@@ -36,17 +36,17 @@ class ModelBase(object):
 
                 device_args['force_gpu_idx'] = io.input_int("Which GPU idx to choose? ( skip: best GPU ) : ", -1, [ x[0] for x in idxs_names_list] )
         self.device_args = device_args
-    
+
         self.device_config = nnlib.DeviceConfig(allow_growth=False, **self.device_args)
-        
+
         io.log_info ("Loading model...")
-            
+
         self.model_path = model_path
         self.model_data_path = Path( self.get_strpath_storage_for_file('data.dat') )
-        
+
         self.training_data_src_path = training_data_src_path
         self.training_data_dst_path = training_data_dst_path
-        
+
         self.src_images_paths = None
         self.dst_images_paths = None
         self.src_yaw_images_paths = None
@@ -60,10 +60,10 @@ class ModelBase(object):
         self.options = {}
         self.loss_history = []
         self.sample_for_preview = None
-        
+
         model_data = {}
-        if self.model_data_path.exists():            
-            model_data = pickle.loads ( self.model_data_path.read_bytes() )            
+        if self.model_data_path.exists():
+            model_data = pickle.loads ( self.model_data_path.read_bytes() )
             self.iter = max( model_data.get('iter',0), model_data.get('epoch',0) )
             if 'epoch' in self.options:
                 self.options.pop('epoch')
@@ -73,101 +73,101 @@ class ModelBase(object):
                 self.sample_for_preview = model_data['sample_for_preview']  if 'sample_for_preview' in model_data.keys() else None
 
         ask_override = self.is_training_mode and self.iter != 0 and io.input_in_time ("Press enter in 2 seconds to override model settings.", 2)
-        
+
         yn_str = {True:'y',False:'n'}
-        
-        if self.iter == 0: 
+
+        if self.iter == 0:
             io.log_info ("\nModel first run. Enter model options as default for each run.")
-        
+
         if ask_write_preview_history and (self.iter == 0 or ask_override):
             default_write_preview_history = False if self.iter == 0 else self.options.get('write_preview_history',False)
             self.options['write_preview_history'] = io.input_bool("Write preview history? (y/n ?:help skip:%s) : " % (yn_str[default_write_preview_history]) , default_write_preview_history, help_message="Preview history will be writed to <ModelName>_history folder.")
         else:
             self.options['write_preview_history'] = self.options.get('write_preview_history', False)
-            
+
         if ask_target_iter and (self.iter == 0 or ask_override):
             self.options['target_iter'] = max(0, io.input_int("Target iteration (skip:unlimited/default) : ", 0))
         else:
             self.options['target_iter'] = max(model_data.get('target_iter',0), self.options.get('target_epoch',0))
             if 'target_epoch' in self.options:
                 self.options.pop('target_epoch')
-        
+
         if ask_batch_size and (self.iter == 0 or ask_override):
             default_batch_size = 0 if self.iter == 0 else self.options.get('batch_size',0)
             self.options['batch_size'] = max(0, io.input_int("Batch_size (?:help skip:%d) : " % (default_batch_size), default_batch_size, help_message="Larger batch size is always better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually."))
         else:
             self.options['batch_size'] = self.options.get('batch_size', 0)
-        
+
         if ask_sort_by_yaw and (self.iter == 0):
             self.options['sort_by_yaw'] = io.input_bool("Feed faces to network sorted by yaw? (y/n ?:help skip:n) : ", False, help_message="NN will not learn src face directions that don't match dst face directions." )
         else:
             self.options['sort_by_yaw'] = self.options.get('sort_by_yaw', False)
-            
+
         if ask_random_flip and (self.iter == 0):
             self.options['random_flip'] = io.input_bool("Flip faces randomly? (y/n ?:help skip:y) : ", True, help_message="Predicted face will look more naturally without this option, but src faceset should cover all face directions as dst faceset.")
         else:
             self.options['random_flip'] = self.options.get('random_flip', True)
-        
+
         if ask_src_scale_mod and (self.iter == 0):
             self.options['src_scale_mod'] = np.clip( io.input_int("Src face scale modifier % ( -30...30, ?:help skip:0) : ", 0, help_message="If src face shape is wider than dst, try to decrease this value to get a better result."), -30, 30)
-        else:            
+        else:
             self.options['src_scale_mod'] = self.options.get('src_scale_mod', 0)
-             
+
         self.write_preview_history = self.options['write_preview_history']
         if not self.options['write_preview_history']:
-            self.options.pop('write_preview_history') 
-        
+            self.options.pop('write_preview_history')
+
         self.target_iter = self.options['target_iter']
         if self.options['target_iter'] == 0:
-            self.options.pop('target_iter') 
-           
+            self.options.pop('target_iter')
+
         self.batch_size = self.options['batch_size']
-        self.sort_by_yaw = self.options['sort_by_yaw']        
+        self.sort_by_yaw = self.options['sort_by_yaw']
         self.random_flip = self.options['random_flip']
-        
+
         self.src_scale_mod = self.options['src_scale_mod']
         if self.src_scale_mod == 0:
-            self.options.pop('src_scale_mod') 
+            self.options.pop('src_scale_mod')
 
         self.onInitializeOptions(self.iter == 0, ask_override)
 
         nnlib.import_all(self.device_config)
         self.keras = nnlib.keras
         self.K = nnlib.keras.backend
-        
+
         self.onInitialize()
-        
+
         self.options['batch_size'] = self.batch_size
-        
+
         if self.debug or self.batch_size == 0:
-            self.batch_size = 1 
-        
+            self.batch_size = 1
+
         if self.is_training_mode:
             if self.write_preview_history:
                 if self.device_args['force_gpu_idx'] == -1:
                     self.preview_history_path = self.model_path / ( '%s_history' % (self.get_model_name()) )
                 else:
                     self.preview_history_path = self.model_path / ( '%d_%s_history' % (self.device_args['force_gpu_idx'], self.get_model_name()) )
-                
+
                 if not self.preview_history_path.exists():
                     self.preview_history_path.mkdir(exist_ok=True)
                 else:
                     if self.iter == 0:
                         for filename in Path_utils.get_image_paths(self.preview_history_path):
                             Path(filename).unlink()
-        
+
             if self.generator_list is None:
                 raise ValueError( 'You didnt set_training_data_generators()')
             else:
                 for i, generator in enumerate(self.generator_list):
                     if not isinstance(generator, SampleGeneratorBase):
                         raise ValueError('training data generator is not subclass of SampleGeneratorBase')
-                        
+
             if (self.sample_for_preview is None) or (self.iter == 0):
                 self.sample_for_preview = self.generate_next_sample()
 
         model_summary_text = []
-                
+
         model_summary_text += ["===== Model summary ====="]
         model_summary_text += ["== Model name: " + self.get_model_name()]
         model_summary_text += ["=="]
@@ -179,41 +179,41 @@ class ModelBase(object):
 
         if self.device_config.multi_gpu:
             model_summary_text += ["== |== multi_gpu : True "]
-        
+
         model_summary_text += ["== Running on:"]
         if self.device_config.cpu_only:
             model_summary_text += ["== |== [CPU]"]
         else:
             for idx in self.device_config.gpu_idxs:
                 model_summary_text += ["== |== [%d : %s]" % (idx, nnlib.device.getDeviceName(idx))]
- 
+
         if not self.device_config.cpu_only and self.device_config.gpu_vram_gb[0] == 2:
             model_summary_text += ["=="]
             model_summary_text += ["== WARNING: You are using 2GB GPU. Result quality may be significantly decreased."]
             model_summary_text += ["== If training does not start, close all programs and try again."]
             model_summary_text += ["== Also you can disable Windows Aero Desktop to get extra free VRAM."]
             model_summary_text += ["=="]
-            
-        model_summary_text += ["========================="]  
-        model_summary_text = "\r\n".join (model_summary_text) 
-        self.model_summary_text = model_summary_text       
+
+        model_summary_text += ["========================="]
+        model_summary_text = "\r\n".join (model_summary_text)
+        self.model_summary_text = model_summary_text
         io.log_info(model_summary_text)
-        
+
     #overridable
     def onInitializeOptions(self, is_first_run, ask_override):
         pass
-       
+
     #overridable
     def onInitialize(self):
         '''
         initialize your keras models
-        
+
         store and retrieve your model options in self.options['']
-        
+
         check example
         '''
         pass
-        
+
     #overridable
     def onSave(self):
         #save your keras models here
@@ -229,59 +229,59 @@ class ModelBase(object):
     #overridable
     def onGetPreview(self, sample):
         #you can return multiple previews
-        #return [ ('preview_name',preview_rgb), ... ]        
+        #return [ ('preview_name',preview_rgb), ... ]
         return []
 
     #overridable if you want model name differs from folder name
     def get_model_name(self):
         return Path(inspect.getmodule(self).__file__).parent.name.rsplit("_", 1)[1]
-        
+
     #overridable
     def get_converter(self):
         raise NotImplementeError
         #return existing or your own converter which derived from base
-     
+
     def get_target_iter(self):
         return self.target_iter
-        
+
     def is_reached_iter_goal(self):
-        return self.target_iter != 0 and self.iter >= self.target_iter    
-     
+        return self.target_iter != 0 and self.iter >= self.target_iter
+
     #multi gpu in keras actually is fake and doesn't work for training https://github.com/keras-team/keras/issues/11976
     #def to_multi_gpu_model_if_possible (self, models_list):
     #    if len(self.device_config.gpu_idxs) > 1:
     #        #make batch_size to divide on GPU count without remainder
     #        self.batch_size = int( self.batch_size / len(self.device_config.gpu_idxs) )
     #        if self.batch_size == 0:
-    #            self.batch_size = 1                
+    #            self.batch_size = 1
     #        self.batch_size *= len(self.device_config.gpu_idxs)
-    #        
+    #
     #        result = []
     #        for model in models_list:
     #            for i in range( len(model.output_names) ):
-    #                model.output_names = 'output_%d' % (i)                 
-    #            result += [ nnlib.keras.utils.multi_gpu_model( model, self.device_config.gpu_idxs ) ]    
-    #            
-    #        return result                
+    #                model.output_names = 'output_%d' % (i)
+    #            result += [ nnlib.keras.utils.multi_gpu_model( model, self.device_config.gpu_idxs ) ]
+    #
+    #        return result
     #    else:
     #        return models_list
-     
-    def get_previews(self):       
+
+    def get_previews(self):
         return self.onGetPreview ( self.last_sample )
-        
-    def get_static_preview(self):        
+
+    def get_static_preview(self):
         return self.onGetPreview (self.sample_for_preview)[0][1] #first preview, and bgr
-       
-    def save(self):    
-        Path( self.get_strpath_storage_for_file('summary.txt') ).write_text(self.model_summary_text)  
+
+    def save(self):
+        Path( self.get_strpath_storage_for_file('summary.txt') ).write_text(self.model_summary_text)
         self.onSave()
-            
+
         model_data = {
             'iter': self.iter,
             'options': self.options,
             'loss_history': self.loss_history,
             'sample_for_preview' : self.sample_for_preview
-        }            
+        }
         self.model_data_path.write_bytes( pickle.dumps(model_data) )
 
     def load_weights_safe(self, model_filename_list, optimizer_filename_list=[]):
@@ -289,17 +289,17 @@ class ModelBase(object):
             filename = self.get_strpath_storage_for_file(filename)
             if Path(filename).exists():
                 model.load_weights(filename)
-            
+
         if len(optimizer_filename_list) != 0:
             opt_filename = self.get_strpath_storage_for_file('opt.h5')
             if Path(opt_filename).exists():
                 try:
                     with open(opt_filename, "rb") as f:
                         d = pickle.loads(f.read())
-                        
+
                     for x in optimizer_filename_list:
                         opt, filename = x
-                        if filename in d:                   
+                        if filename in d:
                             weights = d[filename].get('weights', None)
                             if weights:
                                 opt.set_weights(weights)
@@ -307,16 +307,16 @@ class ModelBase(object):
                 except Exception as e:
                     print ("Unable to load ", opt_filename)
 
-        
+
     def save_weights_safe(self, model_filename_list, optimizer_filename_list=[]):
         for model, filename in model_filename_list:
             filename = self.get_strpath_storage_for_file(filename)
             model.save_weights( filename + '.tmp' )
 
         rename_list = model_filename_list
-        if len(optimizer_filename_list) != 0:            
+        if len(optimizer_filename_list) != 0:
             opt_filename = self.get_strpath_storage_for_file('opt.h5')
-            
+
             try:
                 d = {}
                 for opt, filename in optimizer_filename_list:
@@ -324,54 +324,54 @@ class ModelBase(object):
                     symbolic_weights = getattr(opt, 'weights')
                     if symbolic_weights:
                         fd['weights'] = self.K.batch_get_value(symbolic_weights)
-                    
+
                     d[filename] = fd
-                    
+
                 with open(opt_filename+'.tmp', 'wb') as f:
                     f.write( pickle.dumps(d) )
-                    
+
                 rename_list += [('', 'opt.h5')]
             except Exception as e:
                 print ("Unable to save ", opt_filename)
 
         for _, filename in rename_list:
-            filename = self.get_strpath_storage_for_file(filename)        
+            filename = self.get_strpath_storage_for_file(filename)
             source_filename = Path(filename+'.tmp')
             if source_filename.exists():
                 target_filename = Path(filename)
                 if target_filename.exists():
-                    target_filename.unlink()                    
+                    target_filename.unlink()
                 source_filename.rename ( str(target_filename) )
 
-        
+
     def debug_one_iter(self):
         images = []
-        for generator in self.generator_list:        
+        for generator in self.generator_list:
             for i,batch in enumerate(next(generator)):
                 if len(batch.shape) == 4:
                     images.append( batch[0] )
-        
+
         return image_utils.equalize_and_stack_square (images)
-        
+
     def generate_next_sample(self):
         return [next(generator) for generator in self.generator_list]
 
     def train_one_iter(self):
-        sample = self.generate_next_sample()        
-        iter_time = time.time()        
-        losses = self.onTrainOneIter(sample, self.generator_list)        
+        sample = self.generate_next_sample()
+        iter_time = time.time()
+        losses = self.onTrainOneIter(sample, self.generator_list)
         iter_time = time.time() - iter_time
         self.last_sample = sample
-        
+
         self.loss_history.append ( [float(loss[1]) for loss in losses] )
 
         if self.write_preview_history:
-            if self.iter % 10 == 0:            
+            if self.iter % 10 == 0:
                 preview = self.get_static_preview()
                 preview_lh = ModelBase.get_loss_history_preview(self.loss_history, self.iter, preview.shape[1], preview.shape[2])
                 img = (np.concatenate ( [preview_lh, preview], axis=0 ) * 255).astype(np.uint8)
-                cv2_imwrite ( str (self.preview_history_path / ('%.6d.jpg' %( self.iter) )), img )     
-                
+                cv2_imwrite ( str (self.preview_history_path / ('%.6d.jpg' %( self.iter) )), img )
+
         self.iter += 1
 
         time_str = time.strftime("[%H:%M:%S]")
@@ -383,40 +383,40 @@ class ModelBase(object):
             loss_string += " %s:%.3f" % (loss_name, loss_value)
 
         return loss_string
-        
+
     def pass_one_iter(self):
-        self.last_sample = self.generate_next_sample()     
-        
+        self.last_sample = self.generate_next_sample()
+
     def finalize(self):
         nnlib.finalize_all()
-                
+
     def is_first_run(self):
         return self.iter == 0
-        
+
     def is_debug(self):
         return self.debug
-        
+
     def set_batch_size(self, batch_size):
         self.batch_size = batch_size
-        
+
     def get_batch_size(self):
         return self.batch_size
-        
+
     def get_iter(self):
         return self.iter
-        
+
     def get_loss_history(self):
         return self.loss_history
- 
+
     def set_training_data_generators (self, generator_list):
         self.generator_list = generator_list
-        
+
     def get_training_data_generators (self):
         return self.generator_list
-        
+
     def get_model_root_path(self):
         return self.model_path
-        
+
     def get_strpath_storage_for_file(self, filename):
         if self.device_args['force_gpu_idx'] == -1:
             return str( self.model_path / ( self.get_model_name() + '_' + filename) )
@@ -424,65 +424,65 @@ class ModelBase(object):
             return str( self.model_path / ( str(self.device_args['force_gpu_idx']) + '_' + self.get_model_name() + '_' + filename) )
 
     def set_vram_batch_requirements (self, d):
-        #example d = {2:2,3:4,4:8,5:16,6:32,7:32,8:32,9:48} 
+        #example d = {2:2,3:4,4:8,5:16,6:32,7:32,8:32,9:48}
         keys = [x for x in d.keys()]
-        
+
         if self.device_config.cpu_only:
             if self.batch_size == 0:
                 self.batch_size = 2
         else:
-            if self.batch_size == 0:        
+            if self.batch_size == 0:
                 for x in keys:
                     if self.device_config.gpu_vram_gb[0] <= x:
                         self.batch_size = d[x]
                         break
-                        
+
                 if self.batch_size == 0:
                     self.batch_size = d[ keys[-1] ]
-                    
+
     @staticmethod
     def get_loss_history_preview(loss_history, iter, w, c):
         loss_history = np.array (loss_history.copy())
-                
+
         lh_height = 100
         lh_img = np.ones ( (lh_height,w,c) ) * 0.1
         loss_count = len(loss_history[0])
         lh_len = len(loss_history)
-        
-        l_per_col = lh_len / w                
+
+        l_per_col = lh_len / w
         plist_max = [   [   max (0.0, loss_history[int(col*l_per_col)][p],
-                                            *[  loss_history[i_ab][p] 
-                                                for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )                                         
+                                            *[  loss_history[i_ab][p]
+                                                for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )
                                              ]
-                                ) 
+                                )
                             for p in range(loss_count)
-                        ]  
+                        ]
                         for col in range(w)
                     ]
 
         plist_min = [   [   min (plist_max[col][p], loss_history[int(col*l_per_col)][p],
-                                            *[  loss_history[i_ab][p] 
-                                                for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )                                         
+                                            *[  loss_history[i_ab][p]
+                                                for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )
                                              ]
-                                ) 
-                            for p in range(loss_count) 
-                        ]  
-                        for col in range(w) 
+                                )
+                            for p in range(loss_count)
+                        ]
+                        for col in range(w)
                     ]
 
         plist_abs_max = np.mean(loss_history[ len(loss_history) // 5 : ]) * 2
 
         for col in range(0, w):
-            for p in range(0,loss_count): 
+            for p in range(0,loss_count):
                 point_color = [1.0]*c
                 point_color[0:3] = colorsys.hsv_to_rgb ( p * (1.0/loss_count), 1.0, 1.0 )
-                
+
                 ph_max = int ( (plist_max[col][p] / plist_abs_max) * (lh_height-1) )
                 ph_max = np.clip( ph_max, 0, lh_height-1 )
-                
+
                 ph_min = int ( (plist_min[col][p] / plist_abs_max) * (lh_height-1) )
                 ph_min = np.clip( ph_min, 0, lh_height-1 )
-                
+
                 for ph in range(ph_min, ph_max+1):
                     lh_img[ (lh_height-ph-1), col ] = point_color
 
@@ -490,11 +490,11 @@ class ModelBase(object):
         lh_line_height = (lh_height-1)/lh_lines
         for i in range(0,lh_lines+1):
             lh_img[ int(i*lh_line_height), : ] = (0.8,)*c
-            
+
         last_line_t = int((lh_lines-1)*lh_line_height)
         last_line_b = int(lh_lines*lh_line_height)
-        
+
         lh_text = 'Iter: %d' % (iter) if iter != 0 else ''
-        
+
         lh_img[last_line_t:last_line_b, 0:w] += image_utils.get_text_image (  (w,last_line_b-last_line_t,c), lh_text, color=[0.8]*c )
-        return lh_img
+        return lh_img

models/Model_DF/Model.py

@@ -9,22 +9,22 @@ from interact import interact as io
 class Model(ModelBase):
 
     #override
-    def onInitializeOptions(self, is_first_run, ask_override):        
+    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.")
         else:
             self.options['pixel_loss'] = self.options.get('pixel_loss', False)
-            
+
     #override
     def onInitialize(self):
         exec(nnlib.import_all(), locals(), globals())
         self.set_vram_batch_requirements( {4.5:4} )
-                
+
         ae_input_layer = Input(shape=(128, 128, 3))
         mask_layer = Input(shape=(128, 128, 1)) #same as output
-        
-        self.encoder, self.decoder_src, self.decoder_dst = self.Build(ae_input_layer)    
+
+        self.encoder, self.decoder_src, self.decoder_dst = self.Build(ae_input_layer)
 
         if not self.is_first_run():
             weights_to_load = [  [self.encoder    , 'encoder.h5'],
@@ -38,39 +38,39 @@ class Model(ModelBase):
 
         self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] )
         self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] )
-  
+
         if self.is_training_mode:
             f = SampleProcessor.TypeFlags
-            self.set_training_data_generators ([            
-                    SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, 
-                                                                     debug=self.is_debug(), batch_size=self.batch_size, 
+            self.set_training_data_generators ([
+                    SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None,
+                                                                     debug=self.is_debug(), batch_size=self.batch_size,
                         sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ),
-                        output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], 
-                                              [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], 
+                        output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
+                                              [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
                                               [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] ),
-                                              
-                    SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, 
+
+                    SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
                         sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
-                        output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], 
-                                              [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], 
+                        output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
+                                              [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
                                               [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] )
                 ])
     #override
-    def onSave(self):        
+    def onSave(self):
         self.save_weights_safe( [[self.encoder,     'encoder.h5'],
                                  [self.decoder_src, 'decoder_src.h5'],
                                  [self.decoder_dst, 'decoder_dst.h5']] )
-        
+
     #override
     def onTrainOneIter(self, sample, generators_list):
         warped_src, target_src, target_src_mask = sample[0]
-        warped_dst, target_dst, target_dst_mask = sample[1]    
-  
+        warped_dst, target_dst, target_dst_mask = sample[1]
+
         loss_src = self.autoencoder_src.train_on_batch( [warped_src, target_src_mask], [target_src, target_src_mask] )
         loss_dst = self.autoencoder_dst.train_on_batch( [warped_dst, target_dst_mask], [target_dst, target_dst_mask] )
-        
+
         return ( ('loss_src', loss_src[0]), ('loss_dst', loss_dst[0]) )
-        
+
 
     #override
     def onGetPreview(self, sample):
@@ -78,64 +78,64 @@ class Model(ModelBase):
         test_A_m = sample[0][2][0:4] #first 4 samples
         test_B   = sample[1][1][0:4]
         test_B_m = sample[1][2][0:4]
-        
-        AA, mAA = self.autoencoder_src.predict([test_A, test_A_m])                                       
+
+        AA, mAA = self.autoencoder_src.predict([test_A, test_A_m])
         AB, mAB = self.autoencoder_src.predict([test_B, test_B_m])
         BB, mBB = self.autoencoder_dst.predict([test_B, test_B_m])
-        
+
         mAA = np.repeat ( mAA, (3,), -1)
         mAB = np.repeat ( mAB, (3,), -1)
         mBB = np.repeat ( mBB, (3,), -1)
-        
+
         st = []
         for i in range(0, len(test_A)):
             st.append ( np.concatenate ( (
                 test_A[i,:,:,0:3],
                 AA[i],
                 #mAA[i],
-                test_B[i,:,:,0:3], 
-                BB[i], 
-                #mBB[i],                
+                test_B[i,:,:,0:3],
+                BB[i],
+                #mBB[i],
                 AB[i],
                 #mAB[i]
                 ), axis=1) )
-            
+
         return [ ('DF', np.concatenate ( st, axis=0 ) ) ]
-    
+
     def predictor_func (self, face):
-        
+
         face_128_bgr = face[...,0:3]
         face_128_mask = np.expand_dims(face[...,3],-1)
-        
+
         x, mx = self.autoencoder_src.predict ( [ np.expand_dims(face_128_bgr,0), np.expand_dims(face_128_mask,0) ] )
         x, mx = x[0], mx[0]
-        
+
         return np.concatenate ( (x,mx), -1 )
-        
+
     #override
     def get_converter(self):
-        from converters import ConverterMasked            
-        return ConverterMasked(self.predictor_func, 
-                               predictor_input_size=128, 
-                               output_size=128, 
-                               face_type=FaceType.FULL, 
+        from converters import ConverterMasked
+        return ConverterMasked(self.predictor_func,
+                               predictor_input_size=128,
+                               output_size=128,
+                               face_type=FaceType.FULL,
                                base_erode_mask_modifier=30,
                                base_blur_mask_modifier=0)
-        
+
     def Build(self, input_layer):
         exec(nnlib.code_import_all, locals(), globals())
-    
+
         def downscale (dim):
             def func(x):
                 return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x))
-            return func 
-            
+            return func
+
         def upscale (dim):
             def func(x):
                 return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
-            return func   
-            
-        def Encoder(input_layer):            
+            return func
+
+        def Encoder(input_layer):
             x = input_layer
             x = downscale(128)(x)
             x = downscale(256)(x)
@@ -146,7 +146,7 @@ class Model(ModelBase):
             x = Dense(8 * 8 * 512)(x)
             x = Reshape((8, 8, 512))(x)
             x = upscale(512)(x)
-                
+
             return Model(input_layer, x)
 
         def Decoder():
@@ -155,15 +155,15 @@ class Model(ModelBase):
             x = upscale(512)(x)
             x = upscale(256)(x)
             x = upscale(128)(x)
-            
+
             y = input_  #mask decoder
             y = upscale(512)(y)
             y = upscale(256)(y)
             y = upscale(128)(y)
-                
+
             x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
             y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
-            
+
             return Model(input_, [x,y])
-        
-        return Encoder(input_layer), Decoder(), Decoder()
+
+        return Encoder(input_layer), Decoder(), Decoder()

models/Model_DF/__init__.py

@@ -1 +1 @@
-from .Model import Model
+from .Model import Model

models/Model_FANSegmentator/Model.py

@@ -10,13 +10,13 @@ from interact import interact as io
 class Model(ModelBase):
 
     def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs, 
-                            ask_write_preview_history=False, 
+        super().__init__(*args, **kwargs,
+                            ask_write_preview_history=False,
                             ask_target_iter=False,
                             ask_sort_by_yaw=False,
                             ask_random_flip=False,
                             ask_src_scale_mod=False)
-        
+
     #override
     def onInitialize(self):
         exec(nnlib.import_all(), locals(), globals())
@@ -24,33 +24,33 @@ class Model(ModelBase):
 
         self.resolution = 256
         self.face_type = FaceType.FULL
-        
-        self.fan_seg = FANSegmentator(self.resolution, 
-                                      FaceType.toString(self.face_type), 
+
+        self.fan_seg = FANSegmentator(self.resolution,
+                                      FaceType.toString(self.face_type),
                                       load_weights=not self.is_first_run(),
                                       weights_file_root=self.get_model_root_path() )
 
         if self.is_training_mode:
             f = SampleProcessor.TypeFlags
             f_type = f.FACE_ALIGN_FULL
-            
-            self.set_training_data_generators ([    
-                    SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, 
-                            sample_process_options=SampleProcessor.Options(random_flip=True, normalize_tanh = True ), 
+
+            self.set_training_data_generators ([
+                    SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
+                            sample_process_options=SampleProcessor.Options(random_flip=True, normalize_tanh = True ),
                             output_sample_types=[ [f.TRANSFORMED | f_type | f.MODE_BGR_SHUFFLE, self.resolution],
                                                   [f.TRANSFORMED | f_type | f.MODE_M | f.FACE_MASK_FULL, self.resolution]
                                                 ]),
-                                                
-                    SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, 
-                            sample_process_options=SampleProcessor.Options(random_flip=True, normalize_tanh = True ), 
+
+                    SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
+                            sample_process_options=SampleProcessor.Options(random_flip=True, normalize_tanh = True ),
                             output_sample_types=[ [f.TRANSFORMED | f_type | f.MODE_BGR_SHUFFLE, self.resolution]
                                                 ])
                                                ])
-                
+
     #override
-    def onSave(self):        
+    def onSave(self):
         self.fan_seg.save_weights()
-        
+
     #override
     def onTrainOneIter(self, generators_samples, generators_list):
         target_src, target_src_mask = generators_samples[0]
@@ -58,20 +58,20 @@ class Model(ModelBase):
         loss = self.fan_seg.train_on_batch( [target_src], [target_src_mask] )
 
         return ( ('loss', loss), )
-        
+
     #override
     def onGetPreview(self, sample):
         test_A   = sample[0][0][0:4] #first 4 samples
         test_B   = sample[1][0][0:4] #first 4 samples
-        
+
         mAA = self.fan_seg.extract_from_bgr([test_A])
         mBB = self.fan_seg.extract_from_bgr([test_B])
-        
+
         test_A, test_B, = [ np.clip( (x + 1.0)/2.0, 0.0, 1.0)  for x in [test_A, test_B] ]
-        
+
         mAA = np.repeat ( mAA, (3,), -1)
         mBB = np.repeat ( mBB, (3,), -1)
-        
+
         st = []
         for i in range(0, len(test_A)):
             st.append ( np.concatenate ( (
@@ -79,7 +79,7 @@ class Model(ModelBase):
                 mAA[i],
                 test_A[i,:,:,0:3]*mAA[i],
                 ), axis=1) )
-                
+
         st2 = []
         for i in range(0, len(test_B)):
             st2.append ( np.concatenate ( (
@@ -87,7 +87,7 @@ class Model(ModelBase):
                 mBB[i],
                 test_B[i,:,:,0:3]*mBB[i],
                 ), axis=1) )
-                
+
         return [ ('FANSegmentator', np.concatenate ( st, axis=0 ) ),
                  ('never seen', np.concatenate ( st2, axis=0 ) ),
                  ]

models/Model_FANSegmentator/__init__.py

@@ -1 +1 @@
-from .Model import Model
+from .Model import Model

models/Model_H128/Model.py

@@ -9,7 +9,7 @@ from interact import interact as io
 class Model(ModelBase):
 
     #override
-    def onInitializeOptions(self, is_first_run, ask_override):        
+    def onInitializeOptions(self, is_first_run, ask_override):
         if is_first_run:
             self.options['lighter_ae'] = io.input_bool ("Use lightweight autoencoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight autoencoder is faster, requires less VRAM, sacrificing overall quality. If your GPU VRAM <= 4, you should to choose this option.")
         else:
@@ -17,18 +17,18 @@ class Model(ModelBase):
             if 'created_vram_gb' in self.options.keys():
                 self.options.pop ('created_vram_gb')
             self.options['lighter_ae'] = self.options.get('lighter_ae', default_lighter_ae)
-            
+
         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.")
         else:
             self.options['pixel_loss'] = self.options.get('pixel_loss', False)
-            
+
     #override
     def onInitialize(self):
-        exec(nnlib.import_all(), locals(), globals())        
+        exec(nnlib.import_all(), locals(), globals())
         self.set_vram_batch_requirements( {2.5:4} )
-                
+
         bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build( self.options['lighter_ae'] )
         if not self.is_first_run():
             weights_to_load = [  [self.encoder    , 'encoder.h5'],
@@ -36,120 +36,120 @@ class Model(ModelBase):
                                  [self.decoder_dst, 'decoder_dst.h5']
                               ]
             self.load_weights_safe(weights_to_load)
-            
+
         input_src_bgr = Input(bgr_shape)
         input_src_mask = Input(mask_shape)
         input_dst_bgr = Input(bgr_shape)
         input_dst_mask = Input(mask_shape)
 
-        rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) )        
+        rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) )
         rec_dst_bgr, rec_dst_mask = self.decoder_dst( self.encoder(input_dst_bgr) )
 
         self.ae = Model([input_src_bgr,input_src_mask,input_dst_bgr,input_dst_mask], [rec_src_bgr, rec_src_mask, rec_dst_bgr, rec_dst_mask] )
-            
+
         self.ae.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
                         loss=[ DSSIMMSEMaskLoss(input_src_mask, is_mse=self.options['pixel_loss']), 'mae', DSSIMMSEMaskLoss(input_dst_mask, is_mse=self.options['pixel_loss']), 'mae' ] )
-  
+
         self.src_view = K.function([input_src_bgr],[rec_src_bgr, rec_src_mask])
         self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask])
-        
+
         if self.is_training_mode:
             f = SampleProcessor.TypeFlags
-            self.set_training_data_generators ([            
-                    SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, 
-                                                                     debug=self.is_debug(), batch_size=self.batch_size, 
-                            sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), 
-                            output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128], 
-                                                  [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128], 
+            self.set_training_data_generators ([
+                    SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None,
+                                                                     debug=self.is_debug(), batch_size=self.batch_size,
+                            sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ),
+                            output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128],
+                                                  [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128],
                                                   [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 128] ] ),
-                                                  
-                    SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, 
-                            sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), 
-                            output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128], 
-                                                  [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128], 
+
+                    SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
+                            sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
+                            output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128],
+                                                  [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128],
                                                   [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 128] ] )
                 ])
-            
+
     #override
-    def onSave(self):        
+    def onSave(self):
         self.save_weights_safe( [[self.encoder,     'encoder.h5'],
                                  [self.decoder_src, 'decoder_src.h5'],
                                  [self.decoder_dst, 'decoder_dst.h5']] )
-        
+
     #override
     def onTrainOneIter(self, sample, generators_list):
         warped_src, target_src, target_src_mask = sample[0]
-        warped_dst, target_dst, target_dst_mask = sample[1]    
+        warped_dst, target_dst, target_dst_mask = sample[1]
 
         total, loss_src_bgr, loss_src_mask, loss_dst_bgr, loss_dst_mask = self.ae.train_on_batch( [warped_src, target_src_mask, warped_dst, target_dst_mask], [target_src, target_src_mask, target_dst, target_dst_mask] )
 
         return ( ('loss_src', loss_src_bgr), ('loss_dst', loss_dst_bgr) )
-        
+
     #override
     def onGetPreview(self, sample):
         test_A   = sample[0][1][0:4] #first 4 samples
         test_A_m = sample[0][2][0:4] #first 4 samples
         test_B   = sample[1][1][0:4]
         test_B_m = sample[1][2][0:4]
-        
-        AA, mAA = self.src_view([test_A])                                       
+
+        AA, mAA = self.src_view([test_A])
         AB, mAB = self.src_view([test_B])
         BB, mBB = self.dst_view([test_B])
-        
+
         mAA = np.repeat ( mAA, (3,), -1)
         mAB = np.repeat ( mAB, (3,), -1)
         mBB = np.repeat ( mBB, (3,), -1)
-        
+
         st = []
         for i in range(0, len(test_A)):
             st.append ( np.concatenate ( (
                 test_A[i,:,:,0:3],
                 AA[i],
                 #mAA[i],
-                test_B[i,:,:,0:3], 
-                BB[i], 
-                #mBB[i],                
+                test_B[i,:,:,0:3],
+                BB[i],
+                #mBB[i],
                 AB[i],
                 #mAB[i]
                 ), axis=1) )
-            
+
         return [ ('H128', np.concatenate ( st, axis=0 ) ) ]
 
-    def predictor_func (self, face):        
+    def predictor_func (self, face):
         face_128_bgr = face[...,0:3]
         face_128_mask = np.expand_dims(face[...,3],-1)
-        
+
         x, mx = self.src_view ( [ np.expand_dims(face_128_bgr,0) ] )
-        x, mx = x[0], mx[0]     
-        
+        x, mx = x[0], mx[0]
+
         return np.concatenate ( (x,mx), -1 )
 
     #override
     def get_converter(self):
         from converters import ConverterMasked
-        return ConverterMasked(self.predictor_func, 
-                               predictor_input_size=128, 
-                               output_size=128, 
+        return ConverterMasked(self.predictor_func,
+                               predictor_input_size=128,
+                               output_size=128,
                                face_type=FaceType.HALF,
                                base_erode_mask_modifier=100,
                                base_blur_mask_modifier=100)
-    
+
     def Build(self, lighter_ae):
         exec(nnlib.code_import_all, locals(), globals())
-        
+
         bgr_shape = (128, 128, 3)
         mask_shape = (128, 128, 1)
-        
+
         def downscale (dim):
             def func(x):
                 return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x))
-            return func 
-            
+            return func
+
         def upscale (dim):
             def func(x):
                 return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
-            return func   
-            
+            return func
+
         def Encoder(input_shape):
             input_layer = Input(input_shape)
             x = input_layer
@@ -171,7 +171,7 @@ class Model(ModelBase):
                 x = Dense(8 * 8 * 256)(x)
                 x = Reshape((8, 8, 256))(x)
                 x = upscale(256)(x)
-                
+
             return Model(input_layer, x)
 
         def Decoder():
@@ -181,7 +181,7 @@ class Model(ModelBase):
                 x = upscale(512)(x)
                 x = upscale(256)(x)
                 x = upscale(128)(x)
-                
+
                 y = input_  #mask decoder
                 y = upscale(512)(y)
                 y = upscale(256)(y)
@@ -192,16 +192,16 @@ class Model(ModelBase):
                 x = upscale(256)(x)
                 x = upscale(128)(x)
                 x = upscale(64)(x)
-                
+
                 y = input_  #mask decoder
                 y = upscale(256)(y)
                 y = upscale(128)(y)
                 y = upscale(64)(y)
-            
+
             x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
             y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
-            
-            
+
+
             return Model(input_, [x,y])
-            
+
         return bgr_shape, mask_shape, Encoder(bgr_shape), Decoder(), Decoder()

models/Model_H128/__init__.py

@@ -1 +1 @@
-from .Model import Model
+from .Model import Model

models/Model_H64/Model.py

@@ -9,7 +9,7 @@ from interact import interact as io
 class Model(ModelBase):
 
     #override
-    def onInitializeOptions(self, is_first_run, ask_override):        
+    def onInitializeOptions(self, is_first_run, ask_override):
         if is_first_run:
             self.options['lighter_ae'] = io.input_bool ("Use lightweight autoencoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight autoencoder is faster, requires less VRAM, sacrificing overall quality. If your GPU VRAM <= 4, you should to choose this option.")
         else:
@@ -17,141 +17,141 @@ class Model(ModelBase):
             if 'created_vram_gb' in self.options.keys():
                 self.options.pop ('created_vram_gb')
             self.options['lighter_ae'] = self.options.get('lighter_ae', default_lighter_ae)
-            
+
         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.")
         else:
             self.options['pixel_loss'] = self.options.get('pixel_loss', False)
-            
+
     #override
     def onInitialize(self):
         exec(nnlib.import_all(), locals(), globals())
         self.set_vram_batch_requirements( {1.5:4} )
-        
-        
+
+
         bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build(self.options['lighter_ae'])
-        
+
         if not self.is_first_run():
             weights_to_load = [  [self.encoder    , 'encoder.h5'],
                                  [self.decoder_src, 'decoder_src.h5'],
                                  [self.decoder_dst, 'decoder_dst.h5']
                               ]
             self.load_weights_safe(weights_to_load)
-            
+
         input_src_bgr = Input(bgr_shape)
         input_src_mask = Input(mask_shape)
         input_dst_bgr = Input(bgr_shape)
         input_dst_mask = Input(mask_shape)
-        
-        rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) )        
+
+        rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) )
         rec_dst_bgr, rec_dst_mask = self.decoder_dst( self.encoder(input_dst_bgr) )
 
         self.ae = Model([input_src_bgr,input_src_mask,input_dst_bgr,input_dst_mask], [rec_src_bgr, rec_src_mask, rec_dst_bgr, rec_dst_mask] )
-      
+
         self.ae.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[ DSSIMMSEMaskLoss(input_src_mask, is_mse=self.options['pixel_loss']), 'mae', DSSIMMSEMaskLoss(input_dst_mask, is_mse=self.options['pixel_loss']), 'mae' ] )
-  
+
         self.src_view = K.function([input_src_bgr],[rec_src_bgr, rec_src_mask])
         self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask])
-  
+
         if self.is_training_mode:
             f = SampleProcessor.TypeFlags
-            self.set_training_data_generators ([    
-                    SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, 
-                                                                     debug=self.is_debug(), batch_size=self.batch_size, 
-                            sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), 
-                            output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64], 
-                                                  [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64], 
+            self.set_training_data_generators ([
+                    SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None,
+                                                                     debug=self.is_debug(), batch_size=self.batch_size,
+                            sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ),
+                            output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
+                                                  [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
                                                   [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 64] ] ),
-                                                  
+
                     SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
-                            sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), 
-                            output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64], 
-                                                  [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64], 
+                            sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
+                            output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
+                                                  [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
                                                   [f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 64] ] )
                 ])
-                
+
     #override
-    def onSave(self):        
+    def onSave(self):
         self.save_weights_safe( [[self.encoder,     'encoder.h5'],
                                  [self.decoder_src, 'decoder_src.h5'],
                                  [self.decoder_dst, 'decoder_dst.h5']] )
-        
+
     #override
     def onTrainOneIter(self, sample, generators_list):
         warped_src, target_src, target_src_full_mask = sample[0]
-        warped_dst, target_dst, target_dst_full_mask = sample[1]    
+        warped_dst, target_dst, target_dst_full_mask = sample[1]
 
         total, loss_src_bgr, loss_src_mask, loss_dst_bgr, loss_dst_mask = self.ae.train_on_batch( [warped_src, target_src_full_mask, warped_dst, target_dst_full_mask], [target_src, target_src_full_mask, target_dst, target_dst_full_mask] )
 
         return ( ('loss_src', loss_src_bgr), ('loss_dst', loss_dst_bgr) )
-        
+
     #override
     def onGetPreview(self, sample):
         test_A   = sample[0][1][0:4] #first 4 samples
         test_A_m = sample[0][2][0:4]
         test_B   = sample[1][1][0:4]
         test_B_m = sample[1][2][0:4]
-        
-        AA, mAA = self.src_view([test_A])                                       
+
+        AA, mAA = self.src_view([test_A])
         AB, mAB = self.src_view([test_B])
         BB, mBB = self.dst_view([test_B])
-        
+
         mAA = np.repeat ( mAA, (3,), -1)
         mAB = np.repeat ( mAB, (3,), -1)
         mBB = np.repeat ( mBB, (3,), -1)
-        
+
         st = []
         for i in range(0, len(test_A)):
             st.append ( np.concatenate ( (
                 test_A[i,:,:,0:3],
                 AA[i],
                 #mAA[i],
-                test_B[i,:,:,0:3], 
-                BB[i], 
-                #mBB[i],                
+                test_B[i,:,:,0:3],
+                BB[i],
+                #mBB[i],
                 AB[i],
                 #mAB[i]
                 ), axis=1) )
-            
+
         return [ ('H64', np.concatenate ( st, axis=0 ) ) ]
 
     def predictor_func (self, face):
-        
+
         face_64_bgr = face[...,0:3]
         face_64_mask = np.expand_dims(face[...,3],-1)
-        
+
         x, mx = self.src_view ( [ np.expand_dims(face_64_bgr,0) ] )
-        x, mx = x[0], mx[0]     
-        
+        x, mx = x[0], mx[0]
+
         return np.concatenate ( (x,mx), -1 )
 
     #override
     def get_converter(self):
         from converters import ConverterMasked
         return ConverterMasked(self.predictor_func,
-                               predictor_input_size=64, 
-                               output_size=64, 
-                               face_type=FaceType.HALF, 
+                               predictor_input_size=64,
+                               output_size=64,
+                               face_type=FaceType.HALF,
                                base_erode_mask_modifier=100,
                                base_blur_mask_modifier=100)
-        
+
     def Build(self, lighter_ae):
         exec(nnlib.code_import_all, locals(), globals())
-        
+
         bgr_shape = (64, 64, 3)
         mask_shape = (64, 64, 1)
-        
+
         def downscale (dim):
             def func(x):
                 return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x))
-            return func 
-            
+            return func
+
         def upscale (dim):
             def func(x):
                 return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
-            return func            
-        
+            return func
+
         def Encoder(input_shape):
             input_layer = Input(input_shape)
             x = input_layer
@@ -183,23 +183,23 @@ class Model(ModelBase):
                 x = upscale(512)(x)
                 x = upscale(256)(x)
                 x = upscale(128)(x)
-                
+
             else:
                 input_ = Input(shape=(8, 8, 256))
-                
-                x = input_                
+
+                x = input_
                 x = upscale(256)(x)
                 x = upscale(128)(x)
                 x = upscale(64)(x)
-                
+
             y = input_  #mask decoder
             y = upscale(256)(y)
             y = upscale(128)(y)
             y = upscale(64)(y)
-            
+
             x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
             y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
-            
+
             return Model(input_, [x,y])
-            
-        return bgr_shape, mask_shape, Encoder(bgr_shape), Decoder(), Decoder()
+
+        return bgr_shape, mask_shape, Encoder(bgr_shape), Decoder(), Decoder()

models/Model_H64/__init__.py

@@ -1 +1 @@
-from .Model import Model
+from .Model import Model

models/Model_LIAEF128/Model.py

@@ -9,13 +9,13 @@ from interact import interact as io
 class Model(ModelBase):
 
     #override
-    def onInitializeOptions(self, is_first_run, ask_override):        
+    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.")
         else:
             self.options['pixel_loss'] = self.options.get('pixel_loss', False)
-            
+
     #override
     def onInitialize(self):
         exec(nnlib.import_all(), locals(), globals())
@@ -25,7 +25,7 @@ class Model(ModelBase):
         mask_layer = Input(shape=(128, 128, 1)) #same as output
 
         self.encoder, self.decoder, self.inter_B, self.inter_AB = self.Build(ae_input_layer)
-       
+
         if not self.is_first_run():
             weights_to_load = [  [self.encoder,  'encoder.h5'],
                                  [self.decoder,  'decoder.h5'],
@@ -39,46 +39,46 @@ class Model(ModelBase):
         B = self.inter_B(code)
         self.autoencoder_src = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([AB, AB])) )
         self.autoencoder_dst = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([B, AB])) )
-            
+
         self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] )
         self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] )
-  
+
         if self.is_training_mode:
             f = SampleProcessor.TypeFlags
-            self.set_training_data_generators ([         
-                    
-            
-                    SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, 
+            self.set_training_data_generators ([
+
+
+                    SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None,
                                                                      debug=self.is_debug(), batch_size=self.batch_size,
-                        sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), 
-                        output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], 
-                                              [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], 
+                        sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ),
+                        output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
+                                              [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
                                               [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] ),
-                                              
-                    SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, 
-                        sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), 
-                        output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], 
-                                              [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128], 
+
+                    SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
+                        sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
+                        output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
+                                              [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
                                               [f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] )
                 ])
-            
+
     #override
     def onSave(self):
         self.save_weights_safe( [[self.encoder, 'encoder.h5'],
                                  [self.decoder, 'decoder.h5'],
                                  [self.inter_B, 'inter_B.h5'],
                                  [self.inter_AB, 'inter_AB.h5']] )
-        
+
     #override
     def onTrainOneIter(self, sample, generators_list):
         warped_src, target_src, target_src_mask = sample[0]
-        warped_dst, target_dst, target_dst_mask = sample[1]    
+        warped_dst, target_dst, target_dst_mask = sample[1]
 
         loss_src = self.autoencoder_src.train_on_batch( [warped_src, target_src_mask], [target_src, target_src_mask] )
         loss_dst = self.autoencoder_dst.train_on_batch( [warped_dst, target_dst_mask], [target_dst, target_dst_mask] )
-        
+
         return ( ('loss_src', loss_src[0]), ('loss_dst', loss_dst[0]) )
-        
+
 
     #override
     def onGetPreview(self, sample):
@@ -86,63 +86,63 @@ class Model(ModelBase):
         test_A_m = sample[0][2][0:4] #first 4 samples
         test_B   = sample[1][1][0:4]
         test_B_m = sample[1][2][0:4]
-        
-        AA, mAA = self.autoencoder_src.predict([test_A, test_A_m])                                       
+
+        AA, mAA = self.autoencoder_src.predict([test_A, test_A_m])
         AB, mAB = self.autoencoder_src.predict([test_B, test_B_m])
         BB, mBB = self.autoencoder_dst.predict([test_B, test_B_m])
-        
+
         mAA = np.repeat ( mAA, (3,), -1)
         mAB = np.repeat ( mAB, (3,), -1)
         mBB = np.repeat ( mBB, (3,), -1)
-        
+
         st = []
         for i in range(0, len(test_A)):
             st.append ( np.concatenate ( (
                 test_A[i,:,:,0:3],
                 AA[i],
                 #mAA[i],
-                test_B[i,:,:,0:3], 
-                BB[i], 
-                #mBB[i],                
+                test_B[i,:,:,0:3],
+                BB[i],
+                #mBB[i],
                 AB[i],
                 #mAB[i]
                 ), axis=1) )
-            
+
         return [ ('LIAEF128', np.concatenate ( st, axis=0 ) ) ]
-    
+
     def predictor_func (self, face):
-        
+
         face_128_bgr = face[...,0:3]
         face_128_mask = np.expand_dims(face[...,3],-1)
-        
+
         x, mx = self.autoencoder_src.predict ( [ np.expand_dims(face_128_bgr,0), np.expand_dims(face_128_mask,0) ] )
         x, mx = x[0], mx[0]
-        
+
         return np.concatenate ( (x,mx), -1 )
-        
+
     #override
     def get_converter(self):
         from converters import ConverterMasked
-        return ConverterMasked(self.predictor_func, 
-                               predictor_input_size=128, 
-                               output_size=128, 
-                               face_type=FaceType.FULL, 
+        return ConverterMasked(self.predictor_func,
+                               predictor_input_size=128,
+                               output_size=128,
+                               face_type=FaceType.FULL,
                                base_erode_mask_modifier=30,
                                base_blur_mask_modifier=0)
-      
+
     def Build(self, input_layer):
         exec(nnlib.code_import_all, locals(), globals())
-        
+
         def downscale (dim):
             def func(x):
                 return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x))
-            return func 
-            
+            return func
+
         def upscale (dim):
             def func(x):
                 return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
-            return func   
-            
+            return func
+
         def Encoder():
             x = input_layer
             x = downscale(128)(x)
@@ -161,20 +161,20 @@ class Model(ModelBase):
             x = upscale(512)(x)
             return Model(input_layer, x)
 
-        def Decoder(): 
+        def Decoder():
             input_ = Input(shape=(16, 16, 1024))
             x = input_
             x = upscale(512)(x)
             x = upscale(256)(x)
             x = upscale(128)(x)
             x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
-            
+
             y = input_  #mask decoder
             y = upscale(512)(y)
             y = upscale(256)(y)
             y = upscale(128)(y)
             y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid' )(y)
-            
+
             return Model(input_, [x,y])
 
-        return Encoder(), Decoder(), Intermediate(), Intermediate()
+        return Encoder(), Decoder(), Intermediate(), Intermediate()

models/Model_LIAEF128/__init__.py

@@ -1 +1 @@
-from .Model import Model
+from .Model import Model

models/Model_SAE/Model.py

@@ -9,51 +9,51 @@ from interact import interact as io
 #SAE - Styled AutoEncoder
 class SAEModel(ModelBase):
 
-    encoderH5 = 'encoder.h5'    
+    encoderH5 = 'encoder.h5'
     inter_BH5 = 'inter_B.h5'
     inter_ABH5 = 'inter_AB.h5'
     decoderH5 = 'decoder.h5'
     decodermH5 = 'decoderm.h5'
-    
+
     decoder_srcH5 = 'decoder_src.h5'
     decoder_srcmH5 = 'decoder_srcm.h5'
     decoder_dstH5 = 'decoder_dst.h5'
     decoder_dstmH5 = 'decoder_dstm.h5'
-    
+
     #override
     def onInitializeOptions(self, is_first_run, ask_override):
         yn_str = {True:'y',False:'n'}
-        
+
         default_resolution = 128
         default_archi = 'df'
         default_face_type = 'f'
-        
+
         if is_first_run:
             resolution = io.input_int("Resolution ( 64-256 ?:help skip:128) : ", default_resolution, help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.")
-            resolution = np.clip (resolution, 64, 256)            
+            resolution = np.clip (resolution, 64, 256)
             while np.modf(resolution / 16)[0] != 0.0:
                 resolution -= 1
             self.options['resolution'] = resolution
-            
-            self.options['face_type'] = io.input_str ("Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h','f'], help_message="Half face has better resolution, but covers less area of cheeks.").lower()            
+
+            self.options['face_type'] = io.input_str ("Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h','f'], help_message="Half face has better resolution, but covers less area of cheeks.").lower()
             self.options['learn_mask'] = io.input_bool ("Learn mask? (y/n, ?:help skip:y) : ", True, help_message="Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted. Model with style values can be learned without mask and produce same quality result.")
         else:
             self.options['resolution'] = self.options.get('resolution', default_resolution)
             self.options['face_type'] = self.options.get('face_type', default_face_type)
             self.options['learn_mask'] = self.options.get('learn_mask', True)
-            
-            
+
+
         if is_first_run and 'tensorflow' in self.device_config.backend:
             def_optimizer_mode = self.options.get('optimizer_mode', 1)
             self.options['optimizer_mode'] = io.input_int ("Optimizer mode? ( 1,2,3 ?:help skip:%d) : " % (def_optimizer_mode), def_optimizer_mode, help_message="1 - no changes. 2 - allows you to train x2 bigger network consuming RAM. 3 - allows you to train x3 bigger network consuming huge amount of RAM and slower, depends on CPU power.")
         else:
             self.options['optimizer_mode'] = self.options.get('optimizer_mode', 1)
-        
+
         if is_first_run:
             self.options['archi'] = io.input_str ("AE architecture (df, liae, vg ?:help skip:%s) : " % (default_archi) , default_archi, ['df','liae','vg'], help_message="'df' keeps faces more natural. 'liae' can fix overly different face shapes. 'vg' - currently testing.").lower()
         else:
             self.options['archi'] = self.options.get('archi', default_archi)
-        
+
         default_ae_dims = 256 if self.options['archi'] == 'liae' else 512
         default_ed_ch_dims = 42
         def_ca_weights = False
@@ -65,31 +65,31 @@ class SAEModel(ModelBase):
             self.options['ae_dims'] = self.options.get('ae_dims', default_ae_dims)
             self.options['ed_ch_dims'] = self.options.get('ed_ch_dims', default_ed_ch_dims)
             self.options['ca_weights'] = self.options.get('ca_weights', def_ca_weights)
-            
+
         if is_first_run:
             self.options['lighter_encoder'] = io.input_bool ("Use lightweight encoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight encoder is 35% faster, requires less VRAM, but sacrificing overall quality.")
-            
+
             if self.options['archi'] != 'vg':
                 self.options['multiscale_decoder'] = io.input_bool ("Use multiscale decoder? (y/n, ?:help skip:n) : ", False, help_message="Multiscale decoder helps to get better details.")
         else:
             self.options['lighter_encoder'] = self.options.get('lighter_encoder', False)
-            
+
             if self.options['archi'] != 'vg':
                 self.options['multiscale_decoder'] = self.options.get('multiscale_decoder', False)
-            
-        default_face_style_power = 0.0        
-        default_bg_style_power = 0.0  
+
+        default_face_style_power = 0.0
+        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 15-25k iters to enhance fine details and decrease face jitter.")
-        
+
             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, 
-                                                                               help_message="Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes."), 0.0, 100.0 )            
-                            
+            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,
+                                                                               help_message="Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k iters, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes."), 0.0, 100.0 )
+
             default_bg_style_power = default_bg_style_power if is_first_run else self.options.get('bg_style_power', default_bg_style_power)
-            self.options['bg_style_power'] = np.clip ( io.input_number("Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_bg_style_power), default_bg_style_power, 
-                                                                               help_message="Learn to transfer image around face. This can make face more like dst."), 0.0, 100.0 )            
+            self.options['bg_style_power'] = np.clip ( io.input_number("Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_bg_style_power), default_bg_style_power,
+                                                                               help_message="Learn to transfer image around face. This can make face more like dst."), 0.0, 100.0 )
         else:
             self.options['pixel_loss'] = self.options.get('pixel_loss', False)
             self.options['face_style_power'] = self.options.get('face_style_power', default_face_style_power)
@@ -100,7 +100,7 @@ class SAEModel(ModelBase):
         exec(nnlib.import_all(), locals(), globals())
         SAEModel.initialize_nn_functions()
         self.set_vram_batch_requirements({1.5:4})
-        
+
         resolution = self.options['resolution']
         ae_dims = self.options['ae_dims']
         ed_ch_dims = self.options['ed_ch_dims']
@@ -108,13 +108,13 @@ class SAEModel(ModelBase):
         mask_shape = (resolution, resolution, 1)
 
         self.ms_count = ms_count = 3 if (self.options['archi'] != 'vg' and self.options['multiscale_decoder']) else 1
-        
+
         masked_training = True
-        
+
         warped_src = Input(bgr_shape)
         target_src = Input(bgr_shape)
         target_srcm = Input(mask_shape)
-        
+
         warped_dst = Input(bgr_shape)
         target_dst = Input(bgr_shape)
         target_dstm = Input(mask_shape)
@@ -124,27 +124,28 @@ class SAEModel(ModelBase):
         target_dst_ar  = [ Input ( ( bgr_shape[0] // (2**i) ,)*2 + (bgr_shape[-1],) ) for i in range(ms_count-1, -1, -1)]
         target_dstm_ar = [ Input ( ( mask_shape[0] // (2**i) ,)*2 + (mask_shape[-1],) ) for i in range(ms_count-1, -1, -1)]
 
-    
+        use_bn = True
+
         models_list = []
         weights_to_load = []
         if self.options['archi'] == 'liae':
-            self.encoder = modelify(SAEModel.LIAEEncFlow(resolution, self.options['lighter_encoder'], ed_ch_dims=ed_ch_dims)  ) (Input(bgr_shape))
-            
-            enc_output_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ] 
-            
-            self.inter_B = modelify(SAEModel.LIAEInterFlow(resolution, ae_dims=ae_dims)) (enc_output_Inputs)
-            self.inter_AB = modelify(SAEModel.LIAEInterFlow(resolution, ae_dims=ae_dims)) (enc_output_Inputs)
-            
-            inter_output_Inputs = [ Input( np.array(K.int_shape(x)[1:])*(1,1,2) ) for x in self.inter_B.outputs ] 
+            self.encoder = modelify(SAEModel.LIAEEncFlow(resolution, self.options['lighter_encoder'], ed_ch_dims=ed_ch_dims, use_bn=use_bn)  ) (Input(bgr_shape))
+
+            enc_output_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ]
+
+            self.inter_B = modelify(SAEModel.LIAEInterFlow(resolution, ae_dims=ae_dims, use_bn=use_bn)) (enc_output_Inputs)
+            self.inter_AB = modelify(SAEModel.LIAEInterFlow(resolution, ae_dims=ae_dims, use_bn=use_bn)) (enc_output_Inputs)
+
+            inter_output_Inputs = [ Input( np.array(K.int_shape(x)[1:])*(1,1,2) ) for x in self.inter_B.outputs ]
+
+            self.decoder = modelify(SAEModel.LIAEDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2, multiscale_count=self.ms_count, use_bn=use_bn )) (inter_output_Inputs)
 
-            self.decoder = modelify(SAEModel.LIAEDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2, multiscale_count=self.ms_count )) (inter_output_Inputs)
-            
             models_list += [self.encoder, self.inter_B, self.inter_AB, self.decoder]
-            
+
             if self.options['learn_mask']:
-                self.decoderm = modelify(SAEModel.LIAEDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (inter_output_Inputs)
+                self.decoderm = modelify(SAEModel.LIAEDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5), use_bn=use_bn )) (inter_output_Inputs)
                 models_list += [self.decoderm]
-                
+
             if not self.is_first_run():
                 weights_to_load += [  [self.encoder , 'encoder.h5'],
                                       [self.inter_B , 'inter_B.h5'],
@@ -153,22 +154,22 @@ class SAEModel(ModelBase):
                                     ]
                 if self.options['learn_mask']:
                     weights_to_load += [ [self.decoderm, 'decoderm.h5'] ]
-            
-            warped_src_code = self.encoder (warped_src)            
+
+            warped_src_code = self.encoder (warped_src)
             warped_src_inter_AB_code = self.inter_AB (warped_src_code)
-            warped_src_inter_code = Concatenate()([warped_src_inter_AB_code,warped_src_inter_AB_code])            
-            
+            warped_src_inter_code = Concatenate()([warped_src_inter_AB_code,warped_src_inter_AB_code])
+
             warped_dst_code = self.encoder (warped_dst)
             warped_dst_inter_B_code = self.inter_B (warped_dst_code)
             warped_dst_inter_AB_code = self.inter_AB (warped_dst_code)
             warped_dst_inter_code = Concatenate()([warped_dst_inter_B_code,warped_dst_inter_AB_code])
 
             warped_src_dst_inter_code = Concatenate()([warped_dst_inter_AB_code,warped_dst_inter_AB_code])
-            
+
             pred_src_src = self.decoder(warped_src_inter_code)
-            pred_dst_dst = self.decoder(warped_dst_inter_code)            
+            pred_dst_dst = self.decoder(warped_dst_inter_code)
             pred_src_dst = self.decoder(warped_src_dst_inter_code)
-            
+
             if self.options['learn_mask']:
                 pred_src_srcm = self.decoderm(warped_src_inter_code)
                 pred_dst_dstm = self.decoderm(warped_dst_inter_code)
@@ -177,18 +178,18 @@ class SAEModel(ModelBase):
         elif self.options['archi'] == 'df':
             self.encoder = modelify(SAEModel.DFEncFlow(resolution, self.options['lighter_encoder'], ae_dims=ae_dims, ed_ch_dims=ed_ch_dims)  ) (Input(bgr_shape))
 
-            dec_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ] 
-            
+            dec_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ]
+
             self.decoder_src = modelify(SAEModel.DFDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2, multiscale_count=self.ms_count )) (dec_Inputs)
             self.decoder_dst = modelify(SAEModel.DFDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2, multiscale_count=self.ms_count )) (dec_Inputs)
-            
+
             models_list += [self.encoder, self.decoder_src, self.decoder_dst]
-            
+
             if self.options['learn_mask']:
                 self.decoder_srcm = modelify(SAEModel.DFDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (dec_Inputs)
                 self.decoder_dstm = modelify(SAEModel.DFDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (dec_Inputs)
                 models_list += [self.decoder_srcm, self.decoder_dstm]
-                
+
             if not self.is_first_run():
                 weights_to_load += [  [self.encoder    , 'encoder.h5'],
                                       [self.decoder_src, 'decoder_src.h5'],
@@ -198,37 +199,37 @@ class SAEModel(ModelBase):
                     weights_to_load += [ [self.decoder_srcm, 'decoder_srcm.h5'],
                                          [self.decoder_dstm, 'decoder_dstm.h5'],
                                        ]
-            
-            
-                
-                
-                
+
+
+
+
+
             warped_src_code = self.encoder (warped_src)
             warped_dst_code = self.encoder (warped_dst)
             pred_src_src = self.decoder_src(warped_src_code)
             pred_dst_dst = self.decoder_dst(warped_dst_code)
             pred_src_dst = self.decoder_src(warped_dst_code)
-            
+
             if self.options['learn_mask']:
                 pred_src_srcm = self.decoder_srcm(warped_src_code)
                 pred_dst_dstm = self.decoder_dstm(warped_dst_code)
                 pred_src_dstm = self.decoder_srcm(warped_dst_code)
-                
+
         elif self.options['archi'] == 'vg':
             self.encoder = modelify(SAEModel.VGEncFlow(resolution, self.options['lighter_encoder'], ae_dims=ae_dims, ed_ch_dims=ed_ch_dims)  ) (Input(bgr_shape))
 
-            dec_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ] 
-            
+            dec_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ]
+
             self.decoder_src = modelify(SAEModel.VGDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2 )) (dec_Inputs)
             self.decoder_dst = modelify(SAEModel.VGDecFlow (bgr_shape[2],ed_ch_dims=ed_ch_dims//2 )) (dec_Inputs)
-            
+
             models_list += [self.encoder, self.decoder_src, self.decoder_dst]
-            
+
             if self.options['learn_mask']:
                 self.decoder_srcm = modelify(SAEModel.VGDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (dec_Inputs)
                 self.decoder_dstm = modelify(SAEModel.VGDecFlow (mask_shape[2],ed_ch_dims=int(ed_ch_dims/1.5) )) (dec_Inputs)
                 models_list += [self.decoder_srcm, self.decoder_dstm]
-                
+
             if not self.is_first_run():
                 weights_to_load += [  [self.encoder    , 'encoder.h5'],
                                       [self.decoder_src, 'decoder_src.h5'],
@@ -238,7 +239,7 @@ class SAEModel(ModelBase):
                     weights_to_load += [ [self.decoder_srcm, 'decoder_srcm.h5'],
                                          [self.decoder_dstm, 'decoder_dstm.h5'],
                                        ]
-            
+
             warped_src_code = self.encoder (warped_src)
             warped_dst_code = self.encoder (warped_dst)
             pred_src_src = self.decoder_src(warped_src_code)
@@ -250,30 +251,30 @@ class SAEModel(ModelBase):
                 pred_src_srcm = self.decoder_srcm(warped_src_code)
                 pred_dst_dstm = self.decoder_dstm(warped_dst_code)
                 pred_src_dstm = self.decoder_srcm(warped_dst_code)
-                
+
         if self.is_first_run() and self.options['ca_weights']:
             io.log_info ("Initializing CA weights...")
             conv_weights_list = []
             for model in models_list:
                 for layer in model.layers:
                     if type(layer) == Conv2D:
-                        conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights            
+                        conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights
             CAInitializerMP ( conv_weights_list )
-            
-                
+
+
         pred_src_src, pred_dst_dst, pred_src_dst, = [ [x] if type(x) != list else x for x in [pred_src_src, pred_dst_dst, pred_src_dst, ] ]
-        
+
         if self.options['learn_mask']:
             pred_src_srcm, pred_dst_dstm, pred_src_dstm = [ [x] if type(x) != list else x for x in [pred_src_srcm, pred_dst_dstm, pred_src_dstm] ]
 
         target_srcm_blurred_ar = [ gaussian_blur( max(1, K.int_shape(x)[1] // 32) )(x) for x in target_srcm_ar]
-        target_srcm_sigm_ar = target_srcm_blurred_ar #[ x / 2.0 + 0.5 for x in target_srcm_blurred_ar] 
-        target_srcm_anti_sigm_ar = [ 1.0 - x for x in target_srcm_sigm_ar] 
-    
+        target_srcm_sigm_ar = target_srcm_blurred_ar #[ x / 2.0 + 0.5 for x in target_srcm_blurred_ar]
+        target_srcm_anti_sigm_ar = [ 1.0 - x for x in target_srcm_sigm_ar]
+
         target_dstm_blurred_ar = [ gaussian_blur( max(1, K.int_shape(x)[1] // 32) )(x) for x in target_dstm_ar]
-        target_dstm_sigm_ar = target_dstm_blurred_ar#[ x / 2.0 + 0.5 for x in target_dstm_blurred_ar] 
-        target_dstm_anti_sigm_ar = [ 1.0 - x for x in target_dstm_sigm_ar] 
-        
+        target_dstm_sigm_ar = target_dstm_blurred_ar#[ x / 2.0 + 0.5 for x in target_dstm_blurred_ar]
+        target_dstm_anti_sigm_ar = [ 1.0 - x for x in target_dstm_sigm_ar]
+
         target_src_sigm_ar = target_src_ar#[ x + 1 for x in target_src_ar]
         target_dst_sigm_ar = target_dst_ar#[ x + 1 for x in target_dst_ar]
 
@@ -284,32 +285,32 @@ class SAEModel(ModelBase):
         target_src_masked_ar = [ target_src_sigm_ar[i]*target_srcm_sigm_ar[i]  for i in range(len(target_src_sigm_ar))]
         target_dst_masked_ar = [ target_dst_sigm_ar[i]*target_dstm_sigm_ar[i]  for i in range(len(target_dst_sigm_ar))]
         target_dst_anti_masked_ar = [ target_dst_sigm_ar[i]*target_dstm_anti_sigm_ar[i]  for i in range(len(target_dst_sigm_ar))]
-    
+
         pred_src_src_masked_ar = [ pred_src_src_sigm_ar[i] * target_srcm_sigm_ar[i]  for i in range(len(pred_src_src_sigm_ar))]
         pred_dst_dst_masked_ar = [ pred_dst_dst_sigm_ar[i] * target_dstm_sigm_ar[i]  for i in range(len(pred_dst_dst_sigm_ar))]
-        
+
         target_src_masked_ar_opt = target_src_masked_ar if masked_training else target_src_sigm_ar
         target_dst_masked_ar_opt = target_dst_masked_ar if masked_training else target_dst_sigm_ar
-        
+
         pred_src_src_masked_ar_opt = pred_src_src_masked_ar if masked_training else pred_src_src_sigm_ar
         pred_dst_dst_masked_ar_opt = pred_dst_dst_masked_ar if masked_training else pred_dst_dst_sigm_ar
-        
+
         psd_target_dst_masked_ar = [ pred_src_dst_sigm_ar[i]*target_dstm_sigm_ar[i]  for i in range(len(pred_src_dst_sigm_ar))]
         psd_target_dst_anti_masked_ar = [ pred_src_dst_sigm_ar[i]*target_dstm_anti_sigm_ar[i]  for i in range(len(pred_src_dst_sigm_ar))]
-        
-        if self.is_training_mode:            
+
+        if self.is_training_mode:
             self.src_dst_opt      = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999, tf_cpu_mode=self.options['optimizer_mode']-1)
             self.src_dst_mask_opt = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999, tf_cpu_mode=self.options['optimizer_mode']-1)
 
-            if self.options['archi'] == 'liae':          
+            if self.options['archi'] == 'liae':
                 src_dst_loss_train_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoder.trainable_weights
                 if self.options['learn_mask']:
                     src_dst_mask_loss_train_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoderm.trainable_weights
-            else:   
+            else:
                 src_dst_loss_train_weights = self.encoder.trainable_weights + self.decoder_src.trainable_weights + self.decoder_dst.trainable_weights
                 if self.options['learn_mask']:
                     src_dst_mask_loss_train_weights = self.encoder.trainable_weights + self.decoder_srcm.trainable_weights + self.decoder_dstm.trainable_weights
-            
+
             if not self.options['pixel_loss']:
                 src_loss_batch = sum([ (  100*K.square( dssim(kernel_size=int(resolution/11.6),max_value=1.0)( target_src_masked_ar_opt[i],  pred_src_src_masked_ar_opt[i] ) )) for i in range(len(target_src_masked_ar_opt)) ])
             else:
@@ -318,9 +319,9 @@ class SAEModel(ModelBase):
             src_loss = K.mean(src_loss_batch)
 
             face_style_power = self.options['face_style_power']  / 100.0
-            
-            if face_style_power != 0:    
-                src_loss += style_loss(gaussian_blur_radius=resolution//16, loss_weight=face_style_power, wnd_size=0)( psd_target_dst_masked_ar[-1], target_dst_masked_ar[-1] ) 
+
+            if face_style_power != 0:
+                src_loss += style_loss(gaussian_blur_radius=resolution//16, loss_weight=face_style_power, wnd_size=0)( psd_target_dst_masked_ar[-1], target_dst_masked_ar[-1] )
 
             bg_style_power = self.options['bg_style_power'] / 100.0
             if bg_style_power != 0:
@@ -334,32 +335,32 @@ class SAEModel(ModelBase):
                 dst_loss_batch = sum([ (  100*K.square(dssim(kernel_size=int(resolution/11.6),max_value=1.0)( target_dst_masked_ar_opt[i],  pred_dst_dst_masked_ar_opt[i] ) )) for i in range(len(target_dst_masked_ar_opt)) ])
             else:
                 dst_loss_batch = sum([ K.mean ( 100*K.square( target_dst_masked_ar_opt[i] - pred_dst_dst_masked_ar_opt[i] ), axis=[1,2,3]) for i in range(len(target_dst_masked_ar_opt)) ])
-                
+
             dst_loss = K.mean(dst_loss_batch)
 
-            feed = [warped_src, warped_dst]            
+            feed = [warped_src, warped_dst]
             feed += target_src_ar[::-1]
             feed += target_srcm_ar[::-1]
             feed += target_dst_ar[::-1]
             feed += target_dstm_ar[::-1]
-    
+
             self.src_dst_train = K.function (feed,[src_loss,dst_loss], self.src_dst_opt.get_updates(src_loss+dst_loss, src_dst_loss_train_weights) )
 
             if self.options['learn_mask']:
                 src_mask_loss = sum([ K.mean(K.square(target_srcm_ar[-1]-pred_src_srcm[-1])) for i in range(len(target_srcm_ar)) ])
                 dst_mask_loss = sum([ K.mean(K.square(target_dstm_ar[-1]-pred_dst_dstm[-1])) for i in range(len(target_dstm_ar)) ])
-                
-                feed = [ warped_src, warped_dst]    
+
+                feed = [ warped_src, warped_dst]
                 feed += target_srcm_ar[::-1]
-                feed += target_dstm_ar[::-1]            
-                
+                feed += target_dstm_ar[::-1]
+
                 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]])
             else:
-                self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_src_dst[-1] ] ) 
-                
+                self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_src_dst[-1] ] )
+
             self.load_weights_safe(weights_to_load)#, [ [self.src_dst_opt, 'src_dst_opt'], [self.src_dst_mask_opt, 'src_dst_mask_opt']])
         else:
             self.load_weights_safe(weights_to_load)
@@ -367,30 +368,30 @@ class SAEModel(ModelBase):
                 self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1], pred_src_dstm[-1] ])
             else:
                 self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1] ])
-                
-        
-        if self.is_training_mode:            
+
+
+        if self.is_training_mode:
             self.src_sample_losses = []
             self.dst_sample_losses = []
-            
-            f = SampleProcessor.TypeFlags            
+
+            f = SampleProcessor.TypeFlags
             face_type = f.FACE_ALIGN_FULL if self.options['face_type'] == 'f' else f.FACE_ALIGN_HALF
-            
-            output_sample_types=[ [f.WARPED_TRANSFORMED | face_type | f.MODE_BGR, resolution] ]            
+
+            output_sample_types=[ [f.WARPED_TRANSFORMED | face_type | f.MODE_BGR, resolution] ]
             output_sample_types += [ [f.TRANSFORMED | face_type | f.MODE_BGR, resolution // (2**i) ] for i in range(ms_count)]
             output_sample_types += [ [f.TRANSFORMED | face_type | f.MODE_M | f.FACE_MASK_FULL, resolution // (2**i) ] for i in range(ms_count)]
-            
-            self.set_training_data_generators ([            
-                    SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None, 
-                                                                     debug=self.is_debug(), batch_size=self.batch_size, 
-                        sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ), 
+
+            self.set_training_data_generators ([
+                    SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None,
+                                                                     debug=self.is_debug(), batch_size=self.batch_size,
+                        sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ),
                         output_sample_types=output_sample_types ),
-                                              
+
                     SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
-                        sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, ), 
+                        sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, ),
                         output_sample_types=output_sample_types )
                 ])
-            
+
     #override
     def onSave(self):
         opt_ar = [ [self.src_dst_opt,      'src_dst_opt'],
@@ -413,10 +414,10 @@ class SAEModel(ModelBase):
            if self.options['learn_mask']:
                 ar += [ [self.decoder_srcm, 'decoder_srcm.h5'],
                         [self.decoder_dstm, 'decoder_dstm.h5'] ]
-               
+
         self.save_weights_safe(ar)
-       
-    
+
+
     #override
     def onTrainOneIter(self, generators_samples, generators_list):
         src_samples  = generators_samples[0]
@@ -425,17 +426,17 @@ class SAEModel(ModelBase):
         feed = [src_samples[0], dst_samples[0] ] + \
                src_samples[1:1+self.ms_count*2] + \
                dst_samples[1:1+self.ms_count*2]
-               
+
         src_loss, dst_loss, = self.src_dst_train (feed)
-            
+
         if self.options['learn_mask']:
             feed = [ src_samples[0], dst_samples[0] ] + \
                    src_samples[1+self.ms_count:1+self.ms_count*2] + \
                    dst_samples[1+self.ms_count:1+self.ms_count*2]
             src_mask_loss, dst_mask_loss, = self.src_dst_mask_train (feed)
-        
+
         return ( ('src_loss', src_loss), ('dst_loss', dst_loss) )
-        
+
 
     #override
     def onGetPreview(self, sample):
@@ -454,33 +455,33 @@ class SAEModel(ModelBase):
         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],)            
+            #    ar += (SDM[i],)
             st.append ( np.concatenate ( ar, axis=1) )
-        
+
         return [ ('SAE', np.concatenate (st, axis=0 )), ]
-    
+
     def predictor_func (self, face):
-    
+
         prd = [ x[0] for x in self.AE_convert ( [ face[np.newaxis,:,:,0:3] ] ) ]
 
         if not self.options['learn_mask']:
-            prd += [ face[...,3:4] ] 
-        
+            prd += [ face[...,3:4] ]
+
         return np.concatenate ( prd, -1 )
-        
+
     #override
     def get_converter(self):
         base_erode_mask_modifier = 30 if self.options['face_type'] == 'f' else 100
         base_blur_mask_modifier = 0 if self.options['face_type'] == 'f' else 100
-        
+
         default_erode_mask_modifier = 0
         default_blur_mask_modifier = 100 if (self.options['face_style_power'] or self.options['bg_style_power']) and \
                                                 self.options['face_type'] == 'f' else 0
-        
+
         face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
-        
+
         from converters import ConverterMasked
-        return ConverterMasked(self.predictor_func, 
+        return ConverterMasked(self.predictor_func,
                                predictor_input_size=self.options['resolution'],
                                output_size=self.options['resolution'],
                                face_type=face_type,
@@ -490,30 +491,34 @@ class SAEModel(ModelBase):
                                default_erode_mask_modifier=default_erode_mask_modifier,
                                default_blur_mask_modifier=default_blur_mask_modifier,
                                clip_hborder_mask_per=0.0625 if self.options['face_type'] == 'f' else 0)
-                               
+
     @staticmethod
     def initialize_nn_functions():
         exec (nnlib.import_all(), locals(), globals())
-        
+
+        def BatchNorm():
+            return BatchNormalization(axis=-1,  gamma_initializer=RandomNormal(1., 0.02) )
+
+
         class ResidualBlock(object):
             def __init__(self, filters, kernel_size=3, padding='same', use_reflection_padding=False):
                 self.filters = filters
                 self.kernel_size = kernel_size
                 self.padding = padding #if not use_reflection_padding else 'valid'
                 self.use_reflection_padding = use_reflection_padding
-                
+
             def __call__(self, inp):
                 var_x = LeakyReLU(alpha=0.2)(inp)
-                
+
                 #if self.use_reflection_padding:
                 #    #var_x = ReflectionPadding2D(stride=1, kernel_size=kernel_size)(var_x)
 
-                var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding, kernel_initializer=RandomNormal(0, 0.02) )(var_x)                
+                var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding, kernel_initializer=RandomNormal(0, 0.02) )(var_x)
                 var_x = LeakyReLU(alpha=0.2)(var_x)
-                
+
                 #if self.use_reflection_padding:
                 #    #var_x = ReflectionPadding2D(stride=1, kernel_size=kernel_size)(var_x)
-                    
+
                 var_x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding, kernel_initializer=RandomNormal(0, 0.02) )(var_x)
                 var_x = Scale(gamma_init=keras.initializers.Constant(value=0.1))(var_x)
                 var_x = Add()([var_x, inp])
@@ -521,99 +526,108 @@ class SAEModel(ModelBase):
                 return var_x
         SAEModel.ResidualBlock = ResidualBlock
 
-        def downscale (dim):
+        def downscale (dim, use_bn=False):
             def func(x):
-                return LeakyReLU(0.1)(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x))
-            return func 
+                if use_bn:
+                    return LeakyReLU(0.1)(BatchNorm()(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02), use_bias=False)(x)))
+                else:
+                    return LeakyReLU(0.1)(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x))
+            return func
         SAEModel.downscale = downscale
-        
-        def downscale_sep (dim):
+
+        def downscale_sep (dim, use_bn=False):
             def func(x):
-                return LeakyReLU(0.1)(SeparableConv2D(dim, kernel_size=5, strides=2, padding='same', depthwise_initializer=RandomNormal(0, 0.02), pointwise_initializer=RandomNormal(0, 0.02) )(x))
-            return func 
+                if use_bn:
+                    return LeakyReLU(0.1)(BatchNorm()(SeparableConv2D(dim, kernel_size=5, strides=2, padding='same', depthwise_initializer=RandomNormal(0, 0.02), pointwise_initializer=RandomNormal(0, 0.02), use_bias=False )(x)))
+                else:
+                    return LeakyReLU(0.1)(SeparableConv2D(dim, kernel_size=5, strides=2, padding='same', depthwise_initializer=RandomNormal(0, 0.02), pointwise_initializer=RandomNormal(0, 0.02) )(x))
+            return func
         SAEModel.downscale_sep = downscale_sep
-        
-        def upscale (dim):
+
+        def upscale (dim, use_bn=False):
             def func(x):
-                return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0, 0.02) )(x)))
-            return func 
+                if use_bn:
+                    return SubpixelUpscaler()(LeakyReLU(0.1)(BatchNorm()(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0,0.02), use_bias=False )(x))))
+                else:
+                    return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0, 0.02) )(x)))
+            return func
         SAEModel.upscale = upscale
-        
+
         def to_bgr (output_nc):
             def func(x):
                 return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid', kernel_initializer=RandomNormal(0, 0.02) )(x)
             return func
         SAEModel.to_bgr = to_bgr
-        
-            
+
+
     @staticmethod
-    def LIAEEncFlow(resolution, light_enc, ed_ch_dims=42):
+    def LIAEEncFlow(resolution, light_enc, ed_ch_dims=42, use_bn=False):
         exec (nnlib.import_all(), locals(), globals())
         upscale = SAEModel.upscale
         downscale = SAEModel.downscale
         downscale_sep = SAEModel.downscale_sep
-    
+
         def func(input):
             ed_dims = K.int_shape(input)[-1]*ed_ch_dims
-            
-            x = input            
+
+            x = input
             x = downscale(ed_dims)(x)
-            if not light_enc:                
-                x = downscale(ed_dims*2)(x)
-                x = downscale(ed_dims*4)(x)
-                x = downscale(ed_dims*8)(x)
+            if not light_enc:
+                x = downscale(ed_dims*2, use_bn=use_bn)(x)
+                x = downscale(ed_dims*4, use_bn=use_bn)(x)
+                x = downscale(ed_dims*8, use_bn=use_bn)(x)
             else:
-                x = downscale_sep(ed_dims*2)(x)
-                x = downscale(ed_dims*4)(x)
-                x = downscale_sep(ed_dims*8)(x)
-            
-            x = Flatten()(x)               
+                x = downscale_sep(ed_dims*2, use_bn=use_bn)(x)
+                x = downscale(ed_dims*4, use_bn=use_bn)(x)
+                x = downscale_sep(ed_dims*8, use_bn=use_bn)(x)
+
+            x = Flatten()(x)
             return x
         return func
-    
+
     @staticmethod
-    def LIAEInterFlow(resolution, ae_dims=256):
+    def LIAEInterFlow(resolution, ae_dims=256, use_bn=False):
         exec (nnlib.import_all(), locals(), globals())
         upscale = SAEModel.upscale
         lowest_dense_res=resolution // 16
-        
-        def func(input):   
+
+        def func(input):
             x = input[0]
             x = Dense(ae_dims)(x)
             x = Dense(lowest_dense_res * lowest_dense_res * ae_dims*2)(x)
             x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims*2))(x)
-            x = upscale(ae_dims*2)(x)
+            x = upscale(ae_dims*2, use_bn=use_bn)(x)
             return x
         return func
-        
+
     @staticmethod
-    def LIAEDecFlow(output_nc,ed_ch_dims=21, multiscale_count=1):
+    def LIAEDecFlow(output_nc,ed_ch_dims=21, multiscale_count=1, use_bn=False):
         exec (nnlib.import_all(), locals(), globals())
         upscale = SAEModel.upscale
         to_bgr = SAEModel.to_bgr
         ed_dims = output_nc * ed_ch_dims
-            
-        def func(input):   
+
+        def func(input):
             x = input[0]
-            
+
             outputs = []
-            x1     = upscale(ed_dims*8)( x )       
-            
+            x1     = upscale(ed_dims*8, use_bn=use_bn)( x )
+
             if multiscale_count >= 3:
-                outputs += [ to_bgr(output_nc) ( x1 ) ]  
-                
-            x2     = upscale(ed_dims*4)( x1 )    
-            
+                outputs += [ to_bgr(output_nc) ( x1 ) ]
+
+            x2     = upscale(ed_dims*4, use_bn=use_bn)( x1 )
+
             if multiscale_count >= 2:
                 outputs += [ to_bgr(output_nc) ( x2 ) ]
-                
-            x3     = upscale(ed_dims*2)( x2 )
-            
+
+            x3     = upscale(ed_dims*2, use_bn=use_bn)( x2 )
+
             outputs += [ to_bgr(output_nc) ( x3 ) ]
-        
+
             return outputs
         return func
-    
+
     @staticmethod
     def DFEncFlow(resolution, light_enc, ae_dims=512, ed_ch_dims=42):
         exec (nnlib.import_all(), locals(), globals())
@@ -622,11 +636,11 @@ class SAEModel(ModelBase):
         downscale_sep = SAEModel.downscale_sep
         lowest_dense_res = resolution // 16
 
-        def func(input):     
+        def func(input):
             x = input
-            
+
             ed_dims = K.int_shape(input)[-1]*ed_ch_dims
-            
+
             x = downscale(ed_dims)(x)
             if not light_enc:
                 x = downscale(ed_dims*2)(x)
@@ -636,15 +650,15 @@ class SAEModel(ModelBase):
                 x = downscale_sep(ed_dims*2)(x)
                 x = downscale_sep(ed_dims*4)(x)
                 x = downscale_sep(ed_dims*8)(x)
-    
+
             x = Dense(ae_dims)(Flatten()(x))
             x = Dense(lowest_dense_res * lowest_dense_res * ae_dims)(x)
             x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims))(x)
             x = upscale(ae_dims)(x)
-               
+
             return x
         return func
-    
+
     @staticmethod
     def DFDecFlow(output_nc, ed_ch_dims=21, multiscale_count=1):
         exec (nnlib.import_all(), locals(), globals())
@@ -652,29 +666,29 @@ class SAEModel(ModelBase):
         to_bgr = SAEModel.to_bgr
         ed_dims = output_nc * ed_ch_dims
 
-        def func(input):   
+        def func(input):
             x = input[0]
-            
+
             outputs = []
-            x1     = upscale(ed_dims*8)( x )       
-            
+            x1     = upscale(ed_dims*8)( x )
+
             if multiscale_count >= 3:
-                outputs += [ to_bgr(output_nc) ( x1 ) ]  
-                
-            x2     = upscale(ed_dims*4)( x1 )    
-            
+                outputs += [ to_bgr(output_nc) ( x1 ) ]
+
+            x2     = upscale(ed_dims*4)( x1 )
+
             if multiscale_count >= 2:
                 outputs += [ to_bgr(output_nc) ( x2 ) ]
-                
+
             x3     = upscale(ed_dims*2)( x2 )
-            
+
             outputs += [ to_bgr(output_nc) ( x3 ) ]
-        
-            return outputs            
+
+            return outputs
         return func
-        
-    
-            
+
+
+
     @staticmethod
     def VGEncFlow(resolution, light_enc, ae_dims=512, ed_ch_dims=42):
         exec (nnlib.import_all(), locals(), globals())
@@ -683,78 +697,78 @@ class SAEModel(ModelBase):
         downscale_sep = SAEModel.downscale_sep
         ResidualBlock = SAEModel.ResidualBlock
         lowest_dense_res = resolution // 16
-            
+
         def func(input):
             x = input
             ed_dims = K.int_shape(input)[-1]*ed_ch_dims
             while np.modf(ed_dims / 4)[0] != 0.0:
                 ed_dims -= 1
-            
+
             in_conv_filters = ed_dims# if resolution <= 128 else ed_dims + (resolution//128)*ed_ch_dims
-            
+
             x = tmp_x = Conv2D (in_conv_filters, kernel_size=5, strides=2, padding='same') (x)
 
             for _ in range ( 8 if light_enc else 16 ):
                 x = ResidualBlock(ed_dims)(x)
-                
+
             x = Add()([x, tmp_x])
 
             x = downscale(ed_dims)(x)
             x = SubpixelUpscaler()(x)
-            
+
             x = downscale(ed_dims)(x)
             x = SubpixelUpscaler()(x)
-            
-            x = downscale(ed_dims)(x)           
+
+            x = downscale(ed_dims)(x)
             if light_enc:
                 x = downscale_sep (ed_dims*2)(x)
             else:
                 x = downscale (ed_dims*2)(x)
-                
+
             x = downscale(ed_dims*4)(x)
-            
+
             if light_enc:
                 x = downscale_sep (ed_dims*8)(x)
             else:
                 x = downscale (ed_dims*8)(x)
-            
+
             x = Dense(ae_dims)(Flatten()(x))
             x = Dense(lowest_dense_res * lowest_dense_res * ae_dims)(x)
             x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims))(x)
             x = upscale(ae_dims)(x)
             return x
-            
+
         return func
-        
+
     @staticmethod
     def VGDecFlow(output_nc, ed_ch_dims=21, multiscale_count=1):
         exec (nnlib.import_all(), locals(), globals())
-        upscale = SAEModel.upscale        
+        upscale = SAEModel.upscale
         to_bgr = SAEModel.to_bgr
         ResidualBlock = SAEModel.ResidualBlock
         ed_dims = output_nc * ed_ch_dims
-  
+
         def func(input):
             x = input[0]
-            
+
             x = upscale( ed_dims*8 )(x)
             x = ResidualBlock( ed_dims*8 )(x)
-            
+
             x = upscale( ed_dims*4 )(x)
             x = ResidualBlock( ed_dims*4 )(x)
-            
+
             x = upscale( ed_dims*2 )(x)
             x = ResidualBlock( ed_dims*2 )(x)
-        
-            x = to_bgr(output_nc) (x)        
+
+            x = to_bgr(output_nc) (x)
             return x
-            
+
         return func
-        
+
 Model = SAEModel
 
 
-        
+
 # 'worst' sample booster gives no good result, or I dont know how to filter worst samples properly.
 #
 ##gathering array of sample_losses
@@ -769,7 +783,7 @@ Model = SAEModel
 #    idxs = (x[:,0][ np.argwhere ( b [ b > (np.mean(b)+np.std(b)) ] )[:,0] ]).astype(np.uint)
 #    generators_list[0].repeat_sample_idxs(idxs) #ask generator to repeat these sample idxs
 #    print ("src repeated %d" % (len(idxs)) )
-#    
+#
 #if len(self.dst_sample_losses) >= 128: #array is big enough
 #    #fetching idxs which losses are bigger than average
 #    x = np.array (self.dst_sample_losses)
@@ -777,4 +791,4 @@ Model = SAEModel
 #    b = x[:,1]
 #    idxs = (x[:,0][ np.argwhere ( b [ b > (np.mean(b)+np.std(b)) ] )[:,0] ]).astype(np.uint)
 #    generators_list[1].repeat_sample_idxs(idxs) #ask generator to repeat these sample idxs
-#    print ("dst repeated %d" % (len(idxs)) )
+#    print ("dst repeated %d" % (len(idxs)) )

models/Model_SAE/__init__.py

@@ -1 +1 @@
-from .Model import Model
+from .Model import Model

models/__init__.py

@@ -2,4 +2,4 @@ from .ModelBase import ModelBase
 
 def import_model(name):
     module = __import__('Model_'+name, globals(), locals(), [], 1)
-    return getattr(module, 'Model')
+    return getattr(module, 'Model')

nnlib/CAInitializer.py

@@ -61,7 +61,7 @@ def _scale_filters(filters, variance):
 def CAGenerateWeights ( shape, floatx, data_format, eps_std=0.05, seed=None ):
     if seed is not None:
         np.random.seed(seed)
-    
+
     fan_in, fan_out = _compute_fans(shape, data_format)
     variance = 2 / fan_in
 
@@ -109,4 +109,4 @@ def CAGenerateWeights ( shape, floatx, data_format, eps_std=0.05, seed=None ):
     # Format of array is now: filters, stack, row, column
     init = np.array(init)
     init = _scale_filters(init, variance)
-    return init.transpose(transpose_dimensions)
+    return init.transpose(transpose_dimensions)

nnlib/__init__.py

@@ -1 +1 @@
-from .nnlib import nnlib
+from .nnlib import nnlib

nnlib/device.py

@@ -5,11 +5,11 @@ from .pynvml import *
 
 #you can set DFL_TF_MIN_REQ_CAP manually for your build
 #the reason why we cannot check tensorflow.version is it requires import tensorflow
-tf_min_req_cap = int(os.environ.get("DFL_TF_MIN_REQ_CAP", 35)) 
-    
+tf_min_req_cap = int(os.environ.get("DFL_TF_MIN_REQ_CAP", 35))
+
 class device:
     backend = None
-    class Config():    
+    class Config():
         force_gpu_idx = -1
         multi_gpu = False
         force_gpu_idxs = None
@@ -22,36 +22,36 @@ class device:
         use_fp16 = False
         cpu_only = False
         backend = None
-        def __init__ (self, force_gpu_idx = -1, 
-                            multi_gpu = False, 
-                            force_gpu_idxs = None, 
+        def __init__ (self, force_gpu_idx = -1,
+                            multi_gpu = False,
+                            force_gpu_idxs = None,
                             choose_worst_gpu = False,
                             allow_growth = True,
                             use_fp16 = False,
                             cpu_only = False,
                             **in_options):
-            
+
             self.backend = device.backend
             self.use_fp16 = use_fp16
             self.cpu_only = cpu_only
-            
+
             if not self.cpu_only:
                 self.cpu_only = (self.backend == "tensorflow-cpu")
-            
+
             if not self.cpu_only:
                 self.force_gpu_idx = force_gpu_idx
                 self.multi_gpu = multi_gpu
                 self.force_gpu_idxs = force_gpu_idxs
-                self.choose_worst_gpu = choose_worst_gpu        
+                self.choose_worst_gpu = choose_worst_gpu
                 self.allow_growth = allow_growth
-          
+
                 self.gpu_idxs = []
 
                 if force_gpu_idxs is not None:
                     for idx in force_gpu_idxs.split(','):
                         idx = int(idx)
                         if device.isValidDeviceIdx(idx):
-                            self.gpu_idxs.append(idx)     
+                            self.gpu_idxs.append(idx)
                 else:
                     gpu_idx = force_gpu_idx if (force_gpu_idx >= 0 and device.isValidDeviceIdx(force_gpu_idx)) else device.getBestValidDeviceIdx() if not choose_worst_gpu else device.getWorstValidDeviceIdx()
                     if gpu_idx != -1:
@@ -61,10 +61,10 @@ class device:
                                 self.multi_gpu = False
                         else:
                             self.gpu_idxs = [gpu_idx]
-                            
+
                 self.cpu_only = (len(self.gpu_idxs) == 0)
- 
-            
+
+
             if not self.cpu_only:
                 self.gpu_names = []
                 self.gpu_compute_caps = []
@@ -78,10 +78,10 @@ class device:
                 self.gpu_names = ['CPU']
                 self.gpu_compute_caps = [99]
                 self.gpu_vram_gb = [0]
-            
+
             if self.cpu_only:
                 self.backend = "tensorflow-cpu"
-                
+
     @staticmethod
     def getValidDeviceIdxsEnumerator():
         if device.backend == "plaidML":
@@ -94,8 +94,8 @@ class device:
                     yield gpu_idx
         elif device.backend == "tensorflow-generic":
             yield 0
-        
-    
+
+
     @staticmethod
     def getValidDevicesWithAtLeastTotalMemoryGB(totalmemsize_gb):
         result = []
@@ -111,9 +111,9 @@ class device:
                     result.append (i)
         elif device.backend == "tensorflow-generic":
             return [0]
-                   
+
         return result
-        
+
     @staticmethod
     def getAllDevicesIdxsList():
         if device.backend == "plaidML":
@@ -121,8 +121,8 @@ class device:
         elif device.backend == "tensorflow":
             return [ *range(nvmlDeviceGetCount() ) ]
         elif device.backend == "tensorflow-generic":
-            return [0]  
-            
+            return [0]
+
     @staticmethod
     def getValidDevicesIdxsWithNamesList():
         if device.backend == "plaidML":
@@ -137,17 +137,17 @@ class device:
     @staticmethod
     def getDeviceVRAMTotalGb (idx):
         if device.backend == "plaidML":
-            if idx < plaidML_devices_count: 
+            if idx < plaidML_devices_count:
                 return plaidML_devices[idx]['globalMemSize'] / (1024*1024*1024)
         elif device.backend == "tensorflow":
-            if idx < nvmlDeviceGetCount():    
+            if idx < nvmlDeviceGetCount():
                 memInfo = nvmlDeviceGetMemoryInfo(  nvmlDeviceGetHandleByIndex(idx) )
                 return round ( memInfo.total / (1024*1024*1024) )
 
             return 0
         elif device.backend == "tensorflow-generic":
             return 2
-            
+
     @staticmethod
     def getBestValidDeviceIdx():
         if device.backend == "plaidML":
@@ -172,7 +172,7 @@ class device:
             return idx
         elif device.backend == "tensorflow-generic":
             return 0
-            
+
     @staticmethod
     def getWorstValidDeviceIdx():
         if device.backend == "plaidML":
@@ -197,7 +197,7 @@ class device:
             return idx
         elif device.backend == "tensorflow-generic":
             return 0
-        
+
     @staticmethod
     def isValidDeviceIdx(idx):
         if device.backend == "plaidML":
@@ -206,11 +206,11 @@ class device:
             return idx in [*device.getValidDeviceIdxsEnumerator()]
         elif device.backend == "tensorflow-generic":
             return (idx == 0)
-        
+
     @staticmethod
     def getDeviceIdxsEqualModel(idx):
         if device.backend == "plaidML":
-            result = []  
+            result = []
             idx_name = plaidML_devices[idx]['description']
             for i in device.getValidDeviceIdxsEnumerator():
                 if plaidML_devices[i]['description'] == idx_name:
@@ -218,7 +218,7 @@ class device:
 
             return result
         elif device.backend == "tensorflow":
-            result = []  
+            result = []
             idx_name = nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(idx)).decode()
             for i in device.getValidDeviceIdxsEnumerator():
                 if nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(i)).decode() == idx_name:
@@ -226,60 +226,60 @@ class device:
 
             return result
         elif device.backend == "tensorflow-generic":
-            return [0] if idx == 0 else []            
-            
+            return [0] if idx == 0 else []
+
     @staticmethod
     def getDeviceName (idx):
         if device.backend == "plaidML":
-            if idx < plaidML_devices_count:    
+            if idx < plaidML_devices_count:
                 return plaidML_devices[idx]['description']
         elif device.backend == "tensorflow":
-            if idx < nvmlDeviceGetCount():    
+            if idx < nvmlDeviceGetCount():
                 return nvmlDeviceGetName(nvmlDeviceGetHandleByIndex(idx)).decode()
         elif device.backend == "tensorflow-generic":
             if idx == 0:
                 return "Generic GeForce GPU"
-    
+
         return None
-        
+
     @staticmethod
     def getDeviceID (idx):
         if device.backend == "plaidML":
-            if idx < plaidML_devices_count:    
+            if idx < plaidML_devices_count:
                 return plaidML_devices[idx]['id'].decode()
 
-        return None   
-        
+        return None
+
     @staticmethod
     def getDeviceComputeCapability(idx):
         result = 0
         if device.backend == "plaidML":
             return 99
         elif device.backend == "tensorflow":
-            if idx < nvmlDeviceGetCount():    
+            if idx < nvmlDeviceGetCount():
                 result = nvmlDeviceGetCudaComputeCapability(nvmlDeviceGetHandleByIndex(idx))
         elif device.backend == "tensorflow-generic":
-            return 99 if idx == 0 else 0   
-            
+            return 99 if idx == 0 else 0
+
         return result[0] * 10 + result[1]
 
-        
+
 force_plaidML = os.environ.get("DFL_FORCE_PLAIDML", "0") == "1" #for OpenCL build , forcing using plaidML even if NVIDIA found
 force_tf_cpu = os.environ.get("DFL_FORCE_TF_CPU", "0") == "1"   #for OpenCL build , forcing using tf-cpu if plaidML failed
 has_nvml = False
 has_nvml_cap = False
 
-#use DFL_FORCE_HAS_NVIDIA_DEVICE=1 if 
+#use DFL_FORCE_HAS_NVIDIA_DEVICE=1 if
 #- your NVIDIA cannot be seen by OpenCL
 #- CUDA build of DFL
-has_nvidia_device = os.environ.get("DFL_FORCE_HAS_NVIDIA_DEVICE", "0") == "1" 
+has_nvidia_device = os.environ.get("DFL_FORCE_HAS_NVIDIA_DEVICE", "0") == "1"
 
 plaidML_devices = []
 
 # Using plaidML OpenCL backend to determine system devices and has_nvidia_device
-try:    
+try:
     os.environ['PLAIDML_EXPERIMENTAL'] = 'false' #this enables work plaidML without run 'plaidml-setup'
-    import plaidml        
+    import plaidml
     ctx = plaidml.Context()
     for d in plaidml.devices(ctx, return_all=True)[0]:
         details = json.loads(d.details)
@@ -288,13 +288,13 @@ try:
         if 'nvidia' in details['vendor'].lower():
             has_nvidia_device = True
         plaidML_devices += [ {'id':d.id,
-                              'globalMemSize' : int(details['globalMemSize']), 
+                              'globalMemSize' : int(details['globalMemSize']),
                               'description' : d.description.decode()
                            }]
     ctx.shutdown()
 except:
     pass
-    
+
 plaidML_devices_count = len(plaidML_devices)
 
 #choosing backend
@@ -306,11 +306,11 @@ if device.backend is None and not force_tf_cpu:
         nvmlInit()
         has_nvml = True
         device.backend = "tensorflow"   #set tensorflow backend in order to use device.*device() functions
-        
+
         gpu_idxs = device.getAllDevicesIdxsList()
         gpu_caps = np.array ( [ device.getDeviceComputeCapability(gpu_idx) for gpu_idx in gpu_idxs ] )
-        
-        if len ( np.ndarray.flatten ( np.argwhere (gpu_caps >= tf_min_req_cap) ) ) == 0:        
+
+        if len ( np.ndarray.flatten ( np.argwhere (gpu_caps >= tf_min_req_cap) ) ) == 0:
             if not force_plaidML:
                 print ("No CUDA devices found with minimum required compute capability: %d.%d. Falling back to OpenCL mode." % (tf_min_req_cap // 10, tf_min_req_cap % 10) )
             device.backend = None
@@ -320,7 +320,7 @@ if device.backend is None and not force_tf_cpu:
     except:
         #if no NVSMI installed exception will occur
         device.backend = None
-        has_nvml = False     
+        has_nvml = False
 
 if force_plaidML or (device.backend is None and not has_nvidia_device):
     #tensorflow backend was failed without has_nvidia_device , or forcing plaidML, trying to use plaidML backend
@@ -333,7 +333,7 @@ if force_plaidML or (device.backend is None and not has_nvidia_device):
 if device.backend is None:
     if force_tf_cpu:
         device.backend = "tensorflow-cpu"
-    elif not has_nvml:        
+    elif not has_nvml:
         if has_nvidia_device:
             #some notebook systems have NVIDIA card without NVSMI in official drivers
             #in that case considering we have system with one capable GPU and let tensorflow to choose best GPU
@@ -348,4 +348,3 @@ if device.backend is None:
         else:
             #has NVSMI, no capable CUDA-devices, also plaidML was failed, then CPU only
             device.backend = "tensorflow-cpu"
-

nnlib/nnlib.py

@@ -541,7 +541,6 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
             result = CAInitializerMPSubprocessor ( [ (i, K.int_shape(conv_weights)) for i, conv_weights in enumerate(conv_weights_list) ], K.floatx(), K.image_data_format() ).run()
             for idx, weights in result:
                 K.set_value ( conv_weights_list[idx], weights )
-            
         nnlib.CAInitializerMP = CAInitializerMP
             
       

nnlib/pynvml.py

@@ -3,7 +3,7 @@
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions are met:
-# 
+#
 #    * Redistributions of source code must retain the above copyright notice,
 #      this list of conditions and the following disclaimer.
 #    * Redistributions in binary form must reproduce the above copyright
@@ -18,11 +18,11 @@
 # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
-# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 
-# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 
+# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
-# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 
-# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF 
+# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
 # THE POSSIBILITY OF SUCH DAMAGE.
 #####
 
@@ -35,7 +35,7 @@ import sys
 import os
 import threading
 import string
-    
+
 ## C Type mappings ##
 ## Enums
 _nvmlEnableState_t = c_uint
@@ -155,9 +155,9 @@ NVML_FAN_FAILED             = 1
 _nvmlLedColor_t = c_uint
 NVML_LED_COLOR_GREEN        = 0
 NVML_LED_COLOR_AMBER        = 1
-     
+
 _nvmlGpuOperationMode_t = c_uint
-NVML_GOM_ALL_ON                 = 0 
+NVML_GOM_ALL_ON                 = 0
 NVML_GOM_COMPUTE                = 1
 NVML_GOM_LOW_DP                 = 2
 
@@ -173,7 +173,7 @@ NVML_RESTRICTED_API_COUNT                                 = 2
 
 _nvmlBridgeChipType_t = c_uint
 NVML_BRIDGE_CHIP_PLX = 0
-NVML_BRIDGE_CHIP_BRO4 = 1      
+NVML_BRIDGE_CHIP_BRO4 = 1
 NVML_MAX_PHYSICAL_BRIDGE = 128
 
 _nvmlValueType_t = c_uint
@@ -317,7 +317,7 @@ def _nvmlGetFunctionPointer(name):
 
     if name in _nvmlGetFunctionPointer_cache:
         return _nvmlGetFunctionPointer_cache[name]
-    
+
     libLoadLock.acquire()
     try:
         # ensure library was loaded
@@ -364,7 +364,7 @@ def nvmlFriendlyObjectToStruct(obj, model):
 class struct_c_nvmlUnit_t(Structure):
     pass # opaque handle
 c_nvmlUnit_t = POINTER(struct_c_nvmlUnit_t)
-    
+
 class _PrintableStructure(Structure):
     """
     Abstract class that produces nicer __str__ output than ctypes.Structure.
@@ -373,7 +373,7 @@ class _PrintableStructure(Structure):
       <class_name object at 0x7fdf82fef9e0>
     this class will print
       class_name(field_name: formatted_value, field_name: formatted_value)
-    
+
     _fmt_ dictionary of <str _field_ name> -> <str format>
     e.g. class that has _field_ 'hex_value', c_uint could be formatted with
       _fmt_ = {"hex_value" : "%08X"}
@@ -397,7 +397,7 @@ class _PrintableStructure(Structure):
                 fmt = self._fmt_["<default>"]
             result.append(("%s: " + fmt) % (key, value))
         return self.__class__.__name__ + "(" + string.join(result, ", ") + ")"
-    
+
 class c_nvmlUnitInfo_t(_PrintableStructure):
     _fields_ = [
         ('name', c_char * 96),
@@ -444,7 +444,7 @@ class nvmlPciInfo_t(_PrintableStructure):
         ('bus', c_uint),
         ('device', c_uint),
         ('pciDeviceId', c_uint),
-        
+
         # Added in 2.285
         ('pciSubSystemId', c_uint),
         ('reserved0', c_uint),
@@ -503,7 +503,7 @@ class c_nvmlBridgeChipHierarchy_t(_PrintableStructure):
     _fields_ = [
         ('bridgeCount', c_uint),
         ('bridgeChipInfo', c_nvmlBridgeChipInfo_t * 128),
-    ]    
+    ]
 
 class c_nvmlEccErrorCounts_t(_PrintableStructure):
     _fields_ = [
@@ -582,7 +582,7 @@ nvmlClocksThrottleReasonAll               = (
                                                nvmlClocksThrottleReasonSwPowerCap |
                                                nvmlClocksThrottleReasonHwSlowdown |
                                                nvmlClocksThrottleReasonUnknown
-                                            ) 
+                                            )
 
 class c_nvmlEventData_t(_PrintableStructure):
     _fields_ = [
@@ -606,31 +606,31 @@ class c_nvmlAccountingStats_t(_PrintableStructure):
 ## C function wrappers ##
 def nvmlInit():
     _LoadNvmlLibrary()
-    
+
     #
     # Initialize the library
     #
     fn = _nvmlGetFunctionPointer("nvmlInit_v2")
     ret = fn()
     _nvmlCheckReturn(ret)
-   
+
     # Atomically update refcount
     global _nvmlLib_refcount
     libLoadLock.acquire()
     _nvmlLib_refcount += 1
     libLoadLock.release()
     return None
-    
+
 def _LoadNvmlLibrary():
     '''
     Load the library if it isn't loaded already
     '''
     global nvmlLib
-    
+
     if (nvmlLib == None):
         # lock to ensure only one caller loads the library
         libLoadLock.acquire()
-        
+
         try:
             # ensure the library still isn't loaded
             if (nvmlLib == None):
@@ -649,7 +649,7 @@ def _LoadNvmlLibrary():
         finally:
             # lock is always freed
             libLoadLock.release()
-            
+
 def nvmlShutdown():
     #
     # Leave the library loaded, but shutdown the interface
@@ -657,7 +657,7 @@ def nvmlShutdown():
     fn = _nvmlGetFunctionPointer("nvmlShutdown")
     ret = fn()
     _nvmlCheckReturn(ret)
-    
+
     # Atomically update refcount
     global _nvmlLib_refcount
     libLoadLock.acquire()
@@ -701,19 +701,19 @@ def nvmlSystemGetHicVersion():
     c_count = c_uint(0)
     hics = None
     fn = _nvmlGetFunctionPointer("nvmlSystemGetHicVersion")
-    
+
     # get the count
     ret = fn(byref(c_count), None)
-    
+
     # this should only fail with insufficient size
     if ((ret != NVML_SUCCESS) and
         (ret != NVML_ERROR_INSUFFICIENT_SIZE)):
         raise NVMLError(ret)
-    
+
     # if there are no hics
     if (c_count.value == 0):
         return []
-    
+
     hic_array = c_nvmlHwbcEntry_t * c_count.value
     hics = hic_array()
     ret = fn(byref(c_count), hics)
@@ -770,7 +770,7 @@ def nvmlUnitGetFanSpeedInfo(unit):
     ret = fn(unit, byref(c_speeds))
     _nvmlCheckReturn(ret)
     return c_speeds
-    
+
 # added to API
 def nvmlUnitGetDeviceCount(unit):
     c_count = c_uint(0)
@@ -822,7 +822,7 @@ def nvmlDeviceGetHandleByUUID(uuid):
     ret = fn(c_uuid, byref(device))
     _nvmlCheckReturn(ret)
     return device
-    
+
 def nvmlDeviceGetHandleByPciBusId(pciBusId):
     c_busId = c_char_p(pciBusId)
     device = c_nvmlDevice_t()
@@ -858,7 +858,7 @@ def nvmlDeviceGetBrand(handle):
     ret = fn(handle, byref(c_type))
     _nvmlCheckReturn(ret)
     return c_type.value
-    
+
 def nvmlDeviceGetSerial(handle):
     c_serial = create_string_buffer(NVML_DEVICE_SERIAL_BUFFER_SIZE)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetSerial")
@@ -892,14 +892,14 @@ def nvmlDeviceGetMinorNumber(handle):
     ret = fn(handle, byref(c_minor_number))
     _nvmlCheckReturn(ret)
     return c_minor_number.value
-    
+
 def nvmlDeviceGetUUID(handle):
     c_uuid = create_string_buffer(NVML_DEVICE_UUID_BUFFER_SIZE)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetUUID")
     ret = fn(handle, c_uuid, c_uint(NVML_DEVICE_UUID_BUFFER_SIZE))
     _nvmlCheckReturn(ret)
     return c_uuid.value
-    
+
 def nvmlDeviceGetInforomVersion(handle, infoRomObject):
     c_version = create_string_buffer(NVML_DEVICE_INFOROM_VERSION_BUFFER_SIZE)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetInforomVersion")
@@ -929,7 +929,7 @@ def nvmlDeviceValidateInforom(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceValidateInforom")
     ret = fn(handle)
     _nvmlCheckReturn(ret)
-    return None 
+    return None
 
 def nvmlDeviceGetDisplayMode(handle):
     c_mode = _nvmlEnableState_t()
@@ -937,29 +937,29 @@ def nvmlDeviceGetDisplayMode(handle):
     ret = fn(handle, byref(c_mode))
     _nvmlCheckReturn(ret)
     return c_mode.value
-    
+
 def nvmlDeviceGetDisplayActive(handle):
     c_mode = _nvmlEnableState_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetDisplayActive")
     ret = fn(handle, byref(c_mode))
     _nvmlCheckReturn(ret)
     return c_mode.value
-    
-    
+
+
 def nvmlDeviceGetPersistenceMode(handle):
     c_state = _nvmlEnableState_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPersistenceMode")
     ret = fn(handle, byref(c_state))
     _nvmlCheckReturn(ret)
     return c_state.value
-    
+
 def nvmlDeviceGetPciInfo(handle):
     c_info = nvmlPciInfo_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPciInfo_v2")
     ret = fn(handle, byref(c_info))
     _nvmlCheckReturn(ret)
     return c_info
-    
+
 def nvmlDeviceGetClockInfo(handle, type):
     c_clock = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetClockInfo")
@@ -997,7 +997,7 @@ def nvmlDeviceGetSupportedMemoryClocks(handle):
     c_count = c_uint(0)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedMemoryClocks")
     ret = fn(handle, byref(c_count), None)
-    
+
     if (ret == NVML_SUCCESS):
         # special case, no clocks
         return []
@@ -1005,11 +1005,11 @@ def nvmlDeviceGetSupportedMemoryClocks(handle):
         # typical case
         clocks_array = c_uint * c_count.value
         c_clocks = clocks_array()
-        
+
         # make the call again
         ret = fn(handle, byref(c_count), c_clocks)
         _nvmlCheckReturn(ret)
-        
+
         procs = []
         for i in range(c_count.value):
             procs.append(c_clocks[i])
@@ -1025,7 +1025,7 @@ def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz):
     c_count = c_uint(0)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetSupportedGraphicsClocks")
     ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), None)
-    
+
     if (ret == NVML_SUCCESS):
         # special case, no clocks
         return []
@@ -1033,11 +1033,11 @@ def nvmlDeviceGetSupportedGraphicsClocks(handle, memoryClockMHz):
         # typical case
         clocks_array = c_uint * c_count.value
         c_clocks = clocks_array()
-        
+
         # make the call again
         ret = fn(handle, c_uint(memoryClockMHz), byref(c_count), c_clocks)
         _nvmlCheckReturn(ret)
-        
+
         procs = []
         for i in range(c_count.value):
             procs.append(c_clocks[i])
@@ -1053,7 +1053,7 @@ def nvmlDeviceGetFanSpeed(handle):
     ret = fn(handle, byref(c_speed))
     _nvmlCheckReturn(ret)
     return c_speed.value
-    
+
 def nvmlDeviceGetTemperature(handle, sensor):
     c_temp = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetTemperature")
@@ -1075,7 +1075,7 @@ def nvmlDeviceGetPowerState(handle):
     ret = fn(handle, byref(c_pstate))
     _nvmlCheckReturn(ret)
     return c_pstate.value
-    
+
 def nvmlDeviceGetPerformanceState(handle):
     c_pstate = _nvmlPstates_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPerformanceState")
@@ -1089,7 +1089,7 @@ def nvmlDeviceGetPowerManagementMode(handle):
     ret = fn(handle, byref(c_pcapMode))
     _nvmlCheckReturn(ret)
     return c_pcapMode.value
-    
+
 def nvmlDeviceGetPowerManagementLimit(handle):
     c_limit = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPowerManagementLimit")
@@ -1113,7 +1113,7 @@ def nvmlDeviceGetPowerManagementDefaultLimit(handle):
     ret = fn(handle, byref(c_limit))
     _nvmlCheckReturn(ret)
     return c_limit.value
-    
+
 
 # Added in 331
 def nvmlDeviceGetEnforcedPowerLimit(handle):
@@ -1146,7 +1146,7 @@ def nvmlDeviceGetCurrentGpuOperationMode(handle):
 # Added in 4.304
 def nvmlDeviceGetPendingGpuOperationMode(handle):
     return nvmlDeviceGetGpuOperationMode(handle)[1]
-    
+
 def nvmlDeviceGetMemoryInfo(handle):
     c_memory = c_nvmlMemory_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetMemoryInfo")
@@ -1160,14 +1160,14 @@ def nvmlDeviceGetBAR1MemoryInfo(handle):
     ret = fn(handle, byref(c_bar1_memory))
     _nvmlCheckReturn(ret)
     return c_bar1_memory
-    
+
 def nvmlDeviceGetComputeMode(handle):
     c_mode = _nvmlComputeMode_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeMode")
     ret = fn(handle, byref(c_mode))
     _nvmlCheckReturn(ret)
     return c_mode.value
-    
+
 def nvmlDeviceGetEccMode(handle):
     c_currState = _nvmlEnableState_t()
     c_pendingState = _nvmlEnableState_t()
@@ -1200,7 +1200,7 @@ def nvmlDeviceGetDetailedEccErrors(handle, errorType, counterType):
 	         _nvmlEccCounterType_t(counterType), byref(c_counts))
     _nvmlCheckReturn(ret)
     return c_counts
-    
+
 # Added in 4.304
 def nvmlDeviceGetMemoryErrorCounter(handle, errorType, counterType, locationType):
     c_count = c_ulonglong()
@@ -1212,7 +1212,7 @@ def nvmlDeviceGetMemoryErrorCounter(handle, errorType, counterType, locationType
              byref(c_count))
     _nvmlCheckReturn(ret)
     return c_count.value
-    
+
 def nvmlDeviceGetUtilizationRates(handle):
     c_util = c_nvmlUtilization_t()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetUtilizationRates")
@@ -1273,7 +1273,7 @@ def nvmlDeviceGetComputeRunningProcesses(handle):
     c_count = c_uint(0)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetComputeRunningProcesses")
     ret = fn(handle, byref(c_count), None)
-    
+
     if (ret == NVML_SUCCESS):
         # special case, no running processes
         return []
@@ -1283,11 +1283,11 @@ def nvmlDeviceGetComputeRunningProcesses(handle):
         c_count.value = c_count.value * 2 + 5
         proc_array = c_nvmlProcessInfo_t * c_count.value
         c_procs = proc_array()
-        
+
         # make the call again
         ret = fn(handle, byref(c_count), c_procs)
         _nvmlCheckReturn(ret)
-        
+
         procs = []
         for i in range(c_count.value):
             # use an alternative struct for this object
@@ -1317,11 +1317,11 @@ def nvmlDeviceGetGraphicsRunningProcesses(handle):
         c_count.value = c_count.value * 2 + 5
         proc_array = c_nvmlProcessInfo_t * c_count.value
         c_procs = proc_array()
-        
+
         # make the call again
         ret = fn(handle, byref(c_count), c_procs)
         _nvmlCheckReturn(ret)
-        
+
         procs = []
         for i in range(c_count.value):
             # use an alternative struct for this object
@@ -1351,19 +1351,19 @@ def nvmlUnitSetLedState(unit, color):
     ret = fn(unit, _nvmlLedColor_t(color))
     _nvmlCheckReturn(ret)
     return None
-    
+
 def nvmlDeviceSetPersistenceMode(handle, mode):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetPersistenceMode")
     ret = fn(handle, _nvmlEnableState_t(mode))
     _nvmlCheckReturn(ret)
     return None
-    
+
 def nvmlDeviceSetComputeMode(handle, mode):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetComputeMode")
     ret = fn(handle, _nvmlComputeMode_t(mode))
     _nvmlCheckReturn(ret)
     return None
-    
+
 def nvmlDeviceSetEccMode(handle, mode):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetEccMode")
     ret = fn(handle, _nvmlEnableState_t(mode))
@@ -1381,15 +1381,15 @@ def nvmlDeviceSetDriverModel(handle, model):
     ret = fn(handle, _nvmlDriverModel_t(model))
     _nvmlCheckReturn(ret)
     return None
-    
-def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled): 
+
+def nvmlDeviceSetAutoBoostedClocksEnabled(handle, enabled):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetAutoBoostedClocksEnabled")
     ret = fn(handle, _nvmlEnableState_t(enabled))
     _nvmlCheckReturn(ret)
     return None
     #Throws NVML_ERROR_NOT_SUPPORTED if hardware doesn't support setting auto boosted clocks
 
-def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags): 
+def nvmlDeviceSetDefaultAutoBoostedClocksEnabled(handle, enabled, flags):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetDefaultAutoBoostedClocksEnabled")
     ret = fn(handle, _nvmlEnableState_t(enabled), c_uint(flags))
     _nvmlCheckReturn(ret)
@@ -1402,7 +1402,7 @@ def nvmlDeviceSetApplicationsClocks(handle, maxMemClockMHz, maxGraphicsClockMHz)
     ret = fn(handle, c_uint(maxMemClockMHz), c_uint(maxGraphicsClockMHz))
     _nvmlCheckReturn(ret)
     return None
-    
+
 # Added in 4.304
 def nvmlDeviceResetApplicationsClocks(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceResetApplicationsClocks")
@@ -1416,7 +1416,7 @@ def nvmlDeviceSetPowerManagementLimit(handle, limit):
     ret = fn(handle, c_uint(limit))
     _nvmlCheckReturn(ret)
     return None
-    
+
 # Added in 4.304
 def nvmlDeviceSetGpuOperationMode(handle, mode):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetGpuOperationMode")
@@ -1534,7 +1534,7 @@ def nvmlDeviceGetAccountingMode(handle):
     ret = fn(handle, byref(c_mode))
     _nvmlCheckReturn(ret)
     return c_mode.value
-    
+
 def nvmlDeviceSetAccountingMode(handle, mode):
     fn = _nvmlGetFunctionPointer("nvmlDeviceSetAccountingMode")
     ret = fn(handle, _nvmlEnableState_t(mode))
@@ -1563,7 +1563,7 @@ def nvmlDeviceGetAccountingPids(handle):
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetAccountingPids")
     ret = fn(handle, byref(count), pids)
     _nvmlCheckReturn(ret)
-    return map(int, pids[0:count.value]) 
+    return map(int, pids[0:count.value])
 
 def nvmlDeviceGetAccountingBufferSize(handle):
     bufferSize = c_uint()
@@ -1576,10 +1576,10 @@ def nvmlDeviceGetRetiredPages(device, sourceFilter):
     c_source = _nvmlPageRetirementCause_t(sourceFilter)
     c_count = c_uint(0)
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetRetiredPages")
-    
+
     # First call will get the size
     ret = fn(device, c_source, byref(c_count), None)
-    
+
     # this should only fail with insufficient size
     if ((ret != NVML_SUCCESS) and
         (ret != NVML_ERROR_INSUFFICIENT_SIZE)):
@@ -1651,7 +1651,7 @@ def nvmlDeviceGetViolationStatus(device, perfPolicyType):
     ret = fn(device, c_perfPolicy_type, byref(c_violTime))
     _nvmlCheckReturn(ret)
     return c_violTime
-    
+
 def nvmlDeviceGetPcieThroughput(device, counter):
     c_util = c_uint()
     fn = _nvmlGetFunctionPointer("nvmlDeviceGetPcieThroughput")
@@ -1704,17 +1704,17 @@ def nvmlDeviceGetTopologyCommonAncestor(device1, device2):
 def nvmlDeviceGetCudaComputeCapability(device):
     c_major = c_int()
     c_minor = c_int()
-    
+
     try:
         fn = _nvmlGetFunctionPointer("nvmlDeviceGetCudaComputeCapability")
     except:
         return 9, 9
-    
+
     # get the count
     ret = fn(device, byref(c_major), byref(c_minor))
-    
+
     # this should only fail with insufficient size
     if (ret != NVML_SUCCESS):
         raise NVMLError(ret)
 
-    return c_major.value, c_minor.value
+    return c_major.value, c_minor.value

samples/Sample.py

@@ -5,17 +5,17 @@ from utils.cv2_utils import *
 
 class SampleType(IntEnum):
     IMAGE = 0 #raw image
-    
+
     FACE_BEGIN = 1
     FACE = 1                      #aligned face unsorted
     FACE_YAW_SORTED = 2           #sorted by yaw
     FACE_YAW_SORTED_AS_TARGET = 3 #sorted by yaw and included only yaws which exist in TARGET also automatic mirrored
     FACE_WITH_CLOSE_TO_SELF = 4
     FACE_END = 4
-    
+
     QTY = 5
-    
-class Sample(object):    
+
+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):
         self.sample_type = sample_type if sample_type is not None else SampleType.IMAGE
         self.filename = filename
@@ -26,19 +26,19 @@ class Sample(object):
         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):
-        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(), 
-            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, 
+        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(),
+            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,
             close_target_list=close_target_list if close_target_list is not None else self.close_target_list)
-    
+
     def load_bgr(self):
         img = cv2_imread (self.filename).astype(np.float32) / 255.0
         if self.mirror:
@@ -48,4 +48,4 @@ class Sample(object):
     def get_random_close_target_sample(self):
         if self.close_target_list is None:
             return None
-        return self.close_target_list[randint (0, len(self.close_target_list)-1)]
+        return self.close_target_list[randint (0, len(self.close_target_list)-1)]

samples/SampleGeneratorBase.py

@@ -4,22 +4,21 @@ from pathlib import Path
 You can implement your own SampleGenerator
 '''
 class SampleGeneratorBase(object):
-    
-    
+
+
     def __init__ (self, samples_path, debug, batch_size):
         if samples_path is None:
             raise Exception('samples_path is None')
-            
+
         self.samples_path = Path(samples_path)
         self.debug = debug
-        self.batch_size = 1 if self.debug else batch_size      
-        
+        self.batch_size = 1 if self.debug else batch_size
+
     #overridable
     def __iter__(self):
         #implement your own iterator
         return self
-        
+
     def __next__(self):
         #implement your own iterator
         return None
-    

samples/SampleGeneratorFace.py

@@ -12,9 +12,9 @@ from samples import SampleLoader
 from samples import SampleGeneratorBase
 
 '''
-arg 
+arg
 output_sample_types = [ 
-                        [SampleProcessor.TypeFlags, size, (optional)random_sub_size] , 
+                        [SampleProcessor.TypeFlags, size, (optional)random_sub_size] ,
                         ...
                       ]
 '''
@@ -26,7 +26,7 @@ class SampleGeneratorFace(SampleGeneratorBase):
         self.add_sample_idx = add_sample_idx
         self.add_pitch = add_pitch
         self.add_yaw = add_yaw
-               
+
         if sort_by_yaw_target_samples_path is not None:
             self.sample_type = SampleType.FACE_YAW_SORTED_AS_TARGET
         elif sort_by_yaw:
@@ -34,9 +34,9 @@ class SampleGeneratorFace(SampleGeneratorBase):
         elif with_close_to_self:
             self.sample_type = SampleType.FACE_WITH_CLOSE_TO_SELF
         else:
-            self.sample_type = SampleType.FACE         
-  
-        self.samples = SampleLoader.load (self.sample_type, self.samples_path, sort_by_yaw_target_samples_path)        
+            self.sample_type = SampleType.FACE
+
+        self.samples = SampleLoader.load (self.sample_type, self.samples_path, sort_by_yaw_target_samples_path)
 
         if self.debug:
             self.generators_count = 1
@@ -46,24 +46,24 @@ class SampleGeneratorFace(SampleGeneratorBase):
             self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, i ) for i in range(self.generators_count) ]
 
         self.generators_sq = [ multiprocessing.Queue() for _ in range(self.generators_count) ]
-            
+
         self.generator_counter = -1
-    
+
     def __iter__(self):
         return self
-        
+
     def __next__(self):
         self.generator_counter += 1
         generator = self.generators[self.generator_counter % len(self.generators) ]
         return next(generator)
-    
+
     #forces to repeat these sample idxs as fast as possible
     #currently unused
     def repeat_sample_idxs(self, idxs): # [ idx, ... ]
         #send idxs list to all sub generators.
         for gen_sq in self.generators_sq:
-            gen_sq.put (idxs) 
-    
+            gen_sq.put (idxs)
+
     def batch_func(self, generator_id):
         gen_sq = self.generators_sq[generator_id]
         samples = self.samples
@@ -73,11 +73,11 @@ class SampleGeneratorFace(SampleGeneratorBase):
 
         if len(samples_idxs) == 0:
             raise ValueError('No training data provided.')
-  
+
         if self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
             if all ( [ samples[idx] == None for idx in samples_idxs] ):
                 raise ValueError('Not enough training data. Gather more faces!')
-             
+
         if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
             shuffle_idxs = []
         elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
@@ -89,25 +89,25 @@ class SampleGeneratorFace(SampleGeneratorBase):
                 idxs = gen_sq.get()
                 for idx in idxs:
                     if idx in samples_idxs:
-                        repeat_samples_idxs.append(idx)                
+                        repeat_samples_idxs.append(idx)
 
             batches = None
             for n_batch in range(self.batch_size):
                 while True:
                     sample = None
-                    
+
                     if len(repeat_samples_idxs) > 0:
-                        idx = repeat_samples_idxs.pop()                        
-                        if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:                            
+                        idx = repeat_samples_idxs.pop()
+                        if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
                             sample = samples[idx]
                         elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
                             sample = samples[(idx >> 16) & 0xFFFF][idx & 0xFFFF]
-                    else:        
+                    else:
                         if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
                             if len(shuffle_idxs) == 0:
                                 shuffle_idxs = samples_idxs.copy()
                                 np.random.shuffle(shuffle_idxs)
-                                
+
                             idx = shuffle_idxs.pop()
                             sample = samples[ idx ]
 
@@ -120,18 +120,18 @@ class SampleGeneratorFace(SampleGeneratorBase):
                             if samples[idx] != None:
                                 if len(shuffle_idxs_2D[idx]) == 0:
                                     shuffle_idxs_2D[idx] = random.sample( range(len(samples[idx])), len(samples[idx]) )
-                                    
-                                idx2 = shuffle_idxs_2D[idx].pop()                            
+
+                                idx2 = shuffle_idxs_2D[idx].pop()
                                 sample = samples[idx][idx2]
-                                
+
                                 idx = (idx << 16) | (idx2 & 0xFFFF)
-                            
-                    if sample is not None:          
+
+                    if sample is not None:
                         try:
                             x = SampleProcessor.process (sample, self.sample_process_options, self.output_sample_types, self.debug)
                         except:
                             raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) )
-                        
+
                         if type(x) != tuple and type(x) != list:
                             raise Exception('SampleProcessor.process returns NOT tuple/list')
 
@@ -144,23 +144,23 @@ class SampleGeneratorFace(SampleGeneratorBase):
                                 batches += [ [] ]
                                 i_pitch = len(batches)-1
                             if self.add_yaw:
-                                batches += [ [] ]    
+                                batches += [ [] ]
                                 i_yaw = len(batches)-1
-                                
+
                         for i in range(len(x)):
                             batches[i].append ( x[i] )
 
                         if self.add_sample_idx:
                             batches[i_sample_idx].append (idx)
-                        
+
                         if self.add_pitch or self.add_yaw:
                             pitch, yaw = LandmarksProcessor.estimate_pitch_yaw (sample.landmarks)
-                            
+
                         if self.add_pitch:
                             batches[i_pitch].append ([pitch])
-                            
+
                         if self.add_yaw:
                             batches[i_yaw].append ([yaw])
-                            
+
                         break
             yield [ np.array(batch) for batch in batches]

samples/SampleGeneratorImageTemporal.py

@@ -11,36 +11,36 @@ from samples import SampleLoader
 from samples import SampleGeneratorBase
 
 '''
-output_sample_types = [ 
-                        [SampleProcessor.TypeFlags, size, (optional)random_sub_size] , 
+output_sample_types = [
+                        [SampleProcessor.TypeFlags, size, (optional)random_sub_size] ,
                         ...
                       ]
 '''
 class SampleGeneratorImageTemporal(SampleGeneratorBase):
     def __init__ (self, samples_path, debug, batch_size, temporal_image_count, sample_process_options=SampleProcessor.Options(), output_sample_types=[], **kwargs):
         super().__init__(samples_path, debug, batch_size)
-        
+
         self.temporal_image_count = temporal_image_count
         self.sample_process_options = sample_process_options
         self.output_sample_types = output_sample_types
 
-        self.samples = SampleLoader.load (SampleType.IMAGE, self.samples_path)        
-        
+        self.samples = SampleLoader.load (SampleType.IMAGE, self.samples_path)
+
         self.generator_samples = [ self.samples ]
         self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, 0 )] if self.debug else \
                           [iter_utils.SubprocessGenerator ( self.batch_func, 0 )]
-                
+
         self.generator_counter = -1
-    
+
     def __iter__(self):
         return self
-        
+
     def __next__(self):
         self.generator_counter += 1
         generator = self.generators[self.generator_counter % len(self.generators) ]
         return next(generator)
-        
-    def batch_func(self, generator_id):    
+
+    def batch_func(self, generator_id):
         samples = self.generator_samples[generator_id]
         samples_len = len(samples)
         if samples_len == 0:
@@ -48,20 +48,20 @@ class SampleGeneratorImageTemporal(SampleGeneratorBase):
 
         if samples_len - self.temporal_image_count < 0:
             raise ValueError('Not enough samples to fit temporal line.')
-            
+
         shuffle_idxs = []
         samples_sub_len = samples_len - self.temporal_image_count + 1
-        
-        while True:                
-            
+
+        while True:
+
             batches = None
             for n_batch in range(self.batch_size):
 
                 if len(shuffle_idxs) == 0:
                     shuffle_idxs = random.sample( range(samples_sub_len), samples_sub_len )
-                
+
                 idx = shuffle_idxs.pop()
-                
+
                 temporal_samples = []
 
                 for i in range( self.temporal_image_count ):
@@ -70,11 +70,11 @@ class SampleGeneratorImageTemporal(SampleGeneratorBase):
                         temporal_samples += SampleProcessor.process (sample, self.sample_process_options, self.output_sample_types, self.debug)
                     except:
                         raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) )
-                    
+
                 if batches is None:
                     batches = [ [] for _ in range(len(temporal_samples)) ]
-                    
+
                 for i in range(len(temporal_samples)):
                     batches[i].append ( temporal_samples[i] )
-       
+
             yield [ np.array(batch) for batch in batches]

samples/SampleLoader.py

@@ -17,44 +17,44 @@ from interact import interact as io
 
 class SampleLoader:
     cache = dict()
-    
+
     @staticmethod
     def load(sample_type, samples_path, target_samples_path=None):
         cache = SampleLoader.cache
-        
+
         if str(samples_path) not in cache.keys():
             cache[str(samples_path)] = [None]*SampleType.QTY
-            
+
         datas = cache[str(samples_path)]
 
         if            sample_type == SampleType.IMAGE:
-            if  datas[sample_type] is None:  
+            if  datas[sample_type] is None:
                 datas[sample_type] = [ Sample(filename=filename) for filename in io.progress_bar_generator( Path_utils.get_image_paths(samples_path), "Loading") ]
 
         elif          sample_type == SampleType.FACE:
-            if  datas[sample_type] is None:  
+            if  datas[sample_type] is None:
                 datas[sample_type] = SampleLoader.upgradeToFaceSamples( [ Sample(filename=filename) for filename in Path_utils.get_image_paths(samples_path) ] )
-        
+
         elif          sample_type == SampleType.FACE_YAW_SORTED:
             if  datas[sample_type] is None:
                 datas[sample_type] = SampleLoader.upgradeToFaceYawSortedSamples( SampleLoader.load(SampleType.FACE, samples_path) )
-        
-        elif          sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:            
+
+        elif          sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
             if  datas[sample_type] is None:
                 if target_samples_path is None:
                     raise Exception('target_samples_path is None for FACE_YAW_SORTED_AS_TARGET')
                 datas[sample_type] = SampleLoader.upgradeToFaceYawSortedAsTargetSamples( SampleLoader.load(SampleType.FACE_YAW_SORTED, samples_path), SampleLoader.load(SampleType.FACE_YAW_SORTED, target_samples_path) )
-        elif          sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:            
+        elif          sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
             if  datas[sample_type] is None:
                 datas[sample_type] = SampleLoader.upgradeToFaceCloseToSelfSamples( SampleLoader.load(SampleType.FACE, samples_path) )
-        
-        
+
+
         return datas[sample_type]
-        
+
     @staticmethod
     def upgradeToFaceSamples ( samples ):
         sample_list = []
-        
+
         for s in io.progress_bar_generator(samples, "Loading"):
             s_filename_path = Path(s.filename)
             try:
@@ -64,57 +64,57 @@ class SampleLoader:
                     dflimg = DFLJPG.load ( str(s_filename_path) )
                 else:
                     dflimg = None
-                    
+
                 if dflimg is None:
-                    print ("%s is not a dfl image file required for training" % (s_filename_path.name) ) 
+                    print ("%s is not a dfl image file required for training" % (s_filename_path.name) )
                     continue
-                    
+
                 pitch, yaw = LandmarksProcessor.estimate_pitch_yaw ( dflimg.get_landmarks() )
 
                 sample_list.append( s.copy_and_set(sample_type=SampleType.FACE,
                                                    face_type=FaceType.fromString (dflimg.get_face_type()),
-                                                   shape=dflimg.get_shape(), 
+                                                   shape=dflimg.get_shape(),
                                                    landmarks=dflimg.get_landmarks(),
                                                    pitch=pitch,
                                                    yaw=yaw) )
             except:
                 print ("Unable to load %s , error: %s" % (str(s_filename_path), traceback.format_exc() ) )
-                
-        return sample_list 
-     
+
+        return sample_list
+
     @staticmethod
     def upgradeToFaceCloseToSelfSamples (samples):
         yaw_samples = SampleLoader.upgradeToFaceYawSortedSamples(samples)
         yaw_samples_len = len(yaw_samples)
-        
+
         sample_list = []
         for i in io.progress_bar_generator( range(yaw_samples_len), "Sorting"):
             if yaw_samples[i] is not None:
                 for s in yaw_samples[i]:
                     s_t = []
-                    
-                    for n in range(2000):                    
+
+                    for n in range(2000):
                         yaw_idx = np.clip ( i-10 +np.random.randint(20), 0, yaw_samples_len-1 )
                         if yaw_samples[yaw_idx] is None:
                             continue
-                            
+
                         yaw_idx_samples_len = len(yaw_samples[yaw_idx])
-                        
+
                         yaw_idx_sample = yaw_samples[yaw_idx][ np.random.randint(yaw_idx_samples_len) ]
                         if s.filename == yaw_idx_sample.filename:
                             continue
-                            
+
                         s_t.append ( yaw_idx_sample )
                         if len(s_t) >= 50:
                             break
-                            
+
                     if len(s_t) == 0:
                         s_t = [s]
-                        
+
                     sample_list.append( s.copy_and_set(close_target_list = s_t) )
 
         return sample_list
-        
+
     @staticmethod
     def upgradeToFaceYawSortedSamples( samples ):
 
@@ -123,50 +123,50 @@ class SampleLoader:
         diff_rot_per_grad = abs(highest_yaw-lowest_yaw) / gradations
 
         yaws_sample_list = [None]*gradations
-        
+
         for i in io.progress_bar_generator(range(gradations), "Sorting"):
             yaw = lowest_yaw + i*diff_rot_per_grad
             next_yaw = lowest_yaw + (i+1)*diff_rot_per_grad
 
             yaw_samples = []
-            for s in samples:                
+            for s in samples:
                 s_yaw = s.yaw
                 if (i == 0            and s_yaw < next_yaw) or \
                    (i  < gradations-1 and s_yaw >= yaw and s_yaw < next_yaw) or \
                    (i == gradations-1 and s_yaw >= yaw):
                     yaw_samples.append ( s.copy_and_set(sample_type=SampleType.FACE_YAW_SORTED) )
-                    
+
             if len(yaw_samples) > 0:
                 yaws_sample_list[i] = yaw_samples
-        
+
         return yaws_sample_list
-        
+
     @staticmethod
     def upgradeToFaceYawSortedAsTargetSamples (s, t):
         l = len(s)
         if l != len(t):
             raise Exception('upgradeToFaceYawSortedAsTargetSamples() s_len != t_len')
         b = l // 2
-        
+
         s_idxs = np.argwhere ( np.array ( [ 1 if x != None else 0  for x in s] ) == 1 )[:,0]
         t_idxs = np.argwhere ( np.array ( [ 1 if x != None else 0  for x in t] ) == 1 )[:,0]
-        
-        new_s = [None]*l    
-        
+
+        new_s = [None]*l
+
         for t_idx in t_idxs:
-            search_idxs = []        
+            search_idxs = []
             for i in range(0,l):
                 search_idxs += [t_idx - i, (l-t_idx-1) - i, t_idx + i, (l-t_idx-1) + i]
 
-            for search_idx in search_idxs:            
+            for search_idx in search_idxs:
                 if search_idx in s_idxs:
                     mirrored = ( t_idx != search_idx and ((t_idx < b and search_idx >= b) or (search_idx < b and t_idx >= b)) )
                     new_s[t_idx] = [ sample.copy_and_set(sample_type=SampleType.FACE_YAW_SORTED_AS_TARGET,
-                                                         mirror=True, 
-                                                         yaw=-sample.yaw, 
+                                                         mirror=True,
+                                                         yaw=-sample.yaw,
                                                          landmarks=LandmarksProcessor.mirror_landmarks (sample.landmarks, sample.shape[1] ))
-                                          for sample in s[search_idx] 
-                                        ] if mirrored else s[search_idx]                
+                                          for sample in s[search_idx]
+                                        ] if mirrored else s[search_idx]
                     break
-                 
-        return new_s
+
+        return new_s

samples/SampleProcessor.py

@@ -13,61 +13,61 @@ class SampleProcessor(object):
         WARPED_TRANSFORMED    = 0x00000004,
         TRANSFORMED           = 0x00000008,
         LANDMARKS_ARRAY       = 0x00000010, #currently unused
-                              
+
         RANDOM_CLOSE          = 0x00000020,
         MORPH_TO_RANDOM_CLOSE = 0x00000040,
-                              
+
         FACE_ALIGN_HALF       = 0x00000100,
         FACE_ALIGN_FULL       = 0x00000200,
         FACE_ALIGN_HEAD       = 0x00000400,
-        FACE_ALIGN_AVATAR     = 0x00000800,    
-                              
+        FACE_ALIGN_AVATAR     = 0x00000800,
+
         FACE_MASK_FULL        = 0x00001000,
         FACE_MASK_EYES        = 0x00002000,
-                              
+
         MODE_BGR              = 0x01000000,  #BGR
         MODE_G                = 0x02000000,  #Grayscale
-        MODE_GGG              = 0x04000000,  #3xGrayscale 
+        MODE_GGG              = 0x04000000,  #3xGrayscale
         MODE_M                = 0x08000000,  #mask only
         MODE_BGR_SHUFFLE      = 0x10000000,  #BGR shuffle
-   
-    class Options(object):     
+
+    class Options(object):
         def __init__(self, random_flip = True, normalize_tanh = False, rotation_range=[-10,10], scale_range=[-0.05, 0.05], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05]):
-            self.random_flip = random_flip        
+            self.random_flip = random_flip
             self.normalize_tanh = normalize_tanh
             self.rotation_range = rotation_range
             self.scale_range = scale_range
             self.tx_range = tx_range
-            self.ty_range = ty_range   
-        
+            self.ty_range = ty_range
+
     @staticmethod
     def process (sample, sample_process_options, output_sample_types, debug):
         sample_bgr = sample.load_bgr()
         h,w,c = sample_bgr.shape
 
-        is_face_sample = sample.landmarks is not None 
-        
+        is_face_sample = sample.landmarks is not None
+
         if debug and is_face_sample:
             LandmarksProcessor.draw_landmarks (sample_bgr, sample.landmarks, (0, 1, 0))
-        
+
         close_sample = sample.close_target_list[ np.random.randint(0, len(sample.close_target_list)) ] if sample.close_target_list is not None else None
         close_sample_bgr = close_sample.load_bgr() if close_sample is not None else None
-        
+
         if debug and close_sample_bgr is not None:
-            LandmarksProcessor.draw_landmarks (close_sample_bgr, close_sample.landmarks, (0, 1, 0))        
-        
+            LandmarksProcessor.draw_landmarks (close_sample_bgr, close_sample.landmarks, (0, 1, 0))
+
         params = image_utils.gen_warp_params(sample_bgr, sample_process_options.random_flip, rotation_range=sample_process_options.rotation_range, scale_range=sample_process_options.scale_range, tx_range=sample_process_options.tx_range, ty_range=sample_process_options.ty_range )
 
         images = [[None]*3 for _ in range(30)]
-        
+
         sample_rnd_seed = np.random.randint(0x80000000)
-            
-        outputs = []       
+
+        outputs = []
         for sample_type in output_sample_types:
             f = sample_type[0]
             size = sample_type[1]
             random_sub_size = 0 if len (sample_type) < 3 else min( sample_type[2] , size)
-            
+
             if f & SampleProcessor.TypeFlags.SOURCE != 0:
                 img_type = 0
             elif f & SampleProcessor.TypeFlags.WARPED != 0:
@@ -77,53 +77,53 @@ class SampleProcessor(object):
             elif f & SampleProcessor.TypeFlags.TRANSFORMED != 0:
                 img_type = 3
             elif f & SampleProcessor.TypeFlags.LANDMARKS_ARRAY != 0:
-                img_type = 4                
+                img_type = 4
             else:
                 raise ValueError ('expected SampleTypeFlags type')
-                
+
             if f & SampleProcessor.TypeFlags.RANDOM_CLOSE != 0:
                 img_type += 10
             elif f & SampleProcessor.TypeFlags.MORPH_TO_RANDOM_CLOSE != 0:
                 img_type += 20
-                
+
             face_mask_type = 0
             if f & SampleProcessor.TypeFlags.FACE_MASK_FULL != 0:
-                face_mask_type = 1               
+                face_mask_type = 1
             elif f & SampleProcessor.TypeFlags.FACE_MASK_EYES != 0:
                 face_mask_type = 2
-                  
+
             target_face_type = -1
             if f & SampleProcessor.TypeFlags.FACE_ALIGN_HALF != 0:
-                target_face_type = FaceType.HALF            
+                target_face_type = FaceType.HALF
             elif f & SampleProcessor.TypeFlags.FACE_ALIGN_FULL != 0:
                 target_face_type = FaceType.FULL
             elif f & SampleProcessor.TypeFlags.FACE_ALIGN_HEAD != 0:
                 target_face_type = FaceType.HEAD
             elif f & SampleProcessor.TypeFlags.FACE_ALIGN_AVATAR != 0:
                 target_face_type = FaceType.AVATAR
-            
+
             if img_type == 4:
-                l = sample.landmarks 
+                l = sample.landmarks
                 l = np.concatenate ( [ np.expand_dims(l[:,0] / w,-1), np.expand_dims(l[:,1] / h,-1) ], -1 )
                 l = np.clip(l, 0.0, 1.0)
                 img = l
-            else:                
+            else:
                 if images[img_type][face_mask_type] is None:
                     if img_type >= 10 and img_type <= 19: #RANDOM_CLOSE
                         img_type -= 10
                         img = close_sample_bgr
                         cur_sample = close_sample
-                        
+
                     elif img_type >= 20 and img_type <= 29: #MORPH_TO_RANDOM_CLOSE
                         img_type -= 20
                         res = sample.shape[0]
-                    
-                        s_landmarks = sample.landmarks.copy()                    
-                        d_landmarks = close_sample.landmarks.copy()                        
-                        idxs = list(range(len(s_landmarks)))                        
+
+                        s_landmarks = sample.landmarks.copy()
+                        d_landmarks = close_sample.landmarks.copy()
+                        idxs = list(range(len(s_landmarks)))
                         #remove landmarks near boundaries
                         for i in idxs[:]:
-                            s_l = s_landmarks[i]   
+                            s_l = s_landmarks[i]
                             d_l = d_landmarks[i]
                             if s_l[0] < 5 or s_l[1] < 5 or s_l[0] >= res-5 or s_l[1] >= res-5 or \
                                d_l[0] < 5 or d_l[1] < 5 or d_l[0] >= res-5 or d_l[1] >= res-5:
@@ -139,39 +139,39 @@ class SampleProcessor(object):
                                     diff_l = np.abs(s_l - s_l_2)
                                     if np.sqrt(diff_l.dot(diff_l)) < 5:
                                         idxs.remove(i)
-                                        break                                    
+                                        break
                         s_landmarks = s_landmarks[idxs]
                         d_landmarks = d_landmarks[idxs]
-                        s_landmarks = np.concatenate ( [s_landmarks, [ [0,0], [ res // 2, 0], [ res-1, 0], [0, res//2], [res-1, res//2] ,[0,res-1] ,[res//2, res-1] ,[res-1,res-1] ] ] ) 
+                        s_landmarks = np.concatenate ( [s_landmarks, [ [0,0], [ res // 2, 0], [ res-1, 0], [0, res//2], [res-1, res//2] ,[0,res-1] ,[res//2, res-1] ,[res-1,res-1] ] ] )
                         d_landmarks = np.concatenate ( [d_landmarks, [ [0,0], [ res // 2, 0], [ res-1, 0], [0, res//2], [res-1, res//2] ,[0,res-1] ,[res//2, res-1] ,[res-1,res-1] ] ] )
                         img = image_utils.morph_by_points (sample_bgr, s_landmarks, d_landmarks)
                         cur_sample = close_sample
                     else:
                         img = sample_bgr
                         cur_sample = sample
-                    
+
                     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) ), -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)
                             mask[mask > 0.0] = 1.0
-                            img = np.concatenate( (img, mask ), -1 )               
+                            img = np.concatenate( (img, mask ), -1 )
 
                     images[img_type][face_mask_type] = image_utils.warp_by_params (params, img, (img_type==1 or img_type==2), (img_type==2 or img_type==3), img_type != 0, face_mask_type == 0)
-                    
+
                 img = images[img_type][face_mask_type]
-                
+
                 if is_face_sample and target_face_type != -1:
                     if target_face_type > sample.face_type:
                         raise Exception ('sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample.filename, sample.face_type, target_face_type) )
                     img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, size, target_face_type), (size,size), flags=cv2.INTER_CUBIC )
                 else:
                     img = cv2.resize( img, (size,size), cv2.INTER_CUBIC )
-                    
+
                 if random_sub_size != 0:
-                    sub_size = size - random_sub_size                
+                    sub_size = size - random_sub_size
                     rnd_state = np.random.RandomState (sample_rnd_seed+random_sub_size)
                     start_x = rnd_state.randint(sub_size+1)
                     start_y = rnd_state.randint(sub_size+1)
@@ -195,7 +195,7 @@ class SampleProcessor(object):
                     img = img_mask
                 else:
                     raise ValueError ('expected SampleTypeFlags mode')
-         
+
                 if not debug:
                     if sample_process_options.normalize_tanh:
                         img = np.clip (img * 2.0 - 1.0, -1.0, 1.0)
@@ -213,6 +213,6 @@ class SampleProcessor(object):
                 elif output.shape[2] == 4:
                     result += [output[...,0:3]*output[...,3:4],]
 
-            return result            
+            return result
         else:
-            return outputs
+            return outputs

samples/__init__.py

@@ -4,4 +4,4 @@ from .SampleLoader import SampleLoader
 from .SampleProcessor import SampleProcessor
 from .SampleGeneratorBase import SampleGeneratorBase
 from .SampleGeneratorFace import SampleGeneratorFace
-from .SampleGeneratorImageTemporal import SampleGeneratorImageTemporal
+from .SampleGeneratorImageTemporal import SampleGeneratorImageTemporal

utils/DFLJPG.py

@@ -11,7 +11,7 @@ class DFLJPG(object):
         self.chunks = []
         self.dfl_dict = None
         self.shape = (0,0,0)
-        
+
     @staticmethod
     def load_raw(filename):
         try:
@@ -19,7 +19,7 @@ class DFLJPG(object):
                 data = f.read()
         except:
             raise FileNotFoundError(data)
-    
+
         try:
             inst = DFLJPG()
             inst.data = data
@@ -30,23 +30,23 @@ class DFLJPG(object):
             while data_counter < inst_length:
                 chunk_m_l, chunk_m_h = struct.unpack ("BB", data[data_counter:data_counter+2])
                 data_counter += 2
-                
+
                 if chunk_m_l != 0xFF:
                     raise ValueError("No Valid JPG info")
-                
+
                 chunk_name = None
                 chunk_size = None
                 chunk_data = None
                 chunk_ex_data = None
                 is_unk_chunk = False
-                
-                if chunk_m_h & 0xF0 == 0xD0:            
+
+                if chunk_m_h & 0xF0 == 0xD0:
                     n = chunk_m_h & 0x0F
-                    
-                    if n >= 0 and n <= 7: 
+
+                    if n >= 0 and n <= 7:
                         chunk_name = "RST%d" % (n)
                         chunk_size = 0
-                    elif n == 0x8: 
+                    elif n == 0x8:
                         chunk_name = "SOI"
                         chunk_size = 0
                         if len(chunks) != 0:
@@ -54,73 +54,73 @@ class DFLJPG(object):
                     elif n == 0x9:
                         chunk_name = "EOI"
                         chunk_size = 0
-                    elif n == 0xA: 
-                        chunk_name = "SOS"                        
-                    elif n == 0xB: 
+                    elif n == 0xA:
+                        chunk_name = "SOS"
+                    elif n == 0xB:
                         chunk_name = "DQT"
                     elif n == 0xD:
                         chunk_name = "DRI"
                         chunk_size = 2
                     else:
                         is_unk_chunk = True
-                elif chunk_m_h & 0xF0 == 0xC0:            
-                    n = chunk_m_h & 0x0F            
-                    if n == 0: 
+                elif chunk_m_h & 0xF0 == 0xC0:
+                    n = chunk_m_h & 0x0F
+                    if n == 0:
                         chunk_name = "SOF0"
-                    elif n == 2: 
+                    elif n == 2:
                         chunk_name = "SOF2"
-                    elif n == 4: 
+                    elif n == 4:
                         chunk_name = "DHT"
                     else:
                         is_unk_chunk = True
-                elif chunk_m_h & 0xF0 == 0xE0:            
+                elif chunk_m_h & 0xF0 == 0xE0:
                     n = chunk_m_h & 0x0F
                     chunk_name = "APP%d" % (n)
                 else:
                     is_unk_chunk = True
-                    
+
                 if is_unk_chunk:
-                    raise ValueError("Unknown chunk %X" % (chunk_m_h) )    
-        
+                    raise ValueError("Unknown chunk %X" % (chunk_m_h) )
+
                 if chunk_size == None: #variable size
                     chunk_size, = struct.unpack (">H", data[data_counter:data_counter+2])
                     chunk_size -= 2
                     data_counter += 2
-                
+
                 if chunk_size > 0:
                     chunk_data = data[data_counter:data_counter+chunk_size]
                     data_counter += chunk_size
-                
+
                 if chunk_name == "SOS":
-                    c = data_counter                        
+                    c = data_counter
                     while c < inst_length and (data[c] != 0xFF or data[c+1] != 0xD9):
                         c += 1
-                        
+
                     chunk_ex_data = data[data_counter:c]
                     data_counter = c
-                
+
                 chunks.append ({'name' : chunk_name,
                                 'm_h' : chunk_m_h,
                                 'data' : chunk_data,
                                 'ex_data' : chunk_ex_data,
-                                })                                
+                                })
             inst.chunks = chunks
-            
+
             return inst
         except Exception as e:
             raise Exception ("Corrupted JPG file: %s" % (str(e)))
- 
+
     @staticmethod
     def load(filename):
         try:
             inst = DFLJPG.load_raw (filename)
             inst.dfl_dict = None
-            
+
             for chunk in inst.chunks:
                 if chunk['name'] == 'APP0':
                     d, c = chunk['data'], 0
                     c, id, _ = struct_unpack (d, c, "=4sB")
-                    
+
                     if id == b"JFIF":
                         c, ver_major, ver_minor, units, Xdensity, Ydensity, Xthumbnail, Ythumbnail = struct_unpack (d, c, "=BBBHHBB")
                         #if units == 0:
@@ -131,22 +131,22 @@ class DFLJPG(object):
                     d, c = chunk['data'], 0
                     c, precision, height, width = struct_unpack (d, c, ">BHH")
                     inst.shape = (height, width, 3)
-                    
+
                 elif chunk['name'] == 'APP15':
                     if type(chunk['data']) == bytes:
                         inst.dfl_dict = pickle.loads(chunk['data'])
 
             if (inst.dfl_dict is not None) and ('face_type' not in inst.dfl_dict.keys()):
                 inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL)
-            
+
             if inst.dfl_dict == None:
                 return None
-            
+
             return inst
         except Exception as e:
             print (e)
             return None
-        
+
     @staticmethod
     def embed_data(filename, face_type=None,
                              landmarks=None,
@@ -155,7 +155,7 @@ class DFLJPG(object):
                              source_landmarks=None,
                              image_to_face_mat=None
                    ):
-    
+
         inst = DFLJPG.load_raw (filename)
         inst.setDFLDictData ({
                                 'face_type': face_type,
@@ -165,41 +165,41 @@ class DFLJPG(object):
                                 'source_landmarks': source_landmarks,
                                 'image_to_face_mat': image_to_face_mat
                              })
-    
+
         try:
             with open(filename, "wb") as f:
                 f.write ( inst.dump() )
         except:
             raise Exception( 'cannot save %s' % (filename) )
-            
+
     def dump(self):
         data = b""
-        
+
         for chunk in self.chunks:
             data += struct.pack ("BB", 0xFF, chunk['m_h'] )
             chunk_data = chunk['data']
             if chunk_data is not None:
                 data += struct.pack (">H", len(chunk_data)+2 )
                 data += chunk_data
-                
+
             chunk_ex_data = chunk['ex_data']
-            if chunk_ex_data is not None:      
+            if chunk_ex_data is not None:
                 data += chunk_ex_data
 
         return data
-        
-    def get_shape(self):        
+
+    def get_shape(self):
         return self.shape
-        
+
     def get_height(self):
         for chunk in self.chunks:
             if type(chunk) == IHDR:
                 return chunk.height
         return 0
-        
+
     def getDFLDictData(self):
         return self.dfl_dict
-                
+
     def setDFLDictData (self, dict_data=None):
         self.dfl_dict = dict_data
 
@@ -211,17 +211,17 @@ class DFLJPG(object):
         last_app_chunk = 0
         for i, chunk in enumerate (self.chunks):
             if chunk['m_h'] & 0xF0 == 0xE0:
-                last_app_chunk = i            
-        
+                last_app_chunk = i
+
         dflchunk = {'name' : 'APP15',
                     'm_h' : 0xEF,
                     'data' : pickle.dumps(dict_data),
                     'ex_data' : None,
                     }
         self.chunks.insert (last_app_chunk+1, dflchunk)
-       
+
     def get_face_type(self): return self.dfl_dict['face_type']
-    def get_landmarks(self): return np.array ( self.dfl_dict['landmarks'] )   
-    def get_source_filename(self): return self.dfl_dict['source_filename']        
-    def get_source_rect(self): return self.dfl_dict['source_rect']        
+    def get_landmarks(self): return np.array ( self.dfl_dict['landmarks'] )
+    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'] )

utils/DFLPNG.py

@@ -110,7 +110,7 @@ class Chunk(object):
 
     def __str__(self):
         return "<Chunk '{name}' length={length} crc={crc:08X}>".format(**self.__dict__)
-        
+
 class IHDR(Chunk):
 	"""IHDR Chunk
 	width, height, bit_depth, color_type, compression_method,
@@ -189,24 +189,24 @@ class IEND(Chunk):
 class DFLChunk(Chunk):
     def __init__(self, dict_data=None):
         super().__init__("fcWp")
-        self.dict_data = dict_data       
+        self.dict_data = dict_data
 
     def setDictData(self, dict_data):
         self.dict_data = dict_data
-        
+
     def getDictData(self):
         return self.dict_data
-        
+
     @classmethod
     def load(cls, data):
         inst = super().load(data)
-        inst.dict_data = pickle.loads( inst.data )        
+        inst.dict_data = pickle.loads( inst.data )
         return inst
-        
+
     def dump(self):
         self.data = pickle.dumps (self.dict_data)
         return super().dump()
-        
+
 chunk_map = {
     b"IHDR": IHDR,
     b"fcWp": DFLChunk,
@@ -219,7 +219,7 @@ class DFLPNG(object):
         self.length = 0
         self.chunks = []
         self.fcwp_dict = None
-        
+
     @staticmethod
     def load_raw(filename):
         try:
@@ -227,11 +227,11 @@ class DFLPNG(object):
                 data = f.read()
         except:
             raise FileNotFoundError(data)
-    
+
         inst = DFLPNG()
         inst.data = data
         inst.length = len(data)
-        
+
         if data[0:8] != PNG_HEADER:
             msg = "No Valid PNG header"
             raise ValueError(msg)
@@ -244,26 +244,26 @@ class DFLPNG(object):
             chunk = chunk_map.get(chunk_name, Chunk).load(data[chunk_start:chunk_end])
             inst.chunks.append(chunk)
             chunk_start = chunk_end
-        
+
         return inst
-        
+
     @staticmethod
     def load(filename):
         try:
             inst = DFLPNG.load_raw (filename)
             inst.fcwp_dict = inst.getDFLDictData()
-            
+
             if (inst.fcwp_dict is not None) and ('face_type' not in inst.fcwp_dict.keys()):
                 inst.fcwp_dict['face_type'] = FaceType.toString (FaceType.FULL)
-            
+
             if inst.fcwp_dict == None:
                 return None
-            
+
             return inst
         except Exception as e:
             print(e)
             return None
-            
+
     @staticmethod
     def embed_data(filename, face_type=None,
                              landmarks=None,
@@ -271,7 +271,7 @@ class DFLPNG(object):
                              source_rect=None,
                              source_landmarks=None
                    ):
-    
+
         inst = DFLPNG.load_raw (filename)
         inst.setDFLDictData ({
                                 'face_type': face_type,
@@ -280,7 +280,7 @@ class DFLPNG(object):
                                 'source_rect': source_rect,
                                 'source_landmarks': source_landmarks
                              })
-    
+
         try:
             with open(filename, "wb") as f:
                 f.write ( inst.dump() )
@@ -292,7 +292,7 @@ class DFLPNG(object):
         for chunk in self.chunks:
             data += chunk.dump()
         return data
-        
+
     def get_shape(self):
         for chunk in self.chunks:
             if type(chunk) == IHDR:
@@ -301,34 +301,34 @@ class DFLPNG(object):
                 h = chunk.height
                 return (h,w,c)
         return (0,0,0)
-        
+
     def get_height(self):
         for chunk in self.chunks:
             if type(chunk) == IHDR:
                 return chunk.height
         return 0
-        
-    def getDFLDictData(self):        
+
+    def getDFLDictData(self):
         for chunk in self.chunks:
             if type(chunk) == DFLChunk:
                 return chunk.getDictData()
         return None
-                
+
     def setDFLDictData (self, dict_data=None):
         for chunk in self.chunks:
             if type(chunk) == DFLChunk:
                 self.chunks.remove(chunk)
                 break
-    
+
         if not dict_data is None:
             chunk = DFLChunk(dict_data)
             self.chunks.insert(-1, chunk)
-       
-    def get_face_type(self): return self.fcwp_dict['face_type']        
+
+    def get_face_type(self): return self.fcwp_dict['face_type']
     def get_landmarks(self): return np.array ( self.fcwp_dict['landmarks'] )
     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 __str__(self):
         return "<PNG length={length} chunks={}>".format(len(self.chunks), **self.__dict__)

utils/Path_utils.py

@@ -5,8 +5,8 @@ image_extensions = [".jpg", ".jpeg", ".png", ".tif", ".tiff"]
 
 def get_image_paths(dir_path, image_extensions=image_extensions):
     dir_path = Path (dir_path)
-        
-    result = []    
+
+    result = []
     if dir_path.exists():
         for x in list(scandir(str(dir_path))):
             if any([x.name.lower().endswith(ext) for ext in image_extensions]):
@@ -14,25 +14,25 @@ def get_image_paths(dir_path, image_extensions=image_extensions):
     return result
 
 def get_image_unique_filestem_paths(dir_path, verbose_print_func=None):
-    result = get_image_paths(dir_path)   
-    result_dup = set()    
-    
+    result = get_image_paths(dir_path)
+    result_dup = set()
+
     for f in result[:]:
         f_stem = Path(f).stem
-        if f_stem in result_dup:            
+        if f_stem in result_dup:
             result.remove(f)
             if verbose_print_func is not None:
-                verbose_print_func ("Duplicate filenames are not allowed, skipping: %s" % Path(f).name )            
-            continue            
+                verbose_print_func ("Duplicate filenames are not allowed, skipping: %s" % Path(f).name )
+            continue
         result_dup.add(f_stem)
-            
+
     return result
-    
+
 def get_all_dir_names_startswith (dir_path, startswith):
     dir_path = Path (dir_path)
     startswith = startswith.lower()
-        
-    result = []    
+
+    result = []
     if dir_path.exists():
         for x in list(scandir(str(dir_path))):
             if x.name.lower().startswith(startswith):
@@ -42,7 +42,7 @@ def get_all_dir_names_startswith (dir_path, startswith):
 def get_first_file_by_stem (dir_path, stem, exts=None):
     dir_path = Path (dir_path)
     stem = stem.lower()
-    
+
     if dir_path.exists():
         for x in list(scandir(str(dir_path))):
             if not x.is_file():
@@ -50,5 +50,5 @@ def get_first_file_by_stem (dir_path, stem, exts=None):
             xp = Path(x.path)
             if xp.stem.lower() == stem and (exts is None or xp.suffix.lower() in exts):
                 return xp
-                
-    return None
+
+    return None

utils/cv2_utils.py

@@ -11,7 +11,7 @@ def cv2_imread(filename, flags=cv2.IMREAD_UNCHANGED):
             return cv2.imdecode(numpyarray, flags)
     except:
         return None
-    
+
 def cv2_imwrite(filename, img, *args):
     ret, buf = cv2.imencode( Path(filename).suffix, img, *args)
     if ret == True:
@@ -19,4 +19,4 @@ def cv2_imwrite(filename, img, *args):
             with open(filename, "wb") as stream:
                 stream.write( buf )
         except:
-            pass
+            pass

utils/image_utils.py

@@ -21,7 +21,7 @@ def reinhard_color_transfer(target, source, clip=False, preserve_paper=False, so
 		OpenCV image in BGR color space (the source image)
 	target: NumPy array
 		OpenCV image in BGR color space (the target image)
-	clip: Should components of L*a*b* image be scaled by np.clip before 
+	clip: Should components of L*a*b* image be scaled by np.clip before
 		converting back to BGR color space?
 		If False then components will be min-max scaled appropriately.
 		Clipping will keep target image brightness truer to the input.
@@ -32,7 +32,7 @@ def reinhard_color_transfer(target, source, clip=False, preserve_paper=False, so
 		aesthetically pleasing results.
 		If False then L*a*b* components will scaled using the reciprocal of
 		the scaling factor proposed in the paper.  This method seems to produce
-		more consistently aesthetically pleasing results 
+		more consistently aesthetically pleasing results
 
 	Returns:
 	-------
@@ -40,13 +40,13 @@ def reinhard_color_transfer(target, source, clip=False, preserve_paper=False, so
 		OpenCV image (w, h, 3) NumPy array (uint8)
 	"""
 
-        
+
 	# convert the images from the RGB to L*ab* color space, being
 	# sure to utilizing the floating point data type (note: OpenCV
 	# expects floats to be 32-bit, so use that instead of 64-bit)
 	source = cv2.cvtColor(source, cv2.COLOR_BGR2LAB).astype(np.float32)
-	target = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype(np.float32)       
-    
+	target = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype(np.float32)
+
 	# compute color statistics for the source and target images
 	src_input = source if source_mask is None else source*source_mask
 	tgt_input = target if target_mask is None else target*target_mask
@@ -86,7 +86,7 @@ def reinhard_color_transfer(target, source, clip=False, preserve_paper=False, so
 	# type
 	transfer = cv2.merge([l, a, b])
 	transfer = cv2.cvtColor(transfer.astype(np.uint8), cv2.COLOR_LAB2BGR)
-	
+
 	# return the color transferred image
 	return transfer
 
@@ -127,7 +127,7 @@ def linear_color_transfer(target_img, source_img, mode='pca', eps=1e-5):
     matched_img[matched_img>1] = 1
     matched_img[matched_img<0] = 0
     return matched_img
-    
+
 def lab_image_stats(image):
     # compute the mean and standard deviation of each channel
     (l, a, b) = cv2.split(image)
@@ -137,7 +137,7 @@ def lab_image_stats(image):
 
     # return the color statistics
     return (lMean, lStd, aMean, aStd, bMean, bStd)
-    
+
 def _scale_array(arr, clip=True):
     if clip:
         return np.clip(arr, 0, 255)
@@ -145,12 +145,12 @@ def _scale_array(arr, clip=True):
     mn = arr.min()
     mx = arr.max()
     scale_range = (max([mn, 0]), min([mx, 255]))
-   
+
     if mn < scale_range[0] or mx > scale_range[1]:
         return (scale_range[1] - scale_range[0]) * (arr - mn) / (mx - mn) + scale_range[0]
 
     return arr
-    
+
 def channel_hist_match(source, template, hist_match_threshold=255, mask=None):
     # Code borrowed from:
     # https://stackoverflow.com/questions/32655686/histogram-matching-of-two-images-in-python-2-x
@@ -179,22 +179,22 @@ def channel_hist_match(source, template, hist_match_threshold=255, mask=None):
     t_quantiles = 255 * t_quantiles / t_quantiles[-1]
     interp_t_values = np.interp(s_quantiles, t_quantiles, t_values)
 
-    return interp_t_values[bin_idx].reshape(oldshape)    
+    return interp_t_values[bin_idx].reshape(oldshape)
 
 def color_hist_match(src_im, tar_im, hist_match_threshold=255):
     h,w,c = src_im.shape
     matched_R = channel_hist_match(src_im[:,:,0], tar_im[:,:,0], hist_match_threshold, None)
     matched_G = channel_hist_match(src_im[:,:,1], tar_im[:,:,1], hist_match_threshold, None)
     matched_B = channel_hist_match(src_im[:,:,2], tar_im[:,:,2], hist_match_threshold, None)
-    
+
     to_stack = (matched_R, matched_G, matched_B)
     for i in range(3, c):
         to_stack += ( src_im[:,:,i],)
-    
-    
+
+
     matched = np.stack(to_stack, axis=-1).astype(src_im.dtype)
     return matched
-    
+
 
 pil_fonts = {}
 def _get_pil_font (font, size):
@@ -204,65 +204,65 @@ def _get_pil_font (font, size):
         if font_str_id not in pil_fonts.keys():
             pil_fonts[font_str_id] = ImageFont.truetype(font + ".ttf", size=size, encoding="unic")
         pil_font = pil_fonts[font_str_id]
-        return pil_font    
+        return pil_font
     except:
         return ImageFont.load_default()
-                
+
 def get_text_image( shape, text, color=(1,1,1), border=0.2, font=None):
-    try:     
+    try:
         size = shape[1]
         pil_font = _get_pil_font( localization.get_default_ttf_font_name() , size)
         text_width, text_height = pil_font.getsize(text)
-        
+
         canvas = Image.new('RGB', shape[0:2], (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:        
+        if shape[2] != 3:
             result = np.concatenate ( (result, np.ones ( (shape[1],) + (shape[0],) + (shape[2]-3,)) ), axis=2 )
 
         return result
-    except:    
+    except:
         return np.zeros ( (shape[1], shape[0], shape[2]), dtype=np.float32 )
-    
+
 def draw_text( image, rect, text, color=(1,1,1), border=0.2, font=None):
     h,w,c = image.shape
- 
+
     l,t,r,b = rect
     l = np.clip (l, 0, w-1)
     r = np.clip (r, 0, w-1)
     t = np.clip (t, 0, h-1)
     b = np.clip (b, 0, h-1)
-    
+
     image[t:b, l:r] += get_text_image (  (r-l,b-t,c) , text, color, border, font )
-                
+
 def draw_text_lines (image, rect, text_lines, color=(1,1,1), border=0.2, font=None):
     text_lines_len = len(text_lines)
     if text_lines_len == 0:
         return
-        
+
     l,t,r,b = rect
     h = b-t
     h_per_line = h // text_lines_len
-    
+
     for i in range(0, text_lines_len):
         draw_text (image, (l, i*h_per_line, r, (i+1)*h_per_line), text_lines[i], color, border, font)
-        
+
 def get_draw_text_lines ( image, rect, text_lines, color=(1,1,1), border=0.2, font=None):
     image = np.zeros ( image.shape, dtype=np.float )
     draw_text_lines ( image, rect, text_lines, color, border, font)
     return image
-        
-  
+
+
 def draw_polygon (image, points, color, thickness = 1):
     points_len = len(points)
     for i in range (0, points_len):
         p0 = tuple( points[i] )
         p1 = tuple( points[ (i+1) % points_len] )
         cv2.line (image, p0, p1, color, thickness=thickness)
-        
+
 def draw_rect(image, rect, color, thickness=1):
     l,t,r,b = rect
     draw_polygon (image, [ (l,t), (r,t), (r,b), (l,b ) ], color, thickness)
@@ -272,40 +272,40 @@ def rectContains(rect, point) :
 
 def applyAffineTransform(src, srcTri, dstTri, size) :
     warpMat = cv2.getAffineTransform( np.float32(srcTri), np.float32(dstTri) )
-    return cv2.warpAffine( src, warpMat, (size[0], size[1]), None, flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101 )    
-    
-def morphTriangle(dst_img, src_img, st, dt) :                                
+    return cv2.warpAffine( src, warpMat, (size[0], size[1]), None, flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101 )
+
+def morphTriangle(dst_img, src_img, st, dt) :
     (h,w,c) = dst_img.shape
     sr = np.array( cv2.boundingRect(np.float32(st)) )
     dr = np.array( cv2.boundingRect(np.float32(dt)) )
     sRect = st - sr[0:2]
     dRect = dt - dr[0:2]
     d_mask = np.zeros((dr[3], dr[2], c), dtype = np.float32)
-    cv2.fillConvexPoly(d_mask, np.int32(dRect), (1.0,)*c, 8, 0);                                    
-    imgRect = src_img[sr[1]:sr[1] + sr[3], sr[0]:sr[0] + sr[2]]                                    
-    size = (dr[2], dr[3])                                    
-    warpImage1 = applyAffineTransform(imgRect, sRect, dRect, size) 
+    cv2.fillConvexPoly(d_mask, np.int32(dRect), (1.0,)*c, 8, 0);
+    imgRect = src_img[sr[1]:sr[1] + sr[3], sr[0]:sr[0] + sr[2]]
+    size = (dr[2], dr[3])
+    warpImage1 = applyAffineTransform(imgRect, sRect, dRect, size)
 
     if c == 1:
         warpImage1 = np.expand_dims( warpImage1, -1 )
 
     dst_img[dr[1]:dr[1]+dr[3], dr[0]:dr[0]+dr[2]] = dst_img[dr[1]:dr[1]+dr[3], dr[0]:dr[0]+dr[2]]*(1-d_mask) + warpImage1 * d_mask
-    
+
 def morph_by_points (image, sp, dp):
     if sp.shape != dp.shape:
         raise ValueError ('morph_by_points() sp.shape != dp.shape')
-    (h,w,c) = image.shape    
+    (h,w,c) = image.shape
 
     result_image = np.zeros(image.shape, dtype = image.dtype)
 
-    for tri in Delaunay(dp).simplices:                                    
+    for tri in Delaunay(dp).simplices:
         morphTriangle(result_image, image, sp[tri], dp[tri])
-       
+
     return result_image
-    
+
 def equalize_and_stack_square (images, axis=1):
     max_c = max ([ 1 if len(image.shape) == 2 else image.shape[2]  for image in images ] )
-    
+
     target_wh = 99999
     for i,image in enumerate(images):
         if len(image.shape) == 2:
@@ -313,113 +313,112 @@ def equalize_and_stack_square (images, axis=1):
             c = 1
         else:
             h,w,c = image.shape
-        
+
         if h < target_wh:
             target_wh = h
-    
+
         if w < target_wh:
             target_wh = w
-            
+
     for i,image in enumerate(images):
         if len(image.shape) == 2:
             h,w = image.shape
             c = 1
         else:
             h,w,c = image.shape
-            
+
         if c < max_c:
             if c == 1:
                 if len(image.shape) == 2:
-                    image = np.expand_dims ( image, -1 )                
+                    image = np.expand_dims ( image, -1 )
                 image = np.concatenate ( (image,)*max_c, -1 )
             elif c == 2: #GA
                 image = np.expand_dims ( image[...,0], -1 )
-                image = np.concatenate ( (image,)*max_c, -1 )                
+                image = np.concatenate ( (image,)*max_c, -1 )
             else:
                 image = np.concatenate ( (image, np.ones((h,w,max_c - c))), -1 )
 
         if h != target_wh or w != target_wh:
             image = cv2.resize ( image, (target_wh, target_wh) )
             h,w,c = image.shape
-                
+
         images[i] = image
-        
+
     return np.concatenate ( images, axis = 1 )
-    
-def bgr2hsv (img):    
+
+def bgr2hsv (img):
     return cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
-    
-def hsv2bgr (img):    
+
+def hsv2bgr (img):
     return cv2.cvtColor(img, cv2.COLOR_HSV2BGR)
-    
-def bgra2hsva (img):    
+
+def bgra2hsva (img):
     return np.concatenate ( (cv2.cvtColor(img[...,0:3], cv2.COLOR_BGR2HSV ), np.expand_dims (img[...,3], -1)), -1 )
 
 def bgra2hsva_list (imgs):
     return [ bgra2hsva(img) for img in imgs ]
-    
+
 def hsva2bgra (img):
     return np.concatenate ( (cv2.cvtColor(img[...,0:3], cv2.COLOR_HSV2BGR ), np.expand_dims (img[...,3], -1)), -1 )
 
 def hsva2bgra_list (imgs):
     return [ hsva2bgra(img) for img in imgs ]
-    
+
 def gen_warp_params (source, flip, rotation_range=[-10,10], scale_range=[-0.5, 0.5], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05]  ):
     h,w,c = source.shape
     if (h != w) or (w != 64 and w != 128 and w != 256 and w != 512 and w != 1024):
         raise ValueError ('TrainingDataGenerator accepts only square power of 2 images.')
-        
+
     rotation = np.random.uniform( rotation_range[0], rotation_range[1] )
     scale = np.random.uniform(1 +scale_range[0], 1 +scale_range[1])
     tx = np.random.uniform( tx_range[0], tx_range[1] )
-    ty = np.random.uniform( ty_range[0], ty_range[1] ) 
- 
+    ty = np.random.uniform( ty_range[0], ty_range[1] )
+
     #random warp by grid
     cell_size = [ w // (2**i) for i in range(1,4) ] [ np.random.randint(3) ]
     cell_count = w // cell_size + 1
-    
+
     grid_points = np.linspace( 0, w, cell_count)
     mapx = np.broadcast_to(grid_points, (cell_count, cell_count)).copy()
     mapy = mapx.T
-    
+
     mapx[1:-1,1:-1] = mapx[1:-1,1:-1] + random_utils.random_normal( size=(cell_count-2, cell_count-2) )*(cell_size*0.24)
     mapy[1:-1,1:-1] = mapy[1:-1,1:-1] + random_utils.random_normal( size=(cell_count-2, cell_count-2) )*(cell_size*0.24)
 
     half_cell_size = cell_size // 2
-    
+
     mapx = cv2.resize(mapx, (w+cell_size,)*2 )[half_cell_size:-half_cell_size-1,half_cell_size:-half_cell_size-1].astype(np.float32)
     mapy = cv2.resize(mapy, (w+cell_size,)*2 )[half_cell_size:-half_cell_size-1,half_cell_size:-half_cell_size-1].astype(np.float32)
-    
+
     #random transform
     random_transform_mat = cv2.getRotationMatrix2D((w // 2, w // 2), rotation, scale)
     random_transform_mat[:, 2] += (tx*w, ty*w)
-    
+
     params = dict()
     params['mapx'] = mapx
     params['mapy'] = mapy
     params['rmat'] = random_transform_mat
-    params['w'] = w        
+    params['w'] = w
     params['flip'] = flip and np.random.randint(10) < 4
-            
+
     return params
-    
+
 def warp_by_params (params, img, warp, transform, flip, is_border_replicate):
     if warp:
         img = cv2.remap(img, params['mapx'], params['mapy'], cv2.INTER_CUBIC )
     if transform:
-        img = cv2.warpAffine( img, params['rmat'], (params['w'], params['w']), borderMode=(cv2.BORDER_REPLICATE if is_border_replicate else cv2.BORDER_CONSTANT), flags=cv2.INTER_CUBIC )            
+        img = cv2.warpAffine( img, params['rmat'], (params['w'], params['w']), borderMode=(cv2.BORDER_REPLICATE if is_border_replicate else cv2.BORDER_CONSTANT), flags=cv2.INTER_CUBIC )
     if flip and params['flip']:
         img = img[:,::-1,:]
     return img
-    
+
 #n_colors = [0..256]
 def reduce_colors (img_bgr, n_colors):
     img_rgb = (cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) * 255.0).astype(np.uint8)
     img_rgb_pil = Image.fromarray(img_rgb)
     img_rgb_pil_p = img_rgb_pil.convert('P', palette=Image.ADAPTIVE, colors=n_colors)
-    
+
     img_rgb_p = img_rgb_pil_p.convert('RGB')
     img_bgr = cv2.cvtColor( np.array(img_rgb_p, dtype=np.float32) / 255.0, cv2.COLOR_RGB2BGR )
-    
+
     return img_bgr
-  

utils/iter_utils.py

@@ -5,7 +5,7 @@ import time
 
 
 class ThisThreadGenerator(object):
-    def __init__(self, generator_func, user_param=None): 
+    def __init__(self, generator_func, user_param=None):
         super().__init__()
         self.generator_func = generator_func
         self.user_param = user_param
@@ -13,30 +13,30 @@ class ThisThreadGenerator(object):
 
     def __iter__(self):
         return self
-        
+
     def __next__(self):
-        if not self.initialized:  
+        if not self.initialized:
             self.initialized = True
             self.generator_func = self.generator_func(self.user_param)
 
         return next(self.generator_func)
 
 class SubprocessGenerator(object):
-    def __init__(self, generator_func, user_param=None, prefetch=2): 
-        super().__init__()        
+    def __init__(self, generator_func, user_param=None, prefetch=2):
+        super().__init__()
         self.prefetch = prefetch
         self.generator_func = generator_func
         self.user_param = user_param
         self.sc_queue = multiprocessing.Queue()
         self.cs_queue = multiprocessing.Queue()
         self.p = None
-    
+
     def process_func(self):
         self.generator_func = self.generator_func(self.user_param)
-        while True:        
+        while True:
             while self.prefetch > -1:
                 try:
-                    gen_data = next (self.generator_func)              
+                    gen_data = next (self.generator_func)
                 except StopIteration:
                     self.cs_queue.put (None)
                     return
@@ -47,17 +47,17 @@ class SubprocessGenerator(object):
 
     def __iter__(self):
         return self
-        
+
     def __next__(self):
         if self.p == None:
             self.p = multiprocessing.Process(target=self.process_func, args=())
             self.p.daemon = True
             self.p.start()
-    
+
         gen_data = self.cs_queue.get()
         if gen_data is None:
             self.p.terminate()
             self.p.join()
             raise StopIteration()
-        self.sc_queue.put (1)        
-        return gen_data
+        self.sc_queue.put (1)
+        return gen_data

utils/os_utils.py

@@ -4,12 +4,12 @@ import sys
 if sys.platform[0:3] == 'win':
     from ctypes import windll
     from ctypes import wintypes
-        
+
 def set_process_lowest_prio():
     try:
         if sys.platform[0:3] == 'win':
             GetCurrentProcess = windll.kernel32.GetCurrentProcess
-            GetCurrentProcess.restype = wintypes.HANDLE            
+            GetCurrentProcess.restype = wintypes.HANDLE
             SetPriorityClass = windll.kernel32.SetPriorityClass
             SetPriorityClass.argtypes = (wintypes.HANDLE, wintypes.DWORD)
             SetPriorityClass ( GetCurrentProcess(), 0x00000040 )
@@ -19,7 +19,7 @@ def set_process_lowest_prio():
             os.nice(20)
     except:
         print("Unable to set lowest process priority")
-    
+
 def set_process_dpi_aware():
     if sys.platform[0:3] == 'win':
-        windll.user32.SetProcessDPIAware(True)
+        windll.user32.SetProcessDPIAware(True)

utils/random_utils.py

@@ -3,12 +3,12 @@ import numpy as np
 def random_normal( size=(1,), trunc_val = 2.5 ):
     len = np.array(size).prod()
     result = np.empty ( (len,) , dtype=np.float32)
-    
+
     for i in range (len):
         while True:
             x = np.random.normal()
             if x >= -trunc_val and x <= trunc_val:
                 break
         result[i] = (x / trunc_val)
-        
-    return result.reshape ( size )
+
+    return result.reshape ( size )

utils/std_utils.py

@@ -11,26 +11,26 @@ class suppress_stdout_stderr(object):
 
         self.old_stdout_fileno = os.dup ( sys.stdout.fileno() )
         self.old_stderr_fileno = os.dup ( sys.stderr.fileno() )
-        
+
         self.old_stdout = sys.stdout
         self.old_stderr = sys.stderr
-        
+
         os.dup2 ( self.outnull_file.fileno(), self.old_stdout_fileno_undup )
         os.dup2 ( self.errnull_file.fileno(), self.old_stderr_fileno_undup )
-        
-        sys.stdout = self.outnull_file        
+
+        sys.stdout = self.outnull_file
         sys.stderr = self.errnull_file
         return self
-        
-    def __exit__(self, *_):        
+
+    def __exit__(self, *_):
         sys.stdout = self.old_stdout
         sys.stderr = self.old_stderr
-        
+
         os.dup2 ( self.old_stdout_fileno, self.old_stdout_fileno_undup )
         os.dup2 ( self.old_stderr_fileno, self.old_stderr_fileno_undup )
 
         os.close ( self.old_stdout_fileno )
         os.close ( self.old_stderr_fileno )
-        
+
         self.outnull_file.close()
-        self.errnull_file.close()
+        self.errnull_file.close()

utils/struct_utils.py

@@ -1,6 +1,5 @@
 import struct
 
-def struct_unpack(data, counter, fmt):    
+def struct_unpack(data, counter, fmt):
     fmt_size = struct.calcsize(fmt)
     return (counter+fmt_size,) + struct.unpack (fmt, data[counter:counter+fmt_size])
-