diff --git a/core/imagelib/warp.py b/core/imagelib/warp.py index af77579..2c429d0 100644 --- a/core/imagelib/warp.py +++ b/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 diff --git a/main.py b/main.py index 493c7ed..b821910 100644 --- a/main.py +++ b/main.py @@ -98,6 +98,10 @@ if __name__ == "__main__": 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 @@ if __name__ == "__main__": 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 @@ if __name__ == "__main__": 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") diff --git a/mainscripts/Util.py b/mainscripts/Util.py index 4a51e53..ec3dfea 100644 --- a/mainscripts/Util.py +++ b/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] ) + \ No newline at end of file