1

mainscripts/Extractor.py

@@ -27,13 +27,14 @@ DEBUG = False
 
 class ExtractSubprocessor(Subprocessor):
     class Data(object):
-        def __init__(self, filename=None, rects=None, landmarks = None, landmarks_accurate=True, pitch_yaw_roll=None, final_output_files = None):
+        def __init__(self, filename=None, rects=None, landmarks = None, landmarks_accurate=True, pitch_yaw_roll=None, force_output_path=None, final_output_files = None):
             self.filename = filename
             self.rects = rects or []
             self.rects_rotation = 0
             self.landmarks_accurate = landmarks_accurate
             self.landmarks = landmarks or []
             self.pitch_yaw_roll = pitch_yaw_roll
+            self.force_output_path = force_output_path
             self.final_output_files = final_output_files or []
             self.faces_detected = 0
 
@@ -248,13 +249,18 @@ class ExtractSubprocessor(Subprocessor):
                             if self.debug_dir is not None:
                                 LandmarksProcessor.draw_rect_landmarks (debug_image, rect, image_landmarks, self.image_size, self.face_type, transparent_mask=True)
 
+                        final_output_path = self.final_output_path                        
+                        if data.force_output_path is not None:
+                            final_output_path = data.force_output_path
+                        
                         if src_dflimg is not None and filename_path.suffix == '.jpg':
                             #if extracting from dflimg and jpg copy it in order not to lose quality
-                            output_file = str(self.final_output_path / filename_path.name)
+                            output_file = str(final_output_path / filename_path.name)
                             if str(filename_path) != str(output_file):
                                 shutil.copy ( str(filename_path), str(output_file) )
                         else:
