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

FacesetResizer now supports changing face type

Browse source
This commit is contained in:
iperov 2021-05-06 10:52:04 +04:00
parent d204e049d1
commit aa26089032
1 changed files with 91 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

121
mainscripts/FacesetResizer.py
View file

@ -1,21 +1,23 @@
import multiprocessing import multiprocessing
import shutil import shutil
import cv2
from core import pathex from core import pathex
from core.cv2ex import * from core.cv2ex import *
from core.interact import interact as io from core.interact import interact as io
from core.joblib import Subprocessor from core.joblib import Subprocessor
from DFLIMG import * from DFLIMG import *
from facelib import LandmarksProcessor, FaceType from facelib import FaceType, LandmarksProcessor
class FacesetResizerSubprocessor(Subprocessor): class FacesetResizerSubprocessor(Subprocessor):
#override #override
def __init__(self, image_paths, output_dirpath, image_size): def __init__(self, image_paths, output_dirpath, image_size, face_type=None):
self.image_paths = image_paths self.image_paths = image_paths
self.output_dirpath = output_dirpath self.output_dirpath = output_dirpath
self.image_size = image_size self.image_size = image_size
self.face_type = face_type
self.result = [] self.result = []
super().__init__('FacesetResizer', FacesetResizerSubprocessor.Cli, 600) super().__init__('FacesetResizer', FacesetResizerSubprocessor.Cli, 600)
@ -30,7 +32,7 @@ class FacesetResizerSubprocessor(Subprocessor):
#override #override
def process_info_generator(self): def process_info_generator(self):
base_dict = {'output_dirpath':self.output_dirpath, 'image_size':self.image_size} base_dict = {'output_dirpath':self.output_dirpath, 'image_size':self.image_size, 'face_type':self.face_type}
for device_idx in range( min(8, multiprocessing.cpu_count()) ): for device_idx in range( min(8, multiprocessing.cpu_count()) ):
client_dict = base_dict.copy() client_dict = base_dict.copy()
@ -63,6 +65,7 @@ class FacesetResizerSubprocessor(Subprocessor):
def on_initialize(self, client_dict): def on_initialize(self, client_dict):
self.output_dirpath = client_dict['output_dirpath'] self.output_dirpath = client_dict['output_dirpath']
self.image_size = client_dict['image_size'] self.image_size = client_dict['image_size']
self.face_type = client_dict['face_type']
self.log_info (f"Running on { client_dict['device_name'] }") self.log_info (f"Running on { client_dict['device_name'] }")
#override #override
@ -72,39 +75,87 @@ class FacesetResizerSubprocessor(Subprocessor):
if dflimg is None or not dflimg.has_data(): if dflimg is None or not dflimg.has_data():
self.log_err (f"{filepath.name} is not a dfl image file") self.log_err (f"{filepath.name} is not a dfl image file")
else: else:
dfl_dict = dflimg.get_dict()
img = cv2_imread(filepath) img = cv2_imread(filepath)
h,w = img.shape[:2] h,w = img.shape[:2]
if h != w: if h != w:
raise Exception(f'w != h in {filepath}') raise Exception(f'w != h in {filepath}')
image_size = self.image_size image_size = self.image_size
scale = w / image_size face_type = self.face_type
img = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_LANCZOS4)
output_filepath = self.output_dirpath / filepath.name output_filepath = self.output_dirpath / filepath.name
cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )
if face_type is not None:
lmrks = dflimg.get_landmarks()
mat = LandmarksProcessor.get_transform_mat(lmrks, image_size, face_type)
img = cv2.warpAffine(img, mat, (image_size, image_size), flags=cv2.INTER_LANCZOS4 )
img = np.clip(img, 0, 255).astype(np.uint8)
cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )
dflimg = DFLIMG.load (output_filepath) dfl_dict = dflimg.get_dict()
dflimg.set_dict(dfl_dict) dflimg = DFLIMG.load (output_filepath)
dflimg.set_dict(dfl_dict)
lmrks = dflimg.get_landmarks() xseg_mask = dflimg.get_xseg_mask()
lmrks /= scale if xseg_mask is not None:
dflimg.set_landmarks(lmrks) xseg_res = 256
seg_ie_polys = dflimg.get_seg_ie_polys() xseg_lmrks = lmrks.copy()
seg_ie_polys.mult_points( 1.0 / scale) xseg_lmrks *= (xseg_res / w)
dflimg.set_seg_ie_polys(seg_ie_polys) xseg_mat = LandmarksProcessor.get_transform_mat(xseg_lmrks, xseg_res, face_type)
mat = dflimg.get_image_to_face_mat() xseg_mask = cv2.warpAffine(xseg_mask, xseg_mat, (xseg_res, xseg_res), flags=cv2.INTER_LANCZOS4 )
if mat is not None: xseg_mask[xseg_mask < 0.5] = 0
face_type = FaceType.fromString ( dflimg.get_face_type() ) xseg_mask[xseg_mask >= 0.5] = 1
mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type )
dflimg.set_image_to_face_mat(mat) dflimg.set_xseg_mask(xseg_mask)
dflimg.save()
seg_ie_polys = dflimg.get_seg_ie_polys()
for poly in seg_ie_polys.get_polys():
poly_pts = poly.get_pts()
poly_pts = LandmarksProcessor.transform_points(poly_pts, mat)
poly.set_points(poly_pts)
dflimg.set_seg_ie_polys(seg_ie_polys)
lmrks = LandmarksProcessor.transform_points(lmrks, mat)
dflimg.set_landmarks(lmrks)
image_to_face_mat = dflimg.get_image_to_face_mat()
if image_to_face_mat is not None:
image_to_face_mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type )
dflimg.set_image_to_face_mat(image_to_face_mat)
dflimg.save()
else:
dfl_dict = dflimg.get_dict()
scale = w / image_size
img = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_LANCZOS4)
cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] )
dflimg = DFLIMG.load (output_filepath)
dflimg.set_dict(dfl_dict)
lmrks = dflimg.get_landmarks()
lmrks /= scale
dflimg.set_landmarks(lmrks)
seg_ie_polys = dflimg.get_seg_ie_polys()
seg_ie_polys.mult_points( 1.0 / scale)
dflimg.set_seg_ie_polys(seg_ie_polys)
image_to_face_mat = dflimg.get_image_to_face_mat()
if image_to_face_mat is not None:
face_type = FaceType.fromString ( dflimg.get_face_type() )
image_to_face_mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type )
dflimg.set_image_to_face_mat(image_to_face_mat)
dflimg.save()
return (1, filepath, output_filepath) return (1, filepath, output_filepath)
except: except:
@ -115,7 +166,17 @@ class FacesetResizerSubprocessor(Subprocessor):
def process_folder ( dirpath): def process_folder ( dirpath):
image_size = io.input_int(f"New image size", 512, valid_range=[256,2048]) image_size = io.input_int(f"New image size", 512, valid_range=[256,2048])
face_type = io.input_str ("Change face type", 'same', ['h','mf','f','wf','head','same']).lower()
if face_type == 'same':
face_type = None
else:
face_type = {'h' : FaceType.HALF,
'mf' : FaceType.MID_FULL,
'f' : FaceType.FULL,
'wf' : FaceType.WHOLE_FACE,
'head' : FaceType.HEAD}[face_type]
output_dirpath = dirpath.parent / (dirpath.name + '_resized') output_dirpath = dirpath.parent / (dirpath.name + '_resized')
output_dirpath.mkdir (exist_ok=True, parents=True) output_dirpath.mkdir (exist_ok=True, parents=True)
@ -131,7 +192,7 @@ def process_folder ( dirpath):
Path(filename).unlink() Path(filename).unlink()
image_paths = [Path(x) for x in pathex.get_image_paths( dirpath )] image_paths = [Path(x) for x in pathex.get_image_paths( dirpath )]
result = FacesetResizerSubprocessor ( image_paths, output_dirpath, image_size).run() result = FacesetResizerSubprocessor ( image_paths, output_dirpath, image_size, face_type).run()
is_merge = io.input_bool (f"\r\nMerge {output_dirpath_parts} to {dirpath_parts} ?", True) is_merge = io.input_bool (f"\r\nMerge {output_dirpath_parts} to {dirpath_parts} ?", True)
if is_merge: if is_merge: