mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-03-12 20:42:45 -07:00
231 lines
7.4 KiB
Python
231 lines
7.4 KiB
Python
import pickle
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
|
|
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 *
|
|
|
|
|
|
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)
|
|
dflimg = DFLIMG.load (filepath)
|
|
|
|
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)
|
|
|
|
dflimg = DFLIMG.load (filepath)
|
|
if dflimg is None:
|
|
io.log_err ("%s is not a dfl image file" % (filepath.name) )
|
|
return
|
|
|
|
dflimg.remove_ie_polys()
|
|
dflimg.embed_and_set( str(filepath) )
|
|
|
|
|
|
def remove_ie_polys_folder(input_path):
|
|
input_path = Path(input_path)
|
|
|
|
io.log_info ("Removing ie_polys...\r\n")
|
|
|
|
for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Removing"):
|
|
filepath = Path(filepath)
|
|
remove_ie_polys_file(filepath)
|
|
|
|
def remove_fanseg_file (filepath):
|
|
filepath = Path(filepath)
|
|
|
|
dflimg = DFLIMG.load (filepath)
|
|
|
|
if dflimg is None:
|
|
io.log_err ("%s is not a dfl image file" % (filepath.name) )
|
|
return
|
|
|
|
dflimg.remove_fanseg_mask()
|
|
dflimg.embed_and_set( str(filepath) )
|
|
|
|
|
|
def remove_fanseg_folder(input_path):
|
|
input_path = Path(input_path)
|
|
|
|
io.log_info ("Removing fanseg mask...\r\n")
|
|
|
|
for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Removing"):
|
|
filepath = Path(filepath)
|
|
remove_fanseg_file(filepath)
|
|
|
|
def convert_png_to_jpg_file (filepath):
|
|
filepath = Path(filepath)
|
|
|
|
if filepath.suffix != '.png':
|
|
return
|
|
|
|
dflpng = DFLPNG.load (str(filepath) )
|
|
if dflpng is None:
|
|
io.log_err ("%s is not a dfl png image file" % (filepath.name) )
|
|
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), 100])
|
|
|
|
DFLJPG.embed_data( new_filepath,
|
|
face_type=dfl_dict.get('face_type', None),
|
|
landmarks=dfl_dict.get('landmarks', None),
|
|
ie_polys=dfl_dict.get('ie_polys', None),
|
|
source_filename=dfl_dict.get('source_filename', None),
|
|
source_rect=dfl_dict.get('source_rect', None),
|
|
source_landmarks=dfl_dict.get('source_landmarks', None) )
|
|
|
|
filepath.unlink()
|
|
|
|
def convert_png_to_jpg_folder (input_path):
|
|
input_path = Path(input_path)
|
|
|
|
io.log_info ("Converting PNG to JPG...\r\n")
|
|
|
|
for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Converting"):
|
|
filepath = Path(filepath)
|
|
convert_png_to_jpg_file(filepath)
|
|
|
|
def add_landmarks_debug_images(input_path):
|
|
io.log_info ("Adding landmarks debug images...")
|
|
|
|
for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Processing"):
|
|
filepath = Path(filepath)
|
|
|
|
img = cv2_imread(str(filepath))
|
|
|
|
dflimg = DFLIMG.load (filepath)
|
|
|
|
if dflimg is None:
|
|
io.log_err ("%s is not a dfl image file" % (filepath.name) )
|
|
continue
|
|
|
|
if img is not None:
|
|
face_landmarks = dflimg.get_landmarks()
|
|
LandmarksProcessor.draw_landmarks(img, face_landmarks, transparent_mask=True, ie_polys=IEPolys.load(dflimg.get_ie_polys()) )
|
|
|
|
output_file = '{}{}'.format( str(Path(str(input_path)) / filepath.stem), '_debug.jpg')
|
|
cv2_imwrite(output_file, img, [int(cv2.IMWRITE_JPEG_QUALITY), 50] )
|
|
|
|
def recover_original_aligned_filename(input_path):
|
|
io.log_info ("Recovering original aligned filename...")
|
|
|
|
files = []
|
|
for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Processing"):
|
|
filepath = Path(filepath)
|
|
|
|
dflimg = DFLIMG.load (filepath)
|
|
|
|
if dflimg is None:
|
|
io.log_err ("%s is not a dfl image file" % (filepath.name) )
|
|
continue
|
|
|
|
files += [ [filepath, None, dflimg.get_source_filename(), False] ]
|
|
|
|
files_len = len(files)
|
|
for i in io.progress_bar_generator( range(files_len), "Sorting" ):
|
|
fp, _, sf, converted = files[i]
|
|
|
|
if converted:
|
|
continue
|
|
|
|
sf_stem = Path(sf).stem
|
|
|
|
files[i][1] = fp.parent / ( sf_stem + '_0' + fp.suffix )
|
|
files[i][3] = True
|
|
c = 1
|
|
|
|
for j in range(i+1, files_len):
|
|
fp_j, _, sf_j, converted_j = files[j]
|
|
if converted_j:
|
|
continue
|
|
|
|
if sf_j == sf:
|
|
files[j][1] = fp_j.parent / ( sf_stem + ('_%d' % (c)) + fp_j.suffix )
|
|
files[j][3] = True
|
|
c += 1
|
|
|
|
for file in io.progress_bar_generator( files, "Renaming", leave=False ):
|
|
fs, _, _, _ = file
|
|
dst = fs.parent / ( fs.stem + '_tmp' + fs.suffix )
|
|
try:
|
|
fs.rename (dst)
|
|
except:
|
|
io.log_err ('fail to rename %s' % (fs.name) )
|
|
|
|
for file in io.progress_bar_generator( files, "Renaming" ):
|
|
fs, fd, _, _ = file
|
|
fs = fs.parent / ( fs.stem + '_tmp' + fs.suffix )
|
|
try:
|
|
fs.rename (fd)
|
|
except:
|
|
io.log_err ('fail to rename %s' % (fs.name) )
|