-                            output_file = '{}_{}{}'.format(str(self.final_output_path / filename_path.stem), str(face_idx), '.jpg')
+                            
+                            output_file = '{}_{}{}'.format(str(final_output_path / filename_path.stem), str(face_idx), '.jpg')
                             cv2_imwrite(output_file, face_image, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )
 
                         DFLJPG.embed_data(output_file, face_type=FaceType.toString(self.face_type),
@@ -303,7 +309,8 @@ class ExtractSubprocessor(Subprocessor):
         self.devices = ExtractSubprocessor.get_devices_for_config(self.manual, self.type, multi_gpu, cpu_only)
 
         no_response_time_sec = 60 if not self.manual and not DEBUG else 999999
-        super().__init__('Extractor', ExtractSubprocessor.Cli, no_response_time_sec)
+        
+        super().__init__('Extractor', ExtractSubprocessor.Cli, no_response_time_sec, initialize_subprocesses_in_serial=(type != 'final'))
 
     #override
     def on_check_run(self):

mainscripts/dev_misc.py

@@ -5,7 +5,7 @@ from pathlib import Path
 import cv2
 import numpy as np
 
-from facelib import LandmarksProcessor
+from facelib import FaceType, LandmarksProcessor
 from interact import interact as io
 from joblib import Subprocessor
 from utils import Path_utils
@@ -14,8 +14,159 @@ from utils.DFLJPG import DFLJPG
 from utils.DFLPNG import DFLPNG
 
 from . import Extractor, Sorter
+from .Extractor import ExtractSubprocessor
 
 
+def extract_vggface2_dataset(input_dir, device_args={} ):
+    multi_gpu = device_args.get('multi_gpu', False)
+    cpu_only = device_args.get('cpu_only', False)
+
+    input_path = Path(input_dir)
+    if not input_path.exists():
+        raise ValueError('Input directory not found. Please ensure it exists.')
+    
+    bb_csv = input_path / 'loose_bb_train.csv'
+    if not bb_csv.exists():
+        raise ValueError('loose_bb_train.csv found. Please ensure it exists.')
+    
+    bb_lines = bb_csv.read_text().split('\n')
+    bb_lines.pop(0)
+    
+    bb_dict = {}
+    for line in bb_lines:
+        name, l, t, w, h = line.split(',')
+        name = name[1:-1]
+        l, t, w, h = [ int(x) for x in (l, t, w, h) ]        
+        bb_dict[name] = (l,t,w, h)
+    
+
+    output_path = input_path.parent / (input_path.name + '_out')
+    
+    dir_names = Path_utils.get_all_dir_names(input_path)
+    
+    if not output_path.exists():
+        output_path.mkdir(parents=True, exist_ok=True)
+
+    data = []
+    for dir_name in io.progress_bar_generator(dir_names, "Collecting"):
+        cur_input_path = input_path / dir_name
+        cur_output_path = output_path / dir_name
+        
+        if not cur_output_path.exists():
+            cur_output_path.mkdir(parents=True, exist_ok=True)
+            
+        input_path_image_paths = Path_utils.get_image_paths(cur_input_path)
+
+        for filename in input_path_image_paths:
+            filename_path = Path(filename)
+            
+            name = filename_path.parent.name + '/' + filename_path.stem
+            if name not in bb_dict:
+                continue
+
+            l,t,w,h = bb_dict[name]
+            if min(w,h) < 128:
+                continue
+            
+            data += [ ExtractSubprocessor.Data(filename=filename,rects=[ (l,t,l+w,t+h) ], landmarks_accurate=False, force_output_path=cur_output_path ) ]
+            
+    face_type = FaceType.fromString('full_face')
+    
+    io.log_info ('Performing 2nd pass...')
+    data = ExtractSubprocessor (data, 'landmarks', 256, face_type, debug_dir=None, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False).run()
+       
+    io.log_info ('Performing 3rd pass...')
+    ExtractSubprocessor (data, 'final', 256, face_type, debug_dir=None, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, final_output_path=None).run()
+    
+    
+"""
+    import code
+    code.interact(local=dict(globals(), **locals()))
+    
+    data_len = len(data) 
+    i = 0
+    while i < data_len-1:
+        i_name = Path(data[i].filename).parent.name
+        
+        sub_data = []
+        
+        for j in range (i, data_len):
+            j_name = Path(data[j].filename).parent.name
+            if i_name == j_name:
+                sub_data += [ data[j] ]
+            else:
+                break
+        i = j
+        
+        cur_output_path = output_path / i_name        
+        
+        io.log_info (f"Processing: {str(cur_output_path)}, {i}/{data_len} ")
+        
+        if not cur_output_path.exists():
+            cur_output_path.mkdir(parents=True, exist_ok=True)
+            
+            
+        
+
+
+
+    
+    
+    for dir_name in dir_names:
+        
+        cur_input_path = input_path / dir_name
+        cur_output_path = output_path / dir_name
+        
+        input_path_image_paths = Path_utils.get_image_paths(cur_input_path)
+        l = len(input_path_image_paths)
+        #if l < 250 or l > 350:
+        #    continue
+
+        io.log_info (f"Processing: {str(cur_input_path)} ")
+        
+        if not cur_output_path.exists():
+            cur_output_path.mkdir(parents=True, exist_ok=True)
+
+
+        data = []
+        for filename in input_path_image_paths:
+            filename_path = Path(filename)
+            
+            name = filename_path.parent.name + '/' + filename_path.stem
+            if name not in bb_dict:
+                continue
+            
+            bb = bb_dict[name]
+            l,t,w,h = bb
+            if min(w,h) < 128:
+                continue
+            
+            data += [ ExtractSubprocessor.Data(filename=filename,rects=[ (l,t,l+w,t+h) ], landmarks_accurate=False ) ]
+                
+    
+              
+        io.log_info ('Performing 2nd pass...')
+        data = ExtractSubprocessor (data, 'landmarks', 256, face_type, debug_dir=None, multi_gpu=False, cpu_only=False, manual=False).run()
+
+        io.log_info ('Performing 3rd pass...')
+        data = ExtractSubprocessor (data, 'final', 256, face_type, debug_dir=None, multi_gpu=False, cpu_only=False, manual=False, final_output_path=cur_output_path).run()
+        
+        
+        io.log_info (f"Sorting: {str(cur_output_path)} ")
+        Sorter.main (input_path=str(cur_output_path), sort_by_method='hist')
+        
+        import code
+        code.interact(local=dict(globals(), **locals()))
+        
+        #try:
+        #    io.log_info (f"Removing: {str(cur_input_path)} ")
+        #    shutil.rmtree(cur_input_path)
+        #except:
+        #    io.log_info (f"unable to remove: {str(cur_input_path)} ")
+            
+            
+            
+
 def extract_vggface2_dataset(input_dir, device_args={} ):
     multi_gpu = device_args.get('multi_gpu', False)
     cpu_only = device_args.get('cpu_only', False)
@@ -64,7 +215,7 @@ def extract_vggface2_dataset(input_dir, device_args={} ):
         except:
             io.log_info (f"unable to remove: {str(cur_input_path)} ")
             
-            
+"""            
 
 class CelebAMASKHQSubprocessor(Subprocessor):
     class Cli(Subprocessor.Cli):