refactoring

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

DFLIMG/__init__.py

@@ -0,0 +1,3 @@
+from .DFLIMG import DFLIMG
+from .DFLJPG import DFLJPG
+from .DFLPNG import DFLPNG

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 DFLIMG import DFLIMG
-from utils.DFLPNG import DFLPNG
 
 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"):
+            for dflimg in generator():
-                filepath = Path(filepath)
-
-                if filepath.suffix == '.png':
-                    dflimg = DFLPNG.load( str(filepath) )
-                elif filepath.suffix == '.jpg':
-                    dflimg = DFLJPG.load ( str(filepath) )
-                else:
-                    dflimg = None
-
                 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 = DFLIMG.load(filepath)
-                    dflimg = DFLPNG.load( str(filepath) )
-                elif filepath.suffix == '.jpg':
-                    dflimg = DFLJPG.load ( str(filepath) )
-                else:
-                    dflimg = None
-
                 if dflimg is None:
                     io.log_err ("%s is not a dfl image file" % (filepath.name) )
                     continue

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 DFLIMG import *
-from utils.DFLPNG import DFLPNG
 
 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 = DFLIMG.load (filename_path)
-                    src_dflimg = DFLPNG.load ( str(filename_path) )
-                if filename_path.suffix == '.jpg':
-                    src_dflimg = DFLJPG.load ( str(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 = DFLIMG.load ( filepath )
-            dflimg = DFLPNG.load( str(filepath) )
-        elif filepath.suffix == '.jpg':
-            dflimg = DFLJPG.load ( str(filepath) )
-        else:
-            dflimg = None
-
         if dflimg is not None:
             paths_to_extract.append (filepath)
     

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 DFLIMG import *
-from utils.DFLPNG import DFLPNG
 
 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 = DFLIMG.load (Path(filepath))
-                dflimg = DFLPNG.load( str(filepath) )
-            elif filepath.suffix == '.jpg':
-                dflimg = DFLJPG.load ( str(filepath) )
-            else:
-                dflimg = None
 
             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 = DFLIMG.load ( Path(filepath) )
-                dflimg = DFLPNG.load( str(filepath) )
-            elif filepath.suffix == '.jpg':
-                dflimg = DFLJPG.load ( str(filepath) )
-            else:
-                dflimg = None
-
             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 = DFLIMG.load ( Path(filepath) )
-            dflimg = DFLPNG.load( str(filepath) )
-        elif filepath.suffix == '.jpg':
-            dflimg = DFLJPG.load ( str(filepath) )
-        else:
-            dflimg = None
 
         if dflimg is None:
             io.log_err ("%s is not a dfl image file" % (filepath.name) )

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 = DFLIMG.load (filepath)
-                dflimg = DFLPNG.load( str(filepath) )
-            elif filepath.suffix == '.jpg':
-                dflimg = DFLJPG.load ( str(filepath) )
-            else:
-                dflimg = None
 
             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()
-

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 DFLIMG import *
-from utils.DFLPNG import DFLPNG
-
 
 class BlurEstimatorSubprocessor(Subprocessor):
     class Cli(Subprocessor.Cli):
@@ -33,13 +31,7 @@ class BlurEstimatorSubprocessor(Subprocessor):
         #override
         def process_data(self, data):
             filepath = Path( data[0] )
-
+            dflimg = DFLIMG.load (filepath)
-            if filepath.suffix == '.png':
-                dflimg = DFLPNG.load( str(filepath) )
-            elif filepath.suffix == '.jpg':
-                dflimg = DFLJPG.load ( str(filepath) )
-            else:
-                dflimg = None
 
             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 = DFLIMG.load (filepath)
-            dflimg = DFLPNG.load( str(filepath) )
-        elif filepath.suffix == '.jpg':
-            dflimg = DFLJPG.load ( str(filepath) )
-        else:
-            dflimg = None
 
         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 = DFLIMG.load (filepath)
-            dflimg = DFLPNG.load( str(filepath) )
-        elif filepath.suffix == '.jpg':
-            dflimg = DFLJPG.load ( str(filepath) )
-        else:
-            dflimg = None
 
         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 = DFLIMG.load (filepath)
-            dflimg = DFLPNG.load( str(filepath) )
-        elif filepath.suffix == '.jpg':
-            dflimg = DFLJPG.load ( str(filepath) )
-        else:
-            dflimg = None
 
         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 = DFLIMG.load (filepath)
-            dflimg = DFLPNG.load( str(filepath) )
-        elif filepath.suffix == '.jpg':
-            dflimg = DFLJPG.load ( str(filepath) )
-        else:
-            dflimg = None
 
         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 = DFLIMG.load (filepath)
-            dflimg = DFLPNG.load( str(filepath) )
-        elif filepath.suffix == '.jpg':
-            dflimg = DFLJPG.load ( str(filepath) )
-        else:
-            dflimg = None
 
         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 = DFLIMG.load (filepath)
-            dflimg = DFLPNG.load( str(filepath) )
-        elif filepath.suffix == '.jpg':
-            dflimg = DFLJPG.load( str(filepath) )
-        else:
-            dflimg = None
 
         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 = DFLIMG.load (filepath)
-                    dflimg = DFLPNG.load( str(filepath) )
-                elif filepath.suffix == '.jpg':
-                    dflimg = DFLJPG.load( str(filepath) )
-                else:
-                    dflimg = None
 
                 if dflimg is None:
                     self.log_err("%s is not a dfl image file" % (filepath.name))

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 DFLIMG import *
-from utils.DFLPNG import DFLPNG
-
 
 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 = DFLIMG.load (filepath)
-            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 )
@@ -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 = DFLIMG.load (filepath)
-        dflimg = DFLPNG.load( str(filepath) )
-    elif filepath.suffix == '.jpg':
-        dflimg = DFLJPG.load ( str(filepath) )
-    else:
-        return
-
     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 = DFLIMG.load (filepath)
-        dflimg = DFLPNG.load( str(filepath) )
-    elif filepath.suffix == '.jpg':
-        dflimg = DFLJPG.load ( str(filepath) )
-    else:
-        return
 
     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 = DFLIMG.load (filepath)
-            dflimg = DFLPNG.load( str(filepath) )
-        elif filepath.suffix == '.jpg':
-            dflimg = DFLJPG.load ( str(filepath) )
-        else:
-            dflimg = None
 
         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 = DFLIMG.load (filepath)
-            dflimg = DFLPNG.load( str(filepath) )
-        elif filepath.suffix == '.jpg':
-            dflimg = DFLJPG.load ( str(filepath) )
-        else:
-            dflimg = None
 
         if dflimg is None:
             io.log_err ("%s is not a dfl image file" % (filepath.name) )

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 = DFLIMG.load(Path(filename))
-                dflimg = DFLPNG.load( str(filepath) )
-            elif filepath.suffix == '.jpg':
-                dflimg = DFLJPG.load ( str(filepath) )
-            else:
-                dflimg = None
                 
             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 = DFLIMG.load(filepath)
-            dflimg = DFLPNG.load( str(filepath) )
-        elif filepath.suffix == '.jpg':
-            dflimg = DFLJPG.load ( str(filepath) )
-        else:
-            dflimg = None
 
         if dflimg is not None:
             paths_to_extract.append (filepath)

samplelib/Sample.py

@@ -5,8 +5,7 @@ import cv2
 import numpy as np
 
 from utils.cv2_utils import *
-from utils.DFLJPG import DFLJPG
+from DFLIMG import *
-from utils.DFLPNG import DFLPNG
 
 
 class SampleType(IntEnum):
@@ -87,13 +86,8 @@ class Sample(object):
 
     def load_fanseg_mask(self):
         if self.fanseg_mask_exist:
-            filepath = Path(self.filename)
+            filepath = Path(self.filename)            
-            if filepath.suffix == '.png':
+            dflimg = DFLIMG.load (filepath)
-                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
             return dflimg.get_fanseg_mask()
 
         return None

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 DFLIMG import *
-from utils.DFLPNG import DFLPNG
 
 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 = DFLIMG.load (filename_path)                    
-                    dflimg = DFLPNG.load ( str(filename_path) )
-                elif filename_path.suffix == '.jpg':
-                    dflimg = DFLJPG.load ( str(filename_path) )
-                else:
-                    dflimg = None
 
                 if dflimg is None:
                     io.log_err ("load_face_samples: %s is not a dfl image file required for training" % (filename_path.name) )