mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-06 04:52:13 -07:00
FacesetResizer now supports changing face type
This commit is contained in:
iperov
parent
d204e049d1
commit
aa26089032
1 changed files with 91 additions and 30 deletions
|
|
@ -1,21 +1,23 @@
|
|||
import multiprocessing
|
||||
import shutil
|
||||
|
||||
import cv2
|
||||
from core import pathex
|
||||
from core.cv2ex import *
|
||||
from core.interact import interact as io
|
||||
from core.joblib import Subprocessor
|
||||
from DFLIMG import *
|
||||
from facelib import LandmarksProcessor, FaceType
|
||||
from facelib import FaceType, LandmarksProcessor
|
||||
|
||||
|
||||
class FacesetResizerSubprocessor(Subprocessor):
|
||||
|
||||
#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.output_dirpath = output_dirpath
|
||||
self.image_size = image_size
|
||||
self.face_type = face_type
|
||||
self.result = []
|
||||
|
||||
super().__init__('FacesetResizer', FacesetResizerSubprocessor.Cli, 600)
|
||||
|
|
@ -30,7 +32,7 @@ class FacesetResizerSubprocessor(Subprocessor):
|
|||
|
||||
#override
|
||||
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()) ):
|
||||
client_dict = base_dict.copy()
|
||||
|
|
@ -63,6 +65,7 @@ class FacesetResizerSubprocessor(Subprocessor):
|
|||
def on_initialize(self, client_dict):
|
||||
self.output_dirpath = client_dict['output_dirpath']
|
||||
self.image_size = client_dict['image_size']
|
||||
self.face_type = client_dict['face_type']
|
||||
self.log_info (f"Running on { client_dict['device_name'] }")
|
||||
|
||||
#override
|
||||
|
|
@ -72,39 +75,87 @@ class FacesetResizerSubprocessor(Subprocessor):
|
|||
if dflimg is None or not dflimg.has_data():
|
||||
self.log_err (f"{filepath.name} is not a dfl image file")
|
||||
else:
|
||||
dfl_dict = dflimg.get_dict()
|
||||
|
||||
img = cv2_imread(filepath)
|
||||
h,w = img.shape[:2]
|
||||
if h != w:
|
||||
raise Exception(f'w != h in {filepath}')
|
||||
|
||||
image_size = self.image_size
|
||||
scale = w / image_size
|
||||
|
||||
img = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_LANCZOS4)
|
||||
|
||||
face_type = self.face_type
|
||||
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)
|
||||
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)
|
||||
|
||||
mat = dflimg.get_image_to_face_mat()
|
||||
if mat is not None:
|
||||
face_type = FaceType.fromString ( dflimg.get_face_type() )
|
||||
mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type )
|
||||
dflimg.set_image_to_face_mat(mat)
|
||||
dflimg.save()
|
||||
dfl_dict = dflimg.get_dict()
|
||||
dflimg = DFLIMG.load (output_filepath)
|
||||
dflimg.set_dict(dfl_dict)
|
||||
|
||||
xseg_mask = dflimg.get_xseg_mask()
|
||||
if xseg_mask is not None:
|
||||
xseg_res = 256
|
||||
|
||||
xseg_lmrks = lmrks.copy()
|
||||
xseg_lmrks *= (xseg_res / w)
|
||||
xseg_mat = LandmarksProcessor.get_transform_mat(xseg_lmrks, xseg_res, face_type)
|
||||
|
||||
xseg_mask = cv2.warpAffine(xseg_mask, xseg_mat, (xseg_res, xseg_res), flags=cv2.INTER_LANCZOS4 )
|
||||
xseg_mask[xseg_mask < 0.5] = 0
|
||||
xseg_mask[xseg_mask >= 0.5] = 1
|
||||
|
||||
dflimg.set_xseg_mask(xseg_mask)
|
||||
|
||||
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)
|
||||
except:
|
||||
|
|
@ -115,7 +166,17 @@ class FacesetResizerSubprocessor(Subprocessor):
|
|||
def process_folder ( dirpath):
|
||||
|
||||
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.mkdir (exist_ok=True, parents=True)
|
||||
|
|
@ -131,7 +192,7 @@ def process_folder ( dirpath):
|
|||
Path(filename).unlink()
|
||||
|
||||
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)
|
||||
if is_merge:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue