added 'sort by vggface': sorting by face similarity using VGGFace model.

Requires 4GB+ VRAM and internet connection for the first run.

main.py

@@ -112,7 +112,7 @@ if __name__ == "__main__":
 
     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.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", "vggface", "test"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
     p.set_defaults (func=process_sort)
 
     def process_util(arguments):

mainscripts/Sorter.py

@@ -1,19 +1,26 @@
 import os
-import sys
-import operator
-import numpy as np
-import cv2
-from shutil import copyfile
-from pathlib import Path
-from utils import Path_utils
-from utils.DFLPNG import DFLPNG
-from utils.DFLJPG import DFLJPG
-from utils.cv2_utils import *
-from facelib import LandmarksProcessor
-from joblib import Subprocessor
 import multiprocessing
-from interact import interact as io
+import operator
+import sys
+from pathlib import Path
+from shutil import copyfile
+
+import cv2
+import numpy as np
+from numpy import linalg as npla
+
+import imagelib
+from facelib import LandmarksProcessor
+from functools import cmp_to_key
 from imagelib import estimate_sharpness
+from interact import interact as io
+from joblib import Subprocessor
+from nnlib import VGGFace
+from utils import Path_utils
+from utils.cv2_utils import *
+from utils.DFLJPG import DFLJPG
+from utils.DFLPNG import DFLPNG
+
 
 class BlurEstimatorSubprocessor(Subprocessor):
     class Cli(Subprocessor.Cli):
@@ -790,6 +797,79 @@ def sort_final(input_path, include_by_blur=True):
 
     return final_img_list, trash_img_list
 
+
+def sort_by_vggface(input_path):
+    io.log_info ("Sorting by face similarity using VGGFace model...")
+    
+    model = VGGFace()
+
+    final_img_list = []
+    trash_img_list = []
+
+    image_paths = Path_utils.get_image_paths(input_path)
+    img_list = [ (x,) for x in image_paths ]
+    img_list_len = len(img_list)
+    img_list_range = [*range(img_list_len)]
+
+    feats = [None]*img_list_len    
+    for i in io.progress_bar_generator(img_list_range, "Loading"):
+        img = cv2_imread( img_list[i][0] ).astype(np.float32)
+        img = imagelib.normalize_channels (img, 3)
+        img = cv2.resize (img, (224,224) )
+        img = img[..., ::-1]
+        img[..., 0] -= 93.5940
+        img[..., 1] -= 104.7624
+        img[..., 2] -= 129.1863
+        feats[i] = model.predict( img[None,...] )[0]
+
+    tmp = np.zeros( (img_list_len,) )
+    float_inf = float("inf")    
+    for i in io.progress_bar_generator ( range(img_list_len-1), "Sorting" ):  
+        i_feat = feats[i]
+        
+        for j in img_list_range:
+            tmp[j] = npla.norm(i_feat-feats[j]) if j >= i+1 else float_inf
+            
+        idx = np.argmin(tmp)
+        
+        img_list[i+1], img_list[idx] = img_list[idx], img_list[i+1]
+        feats[i+1], feats[idx] = feats[idx], feats[i+1]
+
+    return img_list, trash_img_list
+    
+"""
+    img_list_len = len(img_list)
+    
+    for i in io.progress_bar_generator ( range(img_list_len-1), "Sorting" ):        
+        a = []
+        i_1 = img_list[i][1]
+        
+        
+        for j in range(i+1, img_list_len):
+            a.append ( [ j, np.linalg.norm(i_1-img_list[j][1]) ] )
+        
+        x = sorted(a, key=operator.itemgetter(1) )[0][0]
+        saved = img_list[i+1]
+        img_list[i+1] = img_list[x]
+        img_list[x] = saved
+        
+        
+    q = np.array ( [ x[1] for x in img_list ] )
+    
+    for i in io.progress_bar_generator ( range(img_list_len-1), "Sorting" ):        
+        
+        a = np.linalg.norm( q[i] - q[i+1:], axis=1 )
+        a = i+1+np.argmin(a)        
+        
+        saved = img_list[i+1]
+        img_list[i+1] = img_list[a]
+        img_list[a] = saved
+        
+        saved = q[i+1]
+        q[i+1] = q[a]
+        q[a] = saved
+"""
+
 def final_process(input_path, img_list, trash_img_list):
     if len(trash_img_list) != 0:
         parent_input_path = input_path.parent
@@ -851,6 +931,7 @@ def main (input_path, sort_by_method):
     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 == 'vggface':       img_list, trash_img_list = sort_by_vggface (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)
 

mainscripts/dev_misc.py

@@ -38,6 +38,10 @@ def extract_vggface2_dataset(input_dir, device_args={} ):
         cur_input_path = input_path / dir_name
         cur_output_path = output_path / dir_name
         
