github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-06 04:52:13 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactoring

Browse source
This commit is contained in:
Colombo 2019-12-22 14:05:41 +04:00
parent 246129a964
commit 8153e90ba3
13 changed files with 80 additions and 186 deletions
Download patch file Download diff file Expand all files Collapse all files

15
DFLIMG/DFLIMG.py Normal file
View file

@ -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

0
utils/DFLJPG.py → DFLIMG/DFLJPG.py
View file

0
utils/DFLPNG.py → DFLIMG/DFLPNG.py
View file

3
DFLIMG/__init__.py Normal file
View file

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

49
mainscripts/Converter.py
View file

@ -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

16
mainscripts/Extractor.py
View file

@ -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)

25
mainscripts/FacesetRelighter.py
View file

@ -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) )

12
mainscripts/MaskEditorTool.py
View file

@ -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()

61
mainscripts/Sorter.py
View file

@ -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))

42
mainscripts/Util.py
View file

@ -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) )

19
mainscripts/dev_misc.py
View file

@ -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)

12
samplelib/Sample.py
View file

@ -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

12
samplelib/SampleHost.py
View file

@ -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) )