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

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

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):
@@ -772,24 +779,97 @@ def sort_final(input_path, include_by_blur=True):
             for pg in range(pitch_grads):
                 img_list = pitch_sample_list[pg]
                 if img_list is None:
-                    continue                
+                    continue
                 final_img_list += [ img_list.pop(0) ]
                 if len(img_list) == 0:
-                    pitch_sample_list[pg] = None                
+                    pitch_sample_list[pg] = None
                 n -= 1
                 if n == 0:
                     break
-            if n_prev == n:                
-                break            
+            if n_prev == n:
+                break
 
         for pg in range(pitch_grads):
             img_list = pitch_sample_list[pg]
             if img_list is None:
-                continue                
+                continue
             trash_img_list += img_list
 
     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)