github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-08 05:51:40 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

added sort by "best faces faster"

Browse source 
same as sort by "best faces"
but faces will be sorted by source-rect-area instead of blur, x40 times faster.
This commit is contained in:
Colombo 2020-03-06 21:38:41 +04:00
parent 54548afe1a
commit 9ccdd271a4
2 changed files with 33 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

2
main.py
View file

@ -112,7 +112,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", "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", "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.set_defaults (func=process_sort)
def process_util(arguments):

30
mainscripts/Sorter.py
View file

@ -425,7 +425,7 @@ class FinalLoaderSubprocessor(Subprocessor):
class Cli(Subprocessor.Cli):
#override
def on_initialize(self, client_dict):
self.include_by_blur = client_dict['include_by_blur']
self.faster = client_dict['faster']
#override
def process_data(self, data):
@ -443,7 +443,13 @@ class FinalLoaderSubprocessor(Subprocessor):
raise Exception ("Unable to load %s" % (filepath.name) )
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
sharpness = estimate_sharpness(gray) if self.include_by_blur else 0
if self.faster:
source_rect = dflimg.get_source_rect()
sharpness = mathlib.polygon_area(np.array(source_rect[[0,2,2,0]]).astype(np.float32), np.array(source_rect[[1,1,3,3]]).astype(np.float32))
else:
sharpness = estimate_sharpness(gray)
pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll ( dflimg.get_landmarks(), size=dflimg.get_shape()[1] )
hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
@ -459,10 +465,10 @@ class FinalLoaderSubprocessor(Subprocessor):
return data[0]
#override
def __init__(self, img_list, include_by_blur ):
def __init__(self, img_list, faster ):
self.img_list = img_list
self.include_by_blur = include_by_blur
self.faster = faster
self.result = []
self.result_trash = []
@ -482,7 +488,7 @@ class FinalLoaderSubprocessor(Subprocessor):
io.log_info(f'Running on {cpu_count} CPUs')
for i in range(cpu_count):
yield 'CPU%d' % (i), {}, {'include_by_blur': self.include_by_blur}
yield 'CPU%d' % (i), {}, {'faster': self.faster}
#override
def get_data(self, host_dict):
@ -579,12 +585,17 @@ class FinalHistDissimSubprocessor(Subprocessor):
def get_result(self):
return self.result
def sort_best(input_path, include_by_blur=True):
io.log_info ("Performing sort by best faces.")
def sort_best_faster(input_path):
return sort_best(input_path, faster=True)
def sort_best(input_path, faster=False):
target_count = io.input_int ("Target number of faces?", 2000)
img_list, trash_img_list = FinalLoaderSubprocessor( pathex.get_image_paths(input_path), include_by_blur ).run()
io.log_info ("Performing sort by best faces.")
if faster:
io.log_info("Using faster algorithm. Faces will be sorted by source-rect-area instead of blur.")
img_list, trash_img_list = FinalLoaderSubprocessor( pathex.get_image_paths(input_path), faster ).run()
final_img_list = []
grads = 128
@ -618,7 +629,7 @@ def sort_best(input_path, include_by_blur=True):
imgs_per_grad += total_lack // grads
if include_by_blur:
sharpned_imgs_per_grad = imgs_per_grad*10
for g in io.progress_bar_generator ( range (grads), "Sort by blur"):
img_list = yaws_sample_list[g]
@ -873,6 +884,7 @@ sort_func_methods = {
'oneface': ("one face in image", sort_by_oneface_in_image),
'absdiff': ("absolute pixel difference", sort_by_absdiff),
'final': ("best faces", sort_best),
'final-fast': ("best faces faster", sort_best_faster),
}
def main (input_path, sort_by_method=None):