diff --git a/DFLIMG/DFLIMG.py b/DFLIMG/DFLIMG.py new file mode 100644 index 0000000..b65be11 --- /dev/null +++ b/DFLIMG/DFLIMG.py @@ -0,0 +1,15 @@ +from pathlib import Path + +from .DFLJPG import DFLJPG +from .DFLPNG import DFLPNG + +class DFLIMG(): + + @staticmethod + def load(filepath, loader_func=None): + if filepath.suffix == '.png': + return DFLPNG.load( str(filepath), loader_func=loader_func ) + elif filepath.suffix == '.jpg': + return DFLJPG.load ( str(filepath), loader_func=loader_func ) + else: + return None diff --git a/utils/DFLJPG.py b/DFLIMG/DFLJPG.py similarity index 100% rename from utils/DFLJPG.py rename to DFLIMG/DFLJPG.py diff --git a/utils/DFLPNG.py b/DFLIMG/DFLPNG.py similarity index 100% rename from utils/DFLPNG.py rename to DFLIMG/DFLPNG.py diff --git a/DFLIMG/__init__.py b/DFLIMG/__init__.py new file mode 100644 index 0000000..454369f --- /dev/null +++ b/DFLIMG/__init__.py @@ -0,0 +1,3 @@ +from .DFLIMG import DFLIMG +from .DFLJPG import DFLJPG +from .DFLPNG import DFLPNG \ No newline at end of file diff --git a/mainscripts/Converter.py b/mainscripts/Converter.py index 41df02a..6a55de5 100644 --- a/mainscripts/Converter.py +++ b/mainscripts/Converter.py @@ -14,17 +14,16 @@ import numpy as np import numpy.linalg as npla import imagelib +import samplelib from converters import (ConverterConfig, ConvertFaceAvatar, ConvertMasked, FrameInfo) from facelib import FaceType, LandmarksProcessor -from nnlib import TernausNet - from interact import interact as io from joblib import SubprocessFunctionCaller, Subprocessor +from nnlib import TernausNet from utils import Path_utils from utils.cv2_utils import * -from utils.DFLJPG import DFLJPG -from utils.DFLPNG import DFLPNG +from DFLIMG import DFLIMG from .ConverterScreen import Screen, ScreenManager @@ -670,19 +669,29 @@ def main (args, device_args): io.log_err('Aligned directory not found. Please ensure it exists.') return + packed_samples = None + try: + packed_samples = samplelib.PackedFaceset.load(aligned_path) + except: + io.log_err(f"Error occured while loading samplelib.PackedFaceset.load {str(aligned_path)}, {traceback.format_exc()}") + + + if packed_samples is not None: + io.log_info ("Using packed faceset.") + def generator(): + for sample in io.progress_bar_generator( packed_samples, "Collecting alignments"): + filepath = Path(sample.filename) + yield DFLIMG.load(filepath, loader_func=lambda x: sample.read_raw_file() ) + else: + def generator(): + for filepath in io.progress_bar_generator( Path_utils.get_image_paths(aligned_path), "Collecting alignments"): + filepath = Path(filepath) + yield DFLIMG.load(filepath) + alignments = {} multiple_faces_detected = False - aligned_path_image_paths = Path_utils.get_image_paths(aligned_path) - for filepath in io.progress_bar_generator(aligned_path_image_paths, "Collecting alignments"): - filepath = Path(filepath) - - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None - + + for dflimg in generator(): if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) continue @@ -745,14 +754,8 @@ def main (args, device_args): filesdata = [] for filepath in io.progress_bar_generator(input_path_image_paths, "Collecting info"): filepath = Path(filepath) - - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None - + + dflimg = DFLIMG.load(filepath) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) continue diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index cb0c9b2..8512f0e 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -20,8 +20,7 @@ from joblib import Subprocessor from nnlib import TernausNet, nnlib from utils import Path_utils from utils.cv2_utils import * -from utils.DFLJPG import DFLJPG -from utils.DFLPNG import DFLPNG +from DFLIMG import * DEBUG = False @@ -136,10 +135,7 @@ class ExtractSubprocessor(Subprocessor): h, w, ch = image.shape if h == w: #extracting from already extracted jpg image? - if filename_path.suffix == '.png': - src_dflimg = DFLPNG.load ( str(filename_path) ) - if filename_path.suffix == '.jpg': - src_dflimg = DFLJPG.load ( str(filename_path) ) + src_dflimg = DFLIMG.load (filename_path) if 'rects' in self.type: if min(w,h) < 128: @@ -810,13 +806,7 @@ def extract_fanseg(input_dir, device_args={} ): paths_to_extract = [] for filename in Path_utils.get_image_paths(input_path) : filepath = Path(filename) - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None - + dflimg = DFLIMG.load ( filepath ) if dflimg is not None: paths_to_extract.append (filepath) diff --git a/mainscripts/FacesetRelighter.py b/mainscripts/FacesetRelighter.py index 9371036..238ab74 100644 --- a/mainscripts/FacesetRelighter.py +++ b/mainscripts/FacesetRelighter.py @@ -6,8 +6,7 @@ from interact import interact as io from nnlib import DeepPortraitRelighting from utils import Path_utils from utils.cv2_utils import * -from utils.DFLJPG import DFLJPG -from utils.DFLPNG import DFLPNG +from DFLIMG import * class RelightEditor: def __init__(self, image_paths, dpr, lighten): @@ -183,12 +182,7 @@ def relight(input_dir, lighten=None, random_one=None): filtered_image_paths = [] for filepath in io.progress_bar_generator(image_paths, "Collecting fileinfo"): try: - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None + dflimg = DFLIMG.load (Path(filepath)) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) @@ -210,13 +204,7 @@ def relight(input_dir, lighten=None, random_one=None): for filepath in io.progress_bar_generator(image_paths, "Relighting"): try: - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None - + dflimg = DFLIMG.load ( Path(filepath) ) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) continue @@ -262,12 +250,7 @@ def delete_relighted(input_dir): files_to_delete = [] for filepath in io.progress_bar_generator(image_paths, "Loading"): - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None + dflimg = DFLIMG.load ( Path(filepath) ) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) diff --git a/mainscripts/MaskEditorTool.py b/mainscripts/MaskEditorTool.py index e58ebc0..de48f6f 100644 --- a/mainscripts/MaskEditorTool.py +++ b/mainscripts/MaskEditorTool.py @@ -9,13 +9,13 @@ import numpy as np import numpy.linalg as npl import imagelib +from DFLIMG import * from facelib import LandmarksProcessor from imagelib import IEPolys 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 + class MaskEditor: STATE_NONE=0 @@ -396,12 +396,7 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None, no_default cached_images[path.name] = cv2_imread(str(path)) / 255.0 if filepath is not None: - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None + dflimg = DFLIMG.load (filepath) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) @@ -573,4 +568,3 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None, no_default io.process_messages(0.005) io.destroy_all_windows() - diff --git a/mainscripts/Sorter.py b/mainscripts/Sorter.py index 0afa82a..2465667 100644 --- a/mainscripts/Sorter.py +++ b/mainscripts/Sorter.py @@ -19,9 +19,7 @@ from joblib import Subprocessor from nnlib import VGGFace, nnlib from utils import Path_utils from utils.cv2_utils import * -from utils.DFLJPG import DFLJPG -from utils.DFLPNG import DFLPNG - +from DFLIMG import * class BlurEstimatorSubprocessor(Subprocessor): class Cli(Subprocessor.Cli): @@ -33,13 +31,7 @@ class BlurEstimatorSubprocessor(Subprocessor): #override def process_data(self, data): filepath = Path( data[0] ) - - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None + dflimg = DFLIMG.load (filepath) if dflimg is not None: image = cv2_imread( str(filepath) ) @@ -119,12 +111,7 @@ def sort_by_face(input_path): for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None + dflimg = DFLIMG.load (filepath) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) @@ -160,12 +147,7 @@ def sort_by_face_dissim(input_path): for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None + dflimg = DFLIMG.load (filepath) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) @@ -198,12 +180,7 @@ def sort_by_face_yaw(input_path): for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None + dflimg = DFLIMG.load (filepath) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) @@ -230,12 +207,7 @@ def sort_by_face_pitch(input_path): for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None + dflimg = DFLIMG.load (filepath) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) @@ -424,12 +396,7 @@ def sort_by_hist_dissim(input_path): for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None + dflimg = DFLIMG.load (filepath) image = cv2_imread(str(filepath)) @@ -481,12 +448,7 @@ def sort_by_origname(input_path): for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Loading"): filepath = Path(filepath) - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load( str(filepath) ) - else: - dflimg = None + dflimg = DFLIMG.load (filepath) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) @@ -528,12 +490,7 @@ class FinalLoaderSubprocessor(Subprocessor): filepath = Path(data[0]) try: - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load( str(filepath) ) - else: - dflimg = None + dflimg = DFLIMG.load (filepath) if dflimg is None: self.log_err("%s is not a dfl image file" % (filepath.name)) diff --git a/mainscripts/Util.py b/mainscripts/Util.py index 2ded559..6208412 100644 --- a/mainscripts/Util.py +++ b/mainscripts/Util.py @@ -6,9 +6,7 @@ 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 - +from DFLIMG import * def save_faceset_metadata_folder(input_path): input_path = Path(input_path) @@ -20,12 +18,7 @@ def save_faceset_metadata_folder(input_path): 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 + dflimg = DFLIMG.load (filepath) dfl_dict = dflimg.getDFLDictData() d[filepath.name] = ( dflimg.get_shape(), dfl_dict ) @@ -82,13 +75,7 @@ def restore_faceset_metadata_folder(input_path): def remove_ie_polys_file (filepath): filepath = Path(filepath) - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - return - + dflimg = DFLIMG.load (filepath) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) return @@ -109,12 +96,7 @@ def remove_ie_polys_folder(input_path): def remove_fanseg_file (filepath): filepath = Path(filepath) - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - return + dflimg = DFLIMG.load (filepath) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) @@ -141,7 +123,7 @@ def convert_png_to_jpg_file (filepath): dflpng = DFLPNG.load (str(filepath) ) if dflpng is None: - io.log_err ("%s is not a dfl image file" % (filepath.name) ) + io.log_err ("%s is not a dfl png image file" % (filepath.name) ) return dfl_dict = dflpng.getDFLDictData() @@ -177,12 +159,7 @@ def add_landmarks_debug_images(input_path): img = cv2_imread(str(filepath)) - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None + dflimg = DFLIMG.load (filepath) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) @@ -202,12 +179,7 @@ def recover_original_aligned_filename(input_path): 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: - dflimg = None + dflimg = DFLIMG.load (filepath) if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) diff --git a/mainscripts/dev_misc.py b/mainscripts/dev_misc.py index 5718416..a9897d2 100644 --- a/mainscripts/dev_misc.py +++ b/mainscripts/dev_misc.py @@ -5,13 +5,12 @@ from pathlib import Path import cv2 import numpy as np +from DFLIMG import DFLIMG from facelib import FaceType, LandmarksProcessor from interact import interact as io from joblib import Subprocessor from utils import Path_utils from utils.cv2_utils import * -from utils.DFLJPG import DFLJPG -from utils.DFLPNG import DFLPNG from . import Extractor, Sorter from .Extractor import ExtractSubprocessor @@ -227,13 +226,8 @@ class CelebAMASKHQSubprocessor(Subprocessor): #override def process_data(self, data): filename = data[0] - filepath = Path(filename) - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None + + dflimg = DFLIMG.load(Path(filename)) image_to_face_mat = dflimg.get_image_to_face_mat() src_filename = dflimg.get_source_filename() @@ -330,12 +324,7 @@ def apply_celebamaskhq(input_dir ): paths_to_extract = [] for filename in io.progress_bar_generator(Path_utils.get_image_paths(img_path), desc="Processing"): filepath = Path(filename) - if filepath.suffix == '.png': - dflimg = DFLPNG.load( str(filepath) ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath) ) - else: - dflimg = None + dflimg = DFLIMG.load(filepath) if dflimg is not None: paths_to_extract.append (filepath) diff --git a/samplelib/Sample.py b/samplelib/Sample.py index ef9f8b3..6e481e3 100644 --- a/samplelib/Sample.py +++ b/samplelib/Sample.py @@ -5,8 +5,7 @@ import cv2 import numpy as np from utils.cv2_utils import * -from utils.DFLJPG import DFLJPG -from utils.DFLPNG import DFLPNG +from DFLIMG import * class SampleType(IntEnum): @@ -87,13 +86,8 @@ class Sample(object): def load_fanseg_mask(self): if self.fanseg_mask_exist: - filepath = Path(self.filename) - if filepath.suffix == '.png': - dflimg = DFLPNG.load ( str(filepath), loader_func=self.read_raw_file ) - elif filepath.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filepath), loader_func=self.read_raw_file ) - else: - dflimg = None + filepath = Path(self.filename) + dflimg = DFLIMG.load (filepath) return dflimg.get_fanseg_mask() return None diff --git a/samplelib/SampleHost.py b/samplelib/SampleHost.py index ad9e84e..a165d32 100644 --- a/samplelib/SampleHost.py +++ b/samplelib/SampleHost.py @@ -5,8 +5,7 @@ from pathlib import Path from facelib import FaceType, LandmarksProcessor from interact import interact as io from utils import Path_utils, mp_utils -from utils.DFLJPG import DFLJPG -from utils.DFLPNG import DFLPNG +from DFLIMG import * from .Sample import Sample, SampleType @@ -41,7 +40,7 @@ class SampleHost: io.log_err(f"Error occured while loading samplelib.PackedFaceset.load {str(samples_dat_path)}, {traceback.format_exc()}") if result is not None: - io.log_info (f"Loaded {len(result)} packed samples from {samples_path}") + io.log_info (f"Loaded {len(result)} packed faces from {samples_path}") if result is None: result = SampleHost.load_face_samples( Path_utils.get_image_paths(samples_path) ) @@ -75,12 +74,7 @@ class SampleHost: for filename in (image_paths if silent else io.progress_bar_generator( image_paths, "Loading")): filename_path = Path(filename) try: - if filename_path.suffix == '.png': - dflimg = DFLPNG.load ( str(filename_path) ) - elif filename_path.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filename_path) ) - else: - dflimg = None + dflimg = DFLIMG.load (filename_path) if dflimg is None: io.log_err ("load_face_samples: %s is not a dfl image file required for training" % (filename_path.name) )