fix scale factor in sample generation. added util --export-faceset-mask

Author: iperov

core/imagelib/warp.py

@@ -115,7 +115,7 @@ def gen_warp_params (w, flip=False, rotation_range=[-10,10], scale_range=[-0.5,
         w = 64
 
     rotation = rnd_state.uniform( rotation_range[0], rotation_range[1] )
-    scale = rnd_state.uniform(1 +scale_range[0], 1 +scale_range[1])
+    scale = rnd_state.uniform( 1/(1-scale_range[0]) , 1+scale_range[1] )
     tx = rnd_state.uniform( tx_range[0], tx_range[1] )
     ty = rnd_state.uniform( ty_range[0], ty_range[1] )
     p_flip = flip and rnd_state.randint(10) < 4

main.py

@@ -98,6 +98,10 @@ def process_util(arguments):
             io.log_info ("Performing faceset unpacking...\r\n")
             from samplelib import PackedFaceset
             PackedFaceset.unpack( Path(arguments.input_dir) )
+            
+        if arguments.export_faceset_mask:
+            io.log_info ("Exporting faceset mask..\r\n")
+            Util.export_faceset_mask( Path(arguments.input_dir) )
 
     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.")
@@ -107,6 +111,7 @@ def process_util(arguments):
     p.add_argument('--restore-faceset-metadata', action="store_true", dest="restore_faceset_metadata", default=False, help="Restore faceset metadata to file. Image filenames must be the same as used with save.")
     p.add_argument('--pack-faceset', action="store_true", dest="pack_faceset", default=False, help="")
     p.add_argument('--unpack-faceset', action="store_true", dest="unpack_faceset", default=False, help="")
+    p.add_argument('--export-faceset-mask', action="store_true", dest="export_faceset_mask", default=False, help="")
 
     p.set_defaults (func=process_util)
 
@@ -271,11 +276,11 @@ def process_faceset_resizer(arguments):
         from mainscripts import FacesetResizer
         FacesetResizer.process_folder ( Path(arguments.input_dir) )
     p.set_defaults(func=process_faceset_resizer)
-    
+
     def process_dev_test(arguments):
         osex.set_process_lowest_prio()
         from mainscripts import dev_misc
-        dev_misc.dev_test( arguments.input_dir )
+        dev_misc.dev_gen_mask_files( arguments.input_dir )
 
     p = subparsers.add_parser( "dev_test", help="")
     p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")

mainscripts/Util.py

@@ -159,3 +159,32 @@ def recover_original_aligned_filename(input_path):
             fs.rename (fd)
         except:
             io.log_err ('fail to rename %s' % (fs.name) )
+
+def export_faceset_mask(input_dir):
+    for filename in io.progress_bar_generator(pathex.get_image_paths (input_dir), "Processing"):
+        filepath = Path(filename)
+        
+        if '_mask' in filepath.stem:
+            continue
+        
+        mask_filepath = filepath.parent / (filepath.stem+'_mask'+filepath.suffix)
+        
+        dflimg = DFLJPG.load(filepath)
+        
+        H,W,C = dflimg.shape
+        
+        seg_ie_polys = dflimg.get_seg_ie_polys()
+        
+        if seg_ie_polys.has_polys():
+            mask = np.zeros ((H,W,1), dtype=np.float32)
+            seg_ie_polys.overlay_mask(mask)
+        elif dflimg.has_xseg_mask():
+            mask = dflimg.get_xseg_mask()
+            mask[mask < 0.5] = 0.0
+            mask[mask >= 0.5] = 1.0
+        else:
+            raise Exception(f'no mask in file {filepath}')
+        
+        
+        cv2_imwrite(mask_filepath, (mask*255).astype(np.uint8), [int(cv2.IMWRITE_JPEG_QUALITY), 100] )
+