+        l = len(Path_utils.get_image_paths(cur_input_path))
+        if l < 250 or l > 350:
+            continue
+
         io.log_info (f"Processing: {str(cur_input_path)} ")
         
         if not cur_output_path.exists():

models/Model_DEV_FUNIT/Model.py

@@ -42,7 +42,7 @@ class FUNITModel(ModelBase):
     #override
     def onInitialize(self, batch_size=-1, **in_options):
         exec(nnlib.code_import_all, locals(), globals())
-        self.set_vram_batch_requirements({4:16})
+        self.set_vram_batch_requirements({4:16,11:24})
 
         resolution = self.options['resolution']
         face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
@@ -75,6 +75,7 @@ class FUNITModel(ModelBase):
             face_type = t.FACE_TYPE_FULL if self.options['face_type'] == 'f' else t.FACE_TYPE_HALF
             
             output_sample_types=[ {'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':128, 'normalize_tanh':True} ]
+            output_sample_types1=[ {'types': (t.IMG_SOURCE, face_type, t.MODE_BGR), 'resolution':128, 'normalize_tanh':True} ]
             
             self.set_training_data_generators ([
                         SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
@@ -87,11 +88,11 @@ class FUNITModel(ModelBase):
 
                         SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
                             sample_process_options=SampleProcessor.Options(random_flip=True),
-                            output_sample_types=output_sample_types, person_id_mode=True ),
+                            output_sample_types=output_sample_types1, person_id_mode=True ),
 
                         SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
                             sample_process_options=SampleProcessor.Options(random_flip=True),
-                            output_sample_types=output_sample_types, person_id_mode=True ),
+                            output_sample_types=output_sample_types1, person_id_mode=True ),
                     ])
 
     #override

nnlib/FUNIT.py

@@ -162,10 +162,6 @@ class FUNIT(object):
             for w in weights_list:
                 K.set_value( w, K.get_value(initer(K.int_shape(w)))  )
 
-        #if not self.is_first_run():
-        #    self.load_weights_safe(self.get_model_filename_list())
-
-
 
         if load_weights_locally:
             pass
@@ -188,9 +184,6 @@ class FUNIT(object):
                 [self.D_opt,           'D_opt.h5'],
                 ]
 
-    #def save_weights(self):
-    #    self.model.save_weights (str(self.weights_path))
-
     def train(self, xa,la,xb,lb):
         D_loss, = self.D_train ([xa,la,xb,lb])
         G_loss, = self.G_train ([xa,la,xb,lb])
@@ -209,17 +202,17 @@ class FUNIT(object):
         def ResBlock(dim):
             def func(input):
                 x = input
