diff --git a/mainscripts/FacesetResizer.py b/mainscripts/FacesetResizer.py index f2ec391..2c58e77 100644 --- a/mainscripts/FacesetResizer.py +++ b/mainscripts/FacesetResizer.py @@ -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: