diff --git a/main.py b/main.py index fb64f96..2ba085f 100644 --- a/main.py +++ b/main.py @@ -70,7 +70,7 @@ if __name__ == "__main__": 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('--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) def process_util(arguments): diff --git a/mainscripts/Sorter.py b/mainscripts/Sorter.py index 9b09e43..39eec5e 100644 --- a/mainscripts/Sorter.py +++ b/mainscripts/Sorter.py @@ -23,6 +23,9 @@ from facelib import LandmarksProcessor class BlurEstimatorSubprocessor(Subprocessor): class Cli(Subprocessor.Cli): + def on_initialize(self, client_dict): + self.estimate_motion_blur = client_dict['estimate_motion_blur'] + #override def process_data(self, data): filepath = Path( data[0] ) @@ -36,8 +39,14 @@ class BlurEstimatorSubprocessor(Subprocessor): face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks()) 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 @@ -46,8 +55,9 @@ class BlurEstimatorSubprocessor(Subprocessor): return data[0] #override - def __init__(self, input_data ): + def __init__(self, input_data, estimate_motion_blur=False ): self.input_data = input_data + self.estimate_motion_blur = estimate_motion_blur self.img_list = [] self.trash_img_list = [] super().__init__('BlurEstimator', BlurEstimatorSubprocessor.Cli, 60) @@ -66,7 +76,7 @@ class BlurEstimatorSubprocessor(Subprocessor): io.log_info(f'Running on {cpu_count} CPUs') for i in range(cpu_count): - yield 'CPU%d' % (i), {}, {} + yield 'CPU%d' % (i), {}, {'estimate_motion_blur':self.estimate_motion_blur} #override 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) 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): io.log_info ("Sorting by face yaw...") img_list = [] @@ -876,6 +897,7 @@ def final_process(input_path, img_list, trash_img_list): sort_func_methods = { 'blur': ("blur", sort_by_blur), + 'motion-blur': ("motion_blur", sort_by_motion_blur), 'face-yaw': ("face yaw direction", sort_by_face_yaw), '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), @@ -903,7 +925,7 @@ def main (input_path, sort_by_method=None): io.log_info(f"[{i}] {desc}") 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] else: