From 8e9b4540ac6e76fe91236936a04706271a3b8e13 Mon Sep 17 00:00:00 2001 From: "J\\Jan" Date: Fri, 16 Jul 2021 20:43:55 +0200 Subject: [PATCH] DFLPNG, dfl_dict in custom chunk --- DFLIMG/DFLIMG.py | 5 +++- mainscripts/Extractor.py | 23 +++++++++++++---- mainscripts/FacesetResizer.py | 48 ++++++++++++++++++++++++++++++----- 3 files changed, 63 insertions(+), 13 deletions(-) diff --git a/DFLIMG/DFLIMG.py b/DFLIMG/DFLIMG.py index e213920..bd2b6df 100644 --- a/DFLIMG/DFLIMG.py +++ b/DFLIMG/DFLIMG.py @@ -1,12 +1,15 @@ from pathlib import Path from .DFLJPG import DFLJPG +from .DFLPNG import DFLPNG class DFLIMG(): @staticmethod def load(filepath, loader_func=None): - if filepath.suffix == '.jpg': + if filepath.suffix == '.jpg' : return DFLJPG.load ( str(filepath), loader_func=loader_func ) + elif filepath.suffix == '.png' : + return DFLPNG.load ( str(filepath), loader_func=loader_func ) else: return None diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index 365804f..5ac9c56 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -44,6 +44,7 @@ class ExtractSubprocessor(Subprocessor): def on_initialize(self, client_dict): self.type = client_dict['type'] self.image_size = client_dict['image_size'] + self.image_type = client_dict['image_type'] self.jpeg_quality = client_dict['jpeg_quality'] self.face_type = client_dict['face_type'] self.max_faces_from_image = client_dict['max_faces_from_image'] @@ -116,6 +117,7 @@ class ExtractSubprocessor(Subprocessor): image=image, face_type=self.face_type, image_size=self.image_size, + image_type=self.image_type, jpeg_quality=self.jpeg_quality, output_debug_path=self.output_debug_path, final_output_path=self.final_output_path, @@ -200,6 +202,7 @@ class ExtractSubprocessor(Subprocessor): image, face_type, image_size, + image_type, jpeg_quality, output_debug_path=None, final_output_path=None, @@ -244,10 +247,10 @@ class ExtractSubprocessor(Subprocessor): if data.force_output_path is not None: output_path = data.force_output_path - output_filepath = output_path / f"{filepath.stem}_{face_idx}.jpg" + output_filepath = output_path / f"{filepath.stem}_{face_idx}.{image_type}" cv2_imwrite(output_filepath, face_image, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality ] ) - dflimg = DFLJPG.load(output_filepath) + dflimg = DFLIMG.load(output_filepath) dflimg.set_face_type(FaceType.toString(face_type)) dflimg.set_landmarks(face_image_landmarks.tolist()) dflimg.set_source_filename(filepath.name) @@ -304,7 +307,7 @@ class ExtractSubprocessor(Subprocessor): elif type == 'final': return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in (range(min(8, multiprocessing.cpu_count())) if not DEBUG else [0]) ] - def __init__(self, input_data, type, image_size=None, jpeg_quality=None, face_type=None, output_debug_path=None, manual_window_size=0, max_faces_from_image=0, final_output_path=None, device_config=None): + def __init__(self, input_data, type, image_size=None, image_type=None, jpeg_quality=None, face_type=None, output_debug_path=None, manual_window_size=0, max_faces_from_image=0, final_output_path=None, device_config=None): if type == 'landmarks-manual': for x in input_data: x.manual = True @@ -313,6 +316,7 @@ class ExtractSubprocessor(Subprocessor): self.type = type self.image_size = image_size + self.image_type = image_type self.jpeg_quality = jpeg_quality self.face_type = face_type self.output_debug_path = output_debug_path @@ -364,6 +368,7 @@ class ExtractSubprocessor(Subprocessor): def process_info_generator(self): base_dict = {'type' : self.type, 'image_size': self.image_size, + 'image_type': self.image_type, 'jpeg_quality' : self.jpeg_quality, 'face_type': self.face_type, 'max_faces_from_image':self.max_faces_from_image, @@ -714,6 +719,7 @@ def main(detector=None, face_type='full_face', max_faces_from_image=None, image_size=None, + image_type=None, jpeg_quality=None, cpu_only = False, force_gpu_idxs = None, @@ -772,9 +778,15 @@ def main(detector=None, if image_size is None: image_size = io.input_int(f"Image size", 512 if face_type < FaceType.HEAD else 768, valid_range=[256,2048], help_message="Output image size. The higher image size, the worse face-enhancer works. Use higher than 512 value only if the source image is sharp enough and the face does not need to be enhanced.") + if image_type is None: + image_type = io.input_str(f"image type", "jpg", ['jpg','png'], help_message="image type, jpg faster extraction (factor 10x), but lossy compression") + if jpeg_quality is None: - jpeg_quality = io.input_int(f"Jpeg quality", 90, valid_range=[1,100], help_message="Jpeg quality. The higher jpeg quality the larger the output file size.") - + if image_type == 'jpg': + jpeg_quality = io.input_int(f"Jpeg quality", 90, valid_range=[1,100], help_message="Jpeg quality. The higher jpeg quality the larger the output file size.") + else: + jpeg_quality = 100 + if detector is None: io.log_info ("Choose detector type.") io.log_info ("[0] S3FD") @@ -819,6 +831,7 @@ def main(detector=None, data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_image_paths ], 'all', image_size, + image_type, jpeg_quality, face_type, output_debug_path if output_debug else None, diff --git a/mainscripts/FacesetResizer.py b/mainscripts/FacesetResizer.py index 4bcd1b8..a39496b 100644 --- a/mainscripts/FacesetResizer.py +++ b/mainscripts/FacesetResizer.py @@ -13,11 +13,13 @@ from facelib import FaceType, LandmarksProcessor class FacesetResizerSubprocessor(Subprocessor): #override - def __init__(self, image_paths, output_dirpath, image_size, face_type=None): + def __init__(self, image_paths, output_dirpath, image_size, face_type=None, dfl_file_type=None, jpg_quality=None): self.image_paths = image_paths self.output_dirpath = output_dirpath self.image_size = image_size self.face_type = face_type + self.dfl_file_type = dfl_file_type + self.jpg_quality = jpg_quality self.result = [] super().__init__('FacesetResizer', FacesetResizerSubprocessor.Cli, 600) @@ -32,7 +34,7 @@ class FacesetResizerSubprocessor(Subprocessor): #override def process_info_generator(self): - base_dict = {'output_dirpath':self.output_dirpath, 'image_size':self.image_size, 'face_type':self.face_type} + base_dict = {'output_dirpath':self.output_dirpath, 'image_size':self.image_size, 'face_type':self.face_type, 'dfl_file_type':self.dfl_file_type, 'jpg_quality':self.jpg_quality } for device_idx in range( min(8, multiprocessing.cpu_count()) ): client_dict = base_dict.copy() @@ -66,6 +68,8 @@ class FacesetResizerSubprocessor(Subprocessor): self.output_dirpath = client_dict['output_dirpath'] self.image_size = client_dict['image_size'] self.face_type = client_dict['face_type'] + self.dfl_file_type = client_dict['dfl_file_type'] + self.jpg_quality = client_dict['jpg_quality'] self.log_info (f"Running on { client_dict['device_name'] }") #override @@ -82,7 +86,9 @@ class FacesetResizerSubprocessor(Subprocessor): image_size = self.image_size face_type = self.face_type + dfl_file_type = self.dfl_file_type output_filepath = self.output_dirpath / filepath.name + file_suffix = filepath.suffix[1:] if face_type is not None: lmrks = dflimg.get_landmarks() @@ -91,7 +97,18 @@ class FacesetResizerSubprocessor(Subprocessor): 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] ) + if dfl_file_type is None: + dfl_file_type = file_suffix + + if dfl_file_type == 'jpg': + #get basename + output_filepath = self.output_dirpath / (filepath.stem + '.jpg') + cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), self.jpg_quality] ) + elif dfl_file_type == 'png': + output_filepath = self.output_dirpath / (filepath.stem + '.png') + cv2_imwrite ( str(output_filepath), img ) + + #cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] ) dfl_dict = dflimg.get_dict() dflimg = DFLIMG.load (output_filepath) @@ -137,7 +154,18 @@ class FacesetResizerSubprocessor(Subprocessor): img = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_LANCZOS4) - cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] ) + if dfl_file_type is None: + dfl_file_type = file_suffix + + if dfl_file_type == 'jpg': + output_filepath = self.output_dirpath / (filepath.stem + '.jpg') + cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), self.jpg_quality] ) + elif dfl_file_type == 'png': + output_filepath = self.output_dirpath / (filepath.stem + '.png') + cv2_imwrite ( str(output_filepath), img ) + + + #cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] ) dflimg = DFLIMG.load (output_filepath) dflimg.set_dict(dfl_dict) @@ -177,8 +205,14 @@ def process_folder ( dirpath): 'f' : FaceType.FULL, 'wf' : FaceType.WHOLE_FACE, 'head' : FaceType.HEAD}[face_type] - - + + dfl_file_type = io.input_str ("Change dfl filetype", 'same', ['same', 'jpg', 'png']).lower() + if dfl_file_type == 'same': + dfl_file_type = None + jpg_quality = None + elif dfl_file_type == 'jpg': + jpg_quality = io.input_int(f"jpeg quality", 90, valid_range=[10,100]) + output_dirpath = dirpath.parent / (dirpath.name + '_resized') output_dirpath.mkdir (exist_ok=True, parents=True) @@ -193,7 +227,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, face_type).run() + result = FacesetResizerSubprocessor ( image_paths, output_dirpath, image_size, face_type, dfl_file_type, jpg_quality).run() is_merge = io.input_bool (f"\r\nMerge {output_dirpath_parts} to {dirpath_parts} ?", True) if is_merge: