diff --git a/main.py b/main.py index 2a36300..ae04f2a 100644 --- a/main.py +++ b/main.py @@ -58,7 +58,7 @@ if __name__ == "__main__": sort_parser = subparsers.add_parser( "sort", help="Sort faces in a directory.") sort_parser.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") - sort_parser.add_argument('--by', required=True, dest="sort_by_method", choices=("blur", "face", "face-dissim", "face-yaw", "hist", "hist-dissim", "hist-blur", "ssim", "brightness", "hue", "origname"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." ) + sort_parser.add_argument('--by', required=True, dest="sort_by_method", choices=("blur", "face", "face-dissim", "face-yaw", "hist", "hist-dissim", "hist-blur", "ssim", "brightness", "hue", "black", "origname"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." ) sort_parser.set_defaults (func=process_sort) def process_train(arguments): diff --git a/mainscripts/Sorter.py b/mainscripts/Sorter.py index db766e0..3c17299 100644 --- a/mainscripts/Sorter.py +++ b/mainscripts/Sorter.py @@ -464,7 +464,20 @@ def sort_by_hist_dissim(input_path): img_list = sorted(img_list, key=operator.itemgetter(4), reverse=True) return img_list + +def sort_by_black(input_path): + print ("Sorting by amount of black pixels...") + img_list = [] + for x in tqdm( Path_utils.get_image_paths(input_path), desc="Loading"): + img = cv2.imread(x) + img_list.append ([x, img[(img == 0)].size ]) + + print ("Sorting...") + img_list = sorted(img_list, key=operator.itemgetter(1), reverse=False) + + return img_list + def final_rename(input_path, img_list): for i in tqdm( range(0,len(img_list)), desc="Renaming" , leave=False): src = Path (img_list[i][0]) @@ -529,6 +542,7 @@ def main (input_path, sort_by_method): elif sort_by_method == 'hist-blur': img_list = sort_by_hist_blur (input_path) elif sort_by_method == 'brightness': img_list = sort_by_brightness (input_path) elif sort_by_method == 'hue': img_list = sort_by_hue (input_path) + elif sort_by_method == 'black': img_list = sort_by_black (input_path) elif sort_by_method == 'origname': img_list = sort_by_origname (input_path) final_rename (input_path, img_list) \ No newline at end of file