DFLPNG, dfl_dict in custom chunk

This commit is contained in:
J\Jan 2021-07-16 20:43:55 +02:00
commit 8e9b4540ac
3 changed files with 63 additions and 13 deletions

View file

@ -1,12 +1,15 @@
from pathlib import Path from pathlib import Path
from .DFLJPG import DFLJPG from .DFLJPG import DFLJPG
from .DFLPNG import DFLPNG
class DFLIMG(): class DFLIMG():
@staticmethod @staticmethod
def load(filepath, loader_func=None): def load(filepath, loader_func=None):
if filepath.suffix == '.jpg': if filepath.suffix == '.jpg' :
return DFLJPG.load ( str(filepath), loader_func=loader_func ) return DFLJPG.load ( str(filepath), loader_func=loader_func )
elif filepath.suffix == '.png' :
return DFLPNG.load ( str(filepath), loader_func=loader_func )
else: else:
return None return None

View file

@ -44,6 +44,7 @@ class ExtractSubprocessor(Subprocessor):
def on_initialize(self, client_dict): def on_initialize(self, client_dict):
self.type = client_dict['type'] self.type = client_dict['type']
self.image_size = client_dict['image_size'] self.image_size = client_dict['image_size']
self.image_type = client_dict['image_type']
self.jpeg_quality = client_dict['jpeg_quality'] self.jpeg_quality = client_dict['jpeg_quality']
self.face_type = client_dict['face_type'] self.face_type = client_dict['face_type']
self.max_faces_from_image = client_dict['max_faces_from_image'] self.max_faces_from_image = client_dict['max_faces_from_image']
@ -116,6 +117,7 @@ class ExtractSubprocessor(Subprocessor):
image=image, image=image,
face_type=self.face_type, face_type=self.face_type,
image_size=self.image_size, image_size=self.image_size,
image_type=self.image_type,
jpeg_quality=self.jpeg_quality, jpeg_quality=self.jpeg_quality,
output_debug_path=self.output_debug_path, output_debug_path=self.output_debug_path,
final_output_path=self.final_output_path, final_output_path=self.final_output_path,
@ -200,6 +202,7 @@ class ExtractSubprocessor(Subprocessor):
image, image,
face_type, face_type,
image_size, image_size,
image_type,
jpeg_quality, jpeg_quality,
output_debug_path=None, output_debug_path=None,
final_output_path=None, final_output_path=None,
@ -244,10 +247,10 @@ class ExtractSubprocessor(Subprocessor):
if data.force_output_path is not None: if data.force_output_path is not None:
output_path = data.force_output_path 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 ] ) 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_face_type(FaceType.toString(face_type))
dflimg.set_landmarks(face_image_landmarks.tolist()) dflimg.set_landmarks(face_image_landmarks.tolist())
dflimg.set_source_filename(filepath.name) dflimg.set_source_filename(filepath.name)
@ -304,7 +307,7 @@ class ExtractSubprocessor(Subprocessor):
elif type == 'final': elif type == 'final':
return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in (range(min(8, multiprocessing.cpu_count())) if not DEBUG else [0]) ] 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': if type == 'landmarks-manual':
for x in input_data: for x in input_data:
x.manual = True x.manual = True
@ -313,6 +316,7 @@ class ExtractSubprocessor(Subprocessor):
self.type = type self.type = type
self.image_size = image_size self.image_size = image_size
self.image_type = image_type
self.jpeg_quality = jpeg_quality self.jpeg_quality = jpeg_quality
self.face_type = face_type self.face_type = face_type
self.output_debug_path = output_debug_path self.output_debug_path = output_debug_path
@ -364,6 +368,7 @@ class ExtractSubprocessor(Subprocessor):
def process_info_generator(self): def process_info_generator(self):
base_dict = {'type' : self.type, base_dict = {'type' : self.type,
'image_size': self.image_size, 'image_size': self.image_size,
'image_type': self.image_type,
'jpeg_quality' : self.jpeg_quality, 'jpeg_quality' : self.jpeg_quality,
'face_type': self.face_type, 'face_type': self.face_type,
'max_faces_from_image':self.max_faces_from_image, 'max_faces_from_image':self.max_faces_from_image,
@ -714,6 +719,7 @@ def main(detector=None,
face_type='full_face', face_type='full_face',
max_faces_from_image=None, max_faces_from_image=None,
image_size=None, image_size=None,
image_type=None,
jpeg_quality=None, jpeg_quality=None,
cpu_only = False, cpu_only = False,
force_gpu_idxs = None, force_gpu_idxs = None,
@ -772,8 +778,14 @@ def main(detector=None,
if image_size is 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.") 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: if jpeg_quality is None:
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.") 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: if detector is None:
io.log_info ("Choose detector type.") io.log_info ("Choose detector type.")
@ -819,6 +831,7 @@ def main(detector=None,
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_image_paths ], data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_image_paths ],
'all', 'all',
image_size, image_size,
image_type,
jpeg_quality, jpeg_quality,
face_type, face_type,
output_debug_path if output_debug else None, output_debug_path if output_debug else None,

View file

@ -13,11 +13,13 @@ from facelib import FaceType, LandmarksProcessor
class FacesetResizerSubprocessor(Subprocessor): class FacesetResizerSubprocessor(Subprocessor):
#override #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.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.face_type = face_type
self.dfl_file_type = dfl_file_type
self.jpg_quality = jpg_quality
self.result = [] self.result = []
super().__init__('FacesetResizer', FacesetResizerSubprocessor.Cli, 600) super().__init__('FacesetResizer', FacesetResizerSubprocessor.Cli, 600)
@ -32,7 +34,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, '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()) ): for device_idx in range( min(8, multiprocessing.cpu_count()) ):
client_dict = base_dict.copy() client_dict = base_dict.copy()
@ -66,6 +68,8 @@ class FacesetResizerSubprocessor(Subprocessor):
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.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'] }") self.log_info (f"Running on { client_dict['device_name'] }")
#override #override
@ -82,7 +86,9 @@ class FacesetResizerSubprocessor(Subprocessor):
image_size = self.image_size image_size = self.image_size
face_type = self.face_type face_type = self.face_type
dfl_file_type = self.dfl_file_type
output_filepath = self.output_dirpath / filepath.name output_filepath = self.output_dirpath / filepath.name
file_suffix = filepath.suffix[1:]
if face_type is not None: if face_type is not None:
lmrks = dflimg.get_landmarks() 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 = cv2.warpAffine(img, mat, (image_size, image_size), flags=cv2.INTER_LANCZOS4 )
img = np.clip(img, 0, 255).astype(np.uint8) 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() dfl_dict = dflimg.get_dict()
dflimg = DFLIMG.load (output_filepath) dflimg = DFLIMG.load (output_filepath)
@ -137,7 +154,18 @@ class FacesetResizerSubprocessor(Subprocessor):
img = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_LANCZOS4) 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 = DFLIMG.load (output_filepath)
dflimg.set_dict(dfl_dict) dflimg.set_dict(dfl_dict)
@ -178,6 +206,12 @@ def process_folder ( dirpath):
'wf' : FaceType.WHOLE_FACE, 'wf' : FaceType.WHOLE_FACE,
'head' : FaceType.HEAD}[face_type] '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 = dirpath.parent / (dirpath.name + '_resized')
output_dirpath.mkdir (exist_ok=True, parents=True) output_dirpath.mkdir (exist_ok=True, parents=True)
@ -193,7 +227,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, 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) is_merge = io.input_bool (f"\r\nMerge {output_dirpath_parts} to {dirpath_parts} ?", True)
if is_merge: if is_merge: