diff --git a/main.py b/main.py index 5428b3b..b480568 100644 --- a/main.py +++ b/main.py @@ -133,7 +133,13 @@ if __name__ == "__main__": if arguments.remove_ie_polys: Util.remove_ie_polys_folder (input_path=arguments.input_dir) - + + if arguments.save_faceset_metadata: + Util.save_faceset_metadata_folder (input_path=arguments.input_dir) + + if arguments.restore_faceset_metadata: + Util.restore_faceset_metadata_folder (input_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.") p.add_argument('--convert-png-to-jpg', action="store_true", dest="convert_png_to_jpg", default=False, help="Convert DeepFaceLAB PNG files to JPEG.") @@ -141,6 +147,8 @@ if __name__ == "__main__": p.add_argument('--recover-original-aligned-filename', action="store_true", dest="recover_original_aligned_filename", default=False, help="Recover original aligned filename.") #p.add_argument('--remove-fanseg', action="store_true", dest="remove_fanseg", default=False, help="Remove fanseg mask from aligned faces.") p.add_argument('--remove-ie-polys', action="store_true", dest="remove_ie_polys", default=False, help="Remove ie_polys from aligned faces.") + p.add_argument('--save-faceset-metadata', action="store_true", dest="save_faceset_metadata", default=False, help="Save faceset metadata to file.") + 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.set_defaults (func=process_util) diff --git a/mainscripts/Util.py b/mainscripts/Util.py index fa6f470..2ded559 100644 --- a/mainscripts/Util.py +++ b/mainscripts/Util.py @@ -1,12 +1,84 @@ import cv2 +import pickle 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 interact import interact as io +from utils import Path_utils +from utils.cv2_utils import * +from utils.DFLJPG import DFLJPG +from utils.DFLPNG import DFLPNG + +def save_faceset_metadata_folder(input_path): + input_path = Path(input_path) + + metadata_filepath = input_path / 'meta.dat' + + io.log_info (f"Saving metadata to {str(metadata_filepath)}\r\n") + + d = {} + for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Processing"): + filepath = Path(filepath) + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath) ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath) ) + else: + continue + + dfl_dict = dflimg.getDFLDictData() + d[filepath.name] = ( dflimg.get_shape(), dfl_dict ) + + try: + with open(metadata_filepath, "wb") as f: + f.write ( pickle.dumps(d) ) + except: + raise Exception( 'cannot save %s' % (filename) ) + + io.log_info("Now you can edit images.") + io.log_info("!!! Keep same filenames in the folder.") + io.log_info("You can change size of images, restoring process will downscale back to original size.") + io.log_info("After that, use restore metadata.") + +def restore_faceset_metadata_folder(input_path): + input_path = Path(input_path) + + metadata_filepath = input_path / 'meta.dat' + io.log_info (f"Restoring metadata from {str(metadata_filepath)}.\r\n") + + if not metadata_filepath.exists(): + io.log_err(f"Unable to find {str(metadata_filepath)}.") + + try: + with open(metadata_filepath, "rb") as f: + d = pickle.loads(f.read()) + except: + raise FileNotFoundError(filename) + + for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Processing"): + filepath = Path(filepath) + + shape, dfl_dict = d.get(filepath.name, None) + + img = cv2_imread (str(filepath)) + if img.shape != shape: + img = cv2.resize (img, (shape[1], shape[0]), cv2.INTER_LANCZOS4 ) + + if filepath.suffix == '.png': + cv2_imwrite (str(filepath), img) + elif filepath.suffix == '.jpg': + cv2_imwrite (str(filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] ) + + if filepath.suffix == '.png': + DFLPNG.embed_dfldict( str(filepath), dfl_dict ) + elif filepath.suffix == '.jpg': + DFLJPG.embed_dfldict( str(filepath), dfl_dict ) + else: + continue + + metadata_filepath.unlink() + def remove_ie_polys_file (filepath): filepath = Path(filepath) @@ -76,7 +148,7 @@ def convert_png_to_jpg_file (filepath): img = cv2_imread (str(filepath)) new_filepath = str(filepath.parent / (filepath.stem + '.jpg')) - cv2_imwrite ( new_filepath, img, [int(cv2.IMWRITE_JPEG_QUALITY), 85]) + cv2_imwrite ( new_filepath, img, [int(cv2.IMWRITE_JPEG_QUALITY), 100]) DFLJPG.embed_data( new_filepath, face_type=dfl_dict.get('face_type', None), diff --git a/utils/DFLJPG.py b/utils/DFLJPG.py index 057c8c8..2a35fe6 100644 --- a/utils/DFLJPG.py +++ b/utils/DFLJPG.py @@ -159,6 +159,17 @@ class DFLJPG(object): print (e) return None + @staticmethod + def embed_dfldict(filename, dfl_dict): + inst = DFLJPG.load_raw (filename) + inst.setDFLDictData (dfl_dict) + + try: + with open(filename, "wb") as f: + f.write ( inst.dump() ) + except: + raise Exception( 'cannot save %s' % (filename) ) + @staticmethod def embed_data(filename, face_type=None, landmarks=None, @@ -185,26 +196,18 @@ class DFLJPG(object): io.log_err("Unable to encode fanseg_mask for %s" % (filename) ) fanseg_mask = None - inst = DFLJPG.load_raw (filename) - inst.setDFLDictData ({ - 'face_type': face_type, - 'landmarks': landmarks, - 'ie_polys' : ie_polys.dump() if ie_polys is not None else None, - 'source_filename': source_filename, - 'source_rect': source_rect, - 'source_landmarks': source_landmarks, - 'image_to_face_mat': image_to_face_mat, - 'fanseg_mask' : fanseg_mask, - 'pitch_yaw_roll' : pitch_yaw_roll, - 'eyebrows_expand_mod' : eyebrows_expand_mod, - 'relighted' : relighted - }) - - try: - with open(filename, "wb") as f: - f.write ( inst.dump() ) - except: - raise Exception( 'cannot save %s' % (filename) ) + DFLJPG.embed_dfldict (filename, {'face_type': face_type, + 'landmarks': landmarks, + 'ie_polys' : ie_polys.dump() if ie_polys is not None else None, + 'source_filename': source_filename, + 'source_rect': source_rect, + 'source_landmarks': source_landmarks, + 'image_to_face_mat': image_to_face_mat, + 'fanseg_mask' : fanseg_mask, + 'pitch_yaw_roll' : pitch_yaw_roll, + 'eyebrows_expand_mod' : eyebrows_expand_mod, + 'relighted' : relighted + }) def embed_and_set(self, filename, face_type=None, landmarks=None, @@ -246,7 +249,7 @@ class DFLJPG(object): def remove_fanseg_mask(self): self.dfl_dict['fanseg_mask'] = None - + def remove_source_filename(self): self.dfl_dict['source_filename'] = None diff --git a/utils/DFLPNG.py b/utils/DFLPNG.py index c25cdd0..f74a186 100644 --- a/utils/DFLPNG.py +++ b/utils/DFLPNG.py @@ -275,6 +275,17 @@ class DFLPNG(object): print(e) return None + @staticmethod + def embed_dfldict(filename, dfl_dict): + inst = DFLPNG.load_raw (filename) + inst.setDFLDictData (dfl_dict) + + try: + with open(filename, "wb") as f: + f.write ( inst.dump() ) + except: + raise Exception( 'cannot save %s' % (filename) ) + @staticmethod def embed_data(filename, face_type=None, landmarks=None, @@ -301,26 +312,18 @@ class DFLPNG(object): io.log_err("Unable to encode fanseg_mask for %s" % (filename) ) fanseg_mask = None - inst = DFLPNG.load_raw (filename) - inst.setDFLDictData ({ - 'face_type': face_type, - 'landmarks': landmarks, - 'ie_polys' : ie_polys.dump() if ie_polys is not None else None, - 'source_filename': source_filename, - 'source_rect': source_rect, - 'source_landmarks': source_landmarks, - 'image_to_face_mat':image_to_face_mat, - 'fanseg_mask' : fanseg_mask, - 'pitch_yaw_roll' : pitch_yaw_roll, - 'eyebrows_expand_mod' : eyebrows_expand_mod, - 'relighted' : relighted - }) - - try: - with open(filename, "wb") as f: - f.write ( inst.dump() ) - except: - raise Exception( 'cannot save %s' % (filename) ) + DFLPNG.embed_dfldict (filename, {'face_type': face_type, + 'landmarks': landmarks, + 'ie_polys' : ie_polys.dump() if ie_polys is not None else None, + 'source_filename': source_filename, + 'source_rect': source_rect, + 'source_landmarks': source_landmarks, + 'image_to_face_mat':image_to_face_mat, + 'fanseg_mask' : fanseg_mask, + 'pitch_yaw_roll' : pitch_yaw_roll, + 'eyebrows_expand_mod' : eyebrows_expand_mod, + 'relighted' : relighted + }) def embed_and_set(self, filename, face_type=None, landmarks=None,