diff --git a/main.py b/main.py index 38d75de..9d0267f 100644 --- a/main.py +++ b/main.py @@ -6,6 +6,7 @@ import multiprocessing from utils import Path_utils from utils import os_utils from pathlib import Path +from interact import interact as io if sys.version_info[0] < 3 or (sys.version_info[0] == 3 and sys.version_info[1] < 6): raise Exception("This program requires at least Python 3.6") @@ -134,10 +135,12 @@ if __name__ == "__main__": Util.restore_faceset_metadata_folder (input_path=arguments.input_dir) if arguments.pack_faceset: + io.log_info ("Performing faceset packing...\r\n") from samplelib import PackedFaceset PackedFaceset.pack( Path(arguments.input_dir) ) if arguments.unpack_faceset: + io.log_info ("Performing faceset unpacking...\r\n") from samplelib import PackedFaceset PackedFaceset.unpack( Path(arguments.input_dir) ) diff --git a/samplelib/PackedFaceset.py b/samplelib/PackedFaceset.py index 3805281..96525d5 100644 --- a/samplelib/PackedFaceset.py +++ b/samplelib/PackedFaceset.py @@ -21,39 +21,40 @@ class PackedFaceset(): io.log_info(f"{samples_dat_path} : file already exists !") io.input_bool("Press enter to continue and overwrite.", False) - of = open(samples_dat_path, "wb") - as_person_faceset = False dir_names = Path_utils.get_all_dir_names(samples_path) if len(dir_names) != 0: as_person_faceset = io.input_bool(f"{len(dir_names)} subdirectories found, process as person faceset? (y/n) skip:y : ", True) - + if as_person_faceset: image_paths = [] - + for dir_name in dir_names: image_paths += Path_utils.get_image_paths(samples_path / dir_name) else: image_paths = Path_utils.get_image_paths(samples_path) - samples = samplelib.SampleHost.load_face_samples(image_paths) samples_len = len(samples) samples_configs = [] - for sample in samples: + for sample in io.progress_bar_generator (samples, "Processing"): sample_filepath = Path(sample.filename) sample.filename = sample_filepath.name - + if as_person_faceset: sample.person_name = sample_filepath.parent.name samples_configs.append ( sample.get_config() ) samples_bytes = pickle.dumps(samples_configs, 4) + of = open(samples_dat_path, "wb") of.write ( struct.pack ("Q", PackedFaceset.VERSION ) ) of.write ( struct.pack ("Q", len(samples_bytes) ) ) of.write ( samples_bytes ) + del samples_bytes #just free mem + del samples_configs + sample_data_table_offset = of.tell() of.write ( bytes( 8*(samples_len+1) ) ) #sample data offset table @@ -66,16 +67,14 @@ class PackedFaceset(): sample_path = samples_path / sample.person_name / sample.filename else: sample_path = samples_path / sample.filename - + + with open(sample_path, "rb") as f: - b = f.read() + b = f.read() offsets.append ( of.tell() - data_start_offset ) of.write(b) except: - import code - code.interact(local=dict(globals(), **locals())) - raise Exception(f"error while processing sample {sample_path}") offsets.append ( of.tell() ) @@ -86,11 +85,11 @@ class PackedFaceset(): of.seek(0,2) of.close() - for filename in io.progress_bar_generator(image_paths,"Deleting"): + for filename in io.progress_bar_generator(image_paths, "Deleting files"): Path(filename).unlink() if as_person_faceset: - for dir_name in dir_names: + for dir_name in io.progress_bar_generator(dir_names, "Deleting dirs"): dir_path = samples_path / dir_name try: shutil.rmtree(dir_path) @@ -108,10 +107,10 @@ class PackedFaceset(): for sample in io.progress_bar_generator(samples, "Unpacking"): person_name = sample.person_name - if person_name is not None: + if person_name is not None: person_path = samples_path / person_name person_path.mkdir(parents=True, exist_ok=True) - + target_filepath = person_path / sample.filename else: target_filepath = samples_path / sample.filename @@ -137,8 +136,8 @@ class PackedFaceset(): samples_configs = pickle.loads ( f.read(sizeof_samples_bytes) ) samples = [] for sample_config in samples_configs: - samples.append ( Sample (**sample_config) ) - + samples.append ( Sample (**sample_config) ) + offsets = [ struct.unpack("Q", f.read(8) )[0] for _ in range(len(samples)+1) ] data_start_offset = f.tell() f.close() diff --git a/samplelib/SampleHost.py b/samplelib/SampleHost.py index 54f4450..7e1879e 100644 --- a/samplelib/SampleHost.py +++ b/samplelib/SampleHost.py @@ -91,9 +91,8 @@ class SampleHost: landmarks, ie_polys, eyebrows_expand_mod, - source_filename, - - ) in io.progress_bar_generator(result, "Processing"): + source_filename, + ) in result: sample_list.append( Sample(filename=filename, sample_type=SampleType.FACE, face_type=FaceType.fromString (face_type), @@ -104,31 +103,6 @@ class SampleHost: source_filename=source_filename, )) return sample_list - - """ - sample_list = [] - - for filename in io.progress_bar_generator(image_paths, "Loading"): - filename_path = Path(filename) - try: - dflimg = DFLIMG.load (filename_path) - - - - sample_list.append( Sample(filename=filename, - sample_type=SampleType.FACE, - face_type=FaceType.fromString (dflimg.get_face_type()), - shape=dflimg.get_shape(), - landmarks=dflimg.get_landmarks(), - ie_polys=dflimg.get_ie_polys(), - eyebrows_expand_mod=dflimg.get_eyebrows_expand_mod(), - source_filename=dflimg.get_source_filename(), - )) - except: - io.log_err ("Unable to load %s , error: %s" % (filename, traceback.format_exc() ) ) - - return sample_list - """ @staticmethod def upgradeToFaceTemporalSortedSamples( samples ): diff --git a/utils/mp_utils.py b/utils/mp_utils.py index 9ec5406..02b04bd 100644 --- a/utils/mp_utils.py +++ b/utils/mp_utils.py @@ -4,7 +4,6 @@ import time import numpy as np - class Index2DHost(): """ Provides random shuffled 2D indexes for multiprocesses