-                x = Conv2D(dim, 3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
+                x = Conv2D(dim, 3, strides=1, padding='same')(x)
                 x = InstanceNormalization()(x)
                 x = ReLU()(x)
-                x = Conv2D(dim, 3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
+                x = Conv2D(dim, 3, strides=1, padding='same')(x)
                 x = InstanceNormalization()(x)
 
                 return Add()([x,input])
             return func
 
         def func(x):
-            x = Conv2D (nf, kernel_size=7, strides=1, padding='valid')(ZeroPadding2D(3)(x))
+            x = Conv2D (nf, kernel_size=7, strides=1, padding='same')(x)
             x = InstanceNormalization()(x)
             x = ReLU()(x)
             for i in range(downs):
@@ -237,11 +230,11 @@ class FUNIT(object):
         exec (nnlib.import_all(), locals(), globals())
 
         def func(x):
-            x = Conv2D (nf, kernel_size=7, strides=1, padding='valid', activation='relu')(ZeroPadding2D(3)(x))
+            x = Conv2D (nf, kernel_size=7, strides=1, padding='same', activation='relu')(x)
             for i in range(downs):
                 x = Conv2D (nf * min ( 4, 2**(i+1) ), kernel_size=4, strides=2, padding='valid', activation='relu')(ZeroPadding2D(1)(x))
             x = GlobalAveragePooling2D()(x)
-            x = Dense(nf)(x)
+            x = Dense(latent_dim)(x)
             return x
 
         return func
@@ -250,16 +243,14 @@ class FUNIT(object):
     def DecoderFlow(ups, n_res_blks=2, mlp_blks=2, subpixel_decoder=False ):
         exec (nnlib.import_all(), locals(), globals())
 
-
-
         def ResBlock(dim):
             def func(input):
                 inp, mlp = input
                 x = inp
-                x = Conv2D(dim, 3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
+                x = Conv2D(dim, 3, strides=1, padding='same')(x)
                 x = FUNITAdain(kernel_initializer='he_normal')([x,mlp])
                 x = ReLU()(x)
-                x = Conv2D(dim, 3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
+                x = Conv2D(dim, 3, strides=1, padding='same')(x)
                 x = FUNITAdain(kernel_initializer='he_normal')([x,mlp])
                 return Add()([x,inp])
             return func
@@ -280,16 +271,16 @@ class FUNIT(object):
             for i in range(ups):
 
                 if subpixel_decoder:
-                    x = Conv2D (4* (nf // 2**(i+1)), kernel_size=3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
+                    x = Conv2D (4* (nf // 2**(i+1)), kernel_size=3, strides=1, padding='same')(x)
                     x = SubpixelUpscaler()(x)
                 else:
                     x = UpSampling2D()(x)
-                    x = Conv2D (nf // 2**(i+1), kernel_size=5, strides=1, padding='valid')(ZeroPadding2D(2)(x))
+                    x = Conv2D (nf // 2**(i+1), kernel_size=5, strides=1, padding='same')(x)
 
                 x = InstanceNormalization()(x)
                 x = ReLU()(x)
 
-            rgb = Conv2D (3, kernel_size=7, strides=1, padding='valid', activation='tanh')(ZeroPadding2D(3)(x))
+            rgb = Conv2D (3, kernel_size=7, strides=1, padding='same', activation='tanh')(x)
             return rgb
 
         return func

nnlib/VGGFace.py

@@ -0,0 +1,64 @@
+from nnlib import nnlib
+
+def VGGFace():
+    exec(nnlib.import_all(), locals(), globals())
+    
+    img_input = Input(shape=(224,224,3) )
+
+    # Block 1
+    x = Conv2D(64, (3, 3), activation='relu', padding='same', name='conv1_1')(
+        img_input)
+    x = Conv2D(64, (3, 3), activation='relu', padding='same', name='conv1_2')(x)
+    x = MaxPooling2D((2, 2), strides=(2, 2), name='pool1')(x)
+
+    # Block 2
+    x = Conv2D(128, (3, 3), activation='relu', padding='same', name='conv2_1')(
+        x)
+    x = Conv2D(128, (3, 3), activation='relu', padding='same', name='conv2_2')(
+        x)
+    x = MaxPooling2D((2, 2), strides=(2, 2), name='pool2')(x)
+
+    # Block 3
+    x = Conv2D(256, (3, 3), activation='relu', padding='same', name='conv3_1')(
+        x)
+    x = Conv2D(256, (3, 3), activation='relu', padding='same', name='conv3_2')(
+        x)
+    x = Conv2D(256, (3, 3), activation='relu', padding='same', name='conv3_3')(
+        x)
+    x = MaxPooling2D((2, 2), strides=(2, 2), name='pool3')(x)
+
+    # Block 4
+    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='conv4_1')(
+        x)
+    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='conv4_2')(
+        x)
+    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='conv4_3')(
+        x)
+    x = MaxPooling2D((2, 2), strides=(2, 2), name='pool4')(x)
+
+    # Block 5
+    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='conv5_1')(
+        x)
+    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='conv5_2')(
+        x)
+    x = Conv2D(512, (3, 3), activation='relu', padding='same', name='conv5_3')(
+        x)
+    x = MaxPooling2D((2, 2), strides=(2, 2), name='pool5')(x)
+
+
+    # Classification block
+    x = Flatten(name='flatten')(x)
+    x = Dense(4096, name='fc6')(x)
+    x = Activation('relu', name='fc6/relu')(x)
+    x = Dense(4096, name='fc7')(x)
+    x = Activation('relu', name='fc7/relu')(x)
+    x = Dense(2622, name='fc8')(x)
+    x = Activation('softmax', name='fc8/softmax')(x)
+
+    model = Model(img_input, x, name='vggface_vgg16')
+    weights_path = keras.utils.data_utils.get_file('rcmalli_vggface_tf_vgg16.h5',
+                                                   'https://github.com/rcmalli/keras-vggface/releases/download/v2.0/rcmalli_vggface_tf_vgg16.h5')
+
+    model.load_weights(weights_path, by_name=True)
+
+    return model

nnlib/__init__.py

@@ -1,3 +1,4 @@
 from .nnlib import nnlib
 from .FUNIT import FUNIT
 from .TernausNet import TernausNet
+from .VGGFace import VGGFace

nnlib/nnlib.py

@@ -63,6 +63,7 @@ UpSampling2D = KL.UpSampling2D
 BatchNormalization = KL.BatchNormalization
 PixelNormalization = nnlib.PixelNormalization
 
+Activation = KL.Activation
 LeakyReLU = KL.LeakyReLU
 ELU = KL.ELU
 ReLU = KL.ReLU