StarMirror/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2024-03-22 13:10:55 +08:00
Code Issues Projects Releases Wiki Activity

Sorter: added sort by motion blur

Browse Source
This commit is contained in:
iperov 2020-12-13 09:49:32 +04:00
parent 27f41b3bcc
commit b21bce458f
2 changed files with 28 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
main.py
View File

@ -70,7 +70,7 @@ if __name__ == "__main__":
p = subparsers.add_parser( "sort", help="Sort faces in a directory.") p = subparsers.add_parser( "sort", help="Sort faces in a directory.")
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
p.add_argument('--by', dest="sort_by_method", default=None, choices=("blur", "face-yaw", "face-pitch", "face-source-rect-size", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final", "final-faster", "absdiff"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." ) p.add_argument('--by', dest="sort_by_method", default=None, choices=("blur", "motion-blur", "face-yaw", "face-pitch", "face-source-rect-size", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final-by-blur", "final-by-size", "absdiff"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
p.set_defaults (func=process_sort) p.set_defaults (func=process_sort)
def process_util(arguments): def process_util(arguments):

32
mainscripts/Sorter.py
View File

@ -23,6 +23,9 @@ from facelib import LandmarksProcessor
class BlurEstimatorSubprocessor(Subprocessor): class BlurEstimatorSubprocessor(Subprocessor):
class Cli(Subprocessor.Cli): class Cli(Subprocessor.Cli):
def on_initialize(self, client_dict):
self.estimate_motion_blur = client_dict['estimate_motion_blur']
#override #override
def process_data(self, data): def process_data(self, data):
filepath = Path( data[0] ) filepath = Path( data[0] )
@ -36,8 +39,14 @@ class BlurEstimatorSubprocessor(Subprocessor):
face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks()) face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks())
image = (image*face_mask).astype(np.uint8) image = (image*face_mask).astype(np.uint8)
return [ str(filepath), estimate_sharpness(image) ]
if self.estimate_motion_blur:
value = cv2.Laplacian(image, cv2.CV_64F, ksize=11).var()
else:
value = estimate_sharpness(image)
return [ str(filepath), value ]
#override #override
@ -46,8 +55,9 @@ class BlurEstimatorSubprocessor(Subprocessor):
return data[0] return data[0]
#override #override
def __init__(self, input_data ): def __init__(self, input_data, estimate_motion_blur=False ):
self.input_data = input_data self.input_data = input_data
self.estimate_motion_blur = estimate_motion_blur
self.img_list = [] self.img_list = []
self.trash_img_list = [] self.trash_img_list = []
super().__init__('BlurEstimator', BlurEstimatorSubprocessor.Cli, 60) super().__init__('BlurEstimator', BlurEstimatorSubprocessor.Cli, 60)
@ -66,7 +76,7 @@ class BlurEstimatorSubprocessor(Subprocessor):
io.log_info(f'Running on {cpu_count} CPUs') io.log_info(f'Running on {cpu_count} CPUs')
for i in range(cpu_count): for i in range(cpu_count):
yield 'CPU%d' % (i), {}, {} yield 'CPU%d' % (i), {}, {'estimate_motion_blur':self.estimate_motion_blur}
#override #override
def get_data(self, host_dict): def get_data(self, host_dict):
@ -103,7 +113,18 @@ def sort_by_blur(input_path):
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True) img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
return img_list, trash_img_list return img_list, trash_img_list
def sort_by_motion_blur(input_path):
io.log_info ("Sorting by motion blur...")
img_list = [ (filename,[]) for filename in pathex.get_image_paths(input_path) ]
img_list, trash_img_list = BlurEstimatorSubprocessor (img_list, estimate_motion_blur=True).run()
io.log_info ("Sorting...")
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
return img_list, trash_img_list
def sort_by_face_yaw(input_path): def sort_by_face_yaw(input_path):
io.log_info ("Sorting by face yaw...") io.log_info ("Sorting by face yaw...")
img_list = [] img_list = []
@ -876,6 +897,7 @@ def final_process(input_path, img_list, trash_img_list):
sort_func_methods = { sort_func_methods = {
'blur': ("blur", sort_by_blur), 'blur': ("blur", sort_by_blur),
'motion-blur': ("motion_blur", sort_by_motion_blur),
'face-yaw': ("face yaw direction", sort_by_face_yaw), 'face-yaw': ("face yaw direction", sort_by_face_yaw),
'face-pitch': ("face pitch direction", sort_by_face_pitch), 'face-pitch': ("face pitch direction", sort_by_face_pitch),
'face-source-rect-size' : ("face rect size in source image", sort_by_face_source_rect_size), 'face-source-rect-size' : ("face rect size in source image", sort_by_face_source_rect_size),
@ -903,7 +925,7 @@ def main (input_path, sort_by_method=None):
io.log_info(f"[{i}] {desc}") io.log_info(f"[{i}] {desc}")
io.log_info("") io.log_info("")
id = io.input_int("", 4, valid_list=[*range(len(key_list))] ) id = io.input_int("", 5, valid_list=[*range(len(key_list))] )
sort_by_method = key_list[id] sort_by_method = key_list[id]
else: else: