Permalink
Browse files

DFL now works with JPG files by default. Old PNG files also supported…

…. Added util convertor PNG to JPG.
  • Loading branch information...
iperov committed Feb 4, 2019
1 parent f0a5f97 commit 6d95dd4a9922e012a86b6947776160cc9b586004
Showing with 398 additions and 90 deletions.
  1. +11 −2 main.py
  2. +20 −7 mainscripts/Converter.py
  3. +3 −3 mainscripts/Extractor.py
  4. +69 −53 mainscripts/Sorter.py
  5. +46 −0 mainscripts/Util.py
  6. +1 −14 nnlib/nnlib.py
  7. +13 −11 samples/SampleLoader.py
  8. +229 −0 utils/DFLJPG.py
  9. +6 −0 utils/struct_utils.py
13 main.py
@@ -53,8 +53,6 @@ def process_extract(arguments):
extract_parser.add_argument('--manual-output-debug-fix', action="store_true", dest="manual_output_debug_fix", default=False, help="Performs manual reextract input-dir frames which were deleted from [output_dir]_debug\ dir.")
extract_parser.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.")
extract_parser.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU. Forces to use MT extractor.")


extract_parser.set_defaults (func=process_extract)

def process_sort(arguments):
@@ -66,6 +64,17 @@ def process_sort(arguments):
sort_parser.add_argument('--by', required=True, dest="sort_by_method", choices=("blur", "face", "face-dissim", "face-yaw", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "final", "test"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
sort_parser.set_defaults (func=process_sort)

def process_util(arguments):
from mainscripts import Util

if arguments.convert_png_to_jpg:
Util.convert_png_to_jpg_folder (input_path=arguments.input_dir)

util_parser = subparsers.add_parser( "util", help="Utilities.")
util_parser.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
util_parser.add_argument('--convert-png-to-jpg', action="store_true", dest="convert_png_to_jpg", default=False, help="Convert DeepFaceLAB PNG files to JPEG.")
util_parser.set_defaults (func=process_util)

def process_train(arguments):
from mainscripts import Trainer
Trainer.main (
@@ -6,6 +6,7 @@
import cv2
from tqdm import tqdm
from utils.DFLPNG import DFLPNG
from utils.DFLJPG import DFLJPG
from utils import image_utils
import shutil
import numpy as np
@@ -163,7 +164,14 @@ def onClientProcessData(self, data):
cv2.waitKey(0)
faces_processed = 1
elif self.converter.get_mode() == ConverterBase.MODE_IMAGE_WITH_LANDMARKS:
image_landmarks = DFLPNG.load( str(filename_path), throw_on_no_embedded_data=True ).get_landmarks()
if filename_path.suffix == '.png':
dflimg = DFLPNG.load( str(filename_path), throw_on_no_embedded_data=True )
elif filename_path.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filename_path), throw_on_no_embedded_data=True )
else:
raise Exception ("%s is not a dfl image file" % (filename_path.name) )

image_landmarks = dflimg.get_landmarks()

image = self.converter.convert_image(image, image_landmarks, self.debug)
if self.debug:
@@ -259,16 +267,21 @@ def main (input_dir, output_dir, model_dir, model_name, aligned_dir=None, **in_o
alignments = {}

aligned_path_image_paths = Path_utils.get_image_paths(aligned_path)
for filename in tqdm(aligned_path_image_paths, desc="Collecting alignments", ascii=True ):
dflpng = DFLPNG.load( str(filename), print_on_no_embedded_data=True )
if dflpng is None:
continue
for filepath in tqdm(aligned_path_image_paths, desc="Collecting alignments", ascii=True ):
filepath = Path(filepath)

if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True )
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True )
else:
print ("%s is not a dfl image file" % (filepath.name) )

source_filename_stem = Path( dflpng.get_source_filename() ).stem
source_filename_stem = Path( dflimg.get_source_filename() ).stem
if source_filename_stem not in alignments.keys():
alignments[ source_filename_stem ] = []

alignments[ source_filename_stem ].append (dflpng.get_source_landmarks())
alignments[ source_filename_stem ].append (dflimg.get_source_landmarks())


#interpolate landmarks
@@ -7,7 +7,7 @@
import numpy as np
import cv2
from utils import Path_utils
from utils.DFLPNG import DFLPNG
from utils.DFLJPG import DFLJPG
from utils import image_utils
from facelib import FaceType
import facelib
@@ -311,7 +311,7 @@ def onClientProcessData(self, data):
debug_image = image.copy()

for (face_idx, face) in enumerate(faces):
output_file = '{}_{}{}'.format(str(self.output_path / filename_path.stem), str(face_idx), '.png')
output_file = '{}_{}{}'.format(str(self.output_path / filename_path.stem), str(face_idx), '.jpg')

rect = face[0]
image_landmarks = np.array(face[1])
@@ -329,7 +329,7 @@ def onClientProcessData(self, data):

cv2.imwrite(output_file, face_image)

DFLPNG.embed_data(output_file, face_type = FaceType.toString(self.face_type),
DFLJPG.embed_data(output_file, face_type = FaceType.toString(self.face_type),
landmarks = face_image_landmarks.tolist(),
yaw_value = facelib.LandmarksProcessor.calc_face_yaw (face_image_landmarks),
pitch_value = facelib.LandmarksProcessor.calc_face_pitch (face_image_landmarks),
@@ -10,6 +10,7 @@
from utils import Path_utils
from utils import image_utils
from utils.DFLPNG import DFLPNG
from utils.DFLJPG import DFLJPG
from facelib import LandmarksProcessor
from utils.SubprocessorBase import SubprocessorBase
import multiprocessing
@@ -86,17 +87,24 @@ def onClientFinalize(self):

#override
def onClientProcessData(self, data):
filename_path = Path( data[0] )

dflpng = DFLPNG.load( str(filename_path), print_on_no_embedded_data=True )
if dflpng is not None:
image = cv2.imread( str(filename_path) )
filepath = Path( data[0] )

if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True )
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True )
else:
print ("%s is not a dfl image file" % (filepath.name) )
dflimg = None

if dflimg is not None:
image = cv2.imread( str(filepath) )
image = ( image * \
LandmarksProcessor.get_image_hull_mask (image.shape, dflpng.get_landmarks()) \
LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks()) \
).astype(np.uint8)
return [ str(filename_path), estimate_sharpness( image ) ]
return [ str(filepath), estimate_sharpness( image ) ]
else:
return [ str(filename_path), 0 ]
return [ str(filepath), 0 ]

#override
def onClientGetDataName (self, data):
@@ -151,15 +159,15 @@ def sort_by_face(input_path):
for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Loading", ascii=True):
filepath = Path(filepath)

if filepath.suffix != '.png':
print ("%s is not a png file required for sort_by_face" % (filepath.name) )
continue

dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True)
if dflpng is None:
if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True )
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True )
else:
print ("%s is not a dfl image file" % (filepath.name) )
continue
img_list.append( [str(filepath), dflpng.get_landmarks()] )

img_list.append( [str(filepath), dflimg.get_landmarks()] )


img_list_len = len(img_list)
@@ -187,15 +195,15 @@ def sort_by_face_dissim(input_path):
for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Loading", ascii=True):
filepath = Path(filepath)

if filepath.suffix != '.png':
print ("%s is not a png file required for sort_by_face_dissim" % (filepath.name) )
continue

dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True)
if dflpng is None:
continue
if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True )
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True )
else:
print ("%s is not a dfl image file" % (filepath.name) )
continue

img_list.append( [str(filepath), dflpng.get_landmarks(), 0 ] )
img_list.append( [str(filepath), dflimg.get_landmarks(), 0 ] )

img_list_len = len(img_list)
for i in tqdm( range(0, img_list_len-1), desc="Sorting", ascii=True):
@@ -220,15 +228,15 @@ def sort_by_face_yaw(input_path):
for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Loading", ascii=True):
filepath = Path(filepath)

if filepath.suffix != '.png':
print ("%s is not a png file required for sort_by_face_dissim" % (filepath.name) )
continue

dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True)
if dflpng is None:
if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True )
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True )
else:
print ("%s is not a dfl image file" % (filepath.name) )
continue

img_list.append( [str(filepath), np.array( dflpng.get_yaw_value() ) ] )
img_list.append( [str(filepath), np.array( dflimg.get_yaw_value() ) ] )

print ("Sorting...")
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
@@ -436,15 +444,22 @@ def sort_by_hist_dissim(input_path):
print ("Sorting by histogram dissimilarity...")

img_list = []
for filename_path in tqdm( Path_utils.get_image_paths(input_path), desc="Loading", ascii=True):
image = cv2.imread(filename_path)
for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Loading", ascii=True):
filepath = Path(filepath)

dflpng = DFLPNG.load( str(filename_path) )
if dflpng is not None:
face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflpng.get_landmarks())
image = (image*face_mask).astype(np.uint8)
if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True )
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filepath), print_on_no_embedded_data=True )
else:
print ("%s is not a dfl image file" % (filepath.name) )
continue

image = cv2.imread(str(filepath))
face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks())
image = (image*face_mask).astype(np.uint8)

img_list.append ([filename_path, cv2.calcHist([cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)], [0], None, [256], [0, 256]), 0 ])
img_list.append ([str(filepath), cv2.calcHist([cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)], [0], None, [256], [0, 256]), 0 ])

img_list = HistDissimSubprocessor(img_list).process()

@@ -512,26 +527,27 @@ def onClientProcessData(self, data):
filepath = Path(data[0])

try:
if filepath.suffix != '.png':
raise Exception ("%s is not a png file required for sort_final" % (filepath.name) )

dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True)
if dflpng is None:
raise Exception ("")
if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True )
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load( str(filepath), print_on_no_embedded_data=True )
else:
print ("%s is not a dfl image file" % (filepath.name) )
raise Exception("")

bgr = cv2.imread(str(filepath))
if bgr is None:
raise Exception ("Unable to load %s" % (filepath.name) )

gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
gray_masked = ( gray * LandmarksProcessor.get_image_hull_mask (bgr.shape, dflpng.get_landmarks() )[:,:,0] ).astype(np.uint8)
gray_masked = ( gray * LandmarksProcessor.get_image_hull_mask (bgr.shape, dflimg.get_landmarks() )[:,:,0] ).astype(np.uint8)
sharpness = estimate_sharpness(gray_masked)
hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
except Exception as e:
print (e)
return [ 1, [str(filepath)] ]

return [ 0, [str(filepath), sharpness, hist, dflpng.get_yaw_value() ] ]
return [ 0, [str(filepath), sharpness, hist, dflimg.get_yaw_value() ] ]


#override
@@ -689,15 +705,15 @@ def sort_by_origname(input_path):
for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Loading", ascii=True):
filepath = Path(filepath)

if filepath.suffix != '.png':
print ("%s is not a png file required for sort_by_origname" % (filepath.name) )
continue

dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True)
if dflpng is None:
if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath), print_on_no_embedded_data=True )
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load( str(filepath), print_on_no_embedded_data=True )
else:
print ("%s is not a dfl image file" % (filepath.name) )
continue

img_list.append( [str(filepath), dflpng.get_source_filename()] )
img_list.append( [str(filepath), dflimg.get_source_filename()] )

print ("Sorting...")
img_list = sorted(img_list, key=operator.itemgetter(1))
@@ -0,0 +1,46 @@
import os
import sys
import operator
import numpy as np
import cv2
from tqdm import tqdm
from shutil import copyfile

from pathlib import Path
from utils import Path_utils
from utils import image_utils
from utils.DFLPNG import DFLPNG
from utils.DFLJPG import DFLJPG
from facelib import LandmarksProcessor
from utils.SubprocessorBase import SubprocessorBase
import multiprocessing

def convert_png_to_jpg_file (filepath):
filepath = Path(filepath)

if filepath.suffix != '.png':
return
dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True)
if dflpng is None:
return

dfl_dict = dflpng.getDFLDictData()

img = cv2.imread (str(filepath))
new_filepath = str(filepath.parent / (filepath.stem + '.jpg'))
cv2.imwrite ( new_filepath, img, [int(cv2.IMWRITE_JPEG_QUALITY), 85])
DFLJPG.embed_data( new_filepath, **dfl_dict )
filepath.unlink()


def convert_png_to_jpg_folder (input_path):
if not all(ord(c) < 128 for c in input_path):
print ("Path to directory must contain only non unicode characters.")
return
input_path = Path(input_path)

print ("Converting PNG to JPG...\r\n")

for filepath in tqdm( Path_utils.get_image_paths(input_path), desc="Converting", ascii=True):
filepath = Path(filepath)
convert_png_to_jpg_file(filepath)
@@ -111,6 +111,7 @@ class nnlib(object):
GroupNormalization = keras_contrib.layers.GroupNormalization
InstanceNormalization = keras_contrib.layers.InstanceNormalization
Padam = keras_contrib.optimizers.Padam
PELU = keras_contrib.layers.advanced_activations.PELU
"""
code_import_dlib_string = \
"""
@@ -448,20 +449,6 @@ def __call__(self,y_true, y_pred):
else:
return (1.0 - tf.image.ssim ((y_true/2+0.5), (y_pred/2+0.5), 1.0)) / 2.0
nnlib.DSSIMLoss = DSSIMLoss

class DSSIMLoss(object):
def __init__(self, is_tanh=False):
self.is_tanh = is_tanh

def __call__(self,y_true, y_pred):

if not self.is_tanh:
loss = (1.0 - tf.image.ssim (y_true, y_pred, 1.0)) / 2.0
else:
loss = (1.0 - tf.image.ssim ( (y_true/2+0.5), (y_pred/2+0.5), 1.0)) / 2.0

return loss
nnlib.DSSIMLoss = DSSIMLoss

class DSSIMMaskLoss(object):
def __init__(self, mask_list, is_tanh=False):
Oops, something went wrong.

0 comments on commit 6d95dd4

Please sign in to comment.