github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-06 04:52:13 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

enhanced sort by final: now faces are evenly distributed not only in the direction of yaw, but also in pitch

Browse source
This commit is contained in:
Colombo 2019-10-20 14:18:14 +04:00
parent e63e89c305
commit d3e35c8ce6
1 changed files with 83 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

109
mainscripts/Sorter.py
View file

@ -233,7 +233,7 @@ def sort_by_face_pitch(input_path):
io.log_err ("%s is not a dfl image file" % (filepath.name) )
trash_img_list.append ( [str(filepath)] )
continue
pitch_yaw_roll = dflimg.get_pitch_yaw_roll()
if pitch_yaw_roll is not None:
pitch, yaw, roll = pitch_yaw_roll
@ -544,7 +544,7 @@ class FinalLoaderSubprocessor(Subprocessor):
self.log_err (e)
return [ 1, [str(filepath)] ]
return [ 0, [str(filepath), sharpness, hist, yaw ] ]
return [ 0, [str(filepath), sharpness, hist, yaw, pitch ] ]
#override
def get_data_name (self, data):
@ -608,28 +608,35 @@ class FinalHistDissimSubprocessor(Subprocessor):
#override
def process_data(self, data):
idx, img_list = data
for i in range( len(img_list) ):
score_total = 0
for j in range( len(img_list) ):
if i == j:
continue
score_total += cv2.compareHist(img_list[i][2], img_list[j][2], cv2.HISTCMP_BHATTACHARYYA)
img_list[i][3] = score_total
img_list = sorted(img_list, key=operator.itemgetter(3), reverse=True)
return idx, img_list
idx, pitch_yaw_img_list = data
for p in range ( len(pitch_yaw_img_list) ):
img_list = pitch_yaw_img_list[p]
if img_list is not None:
for i in range( len(img_list) ):
score_total = 0
for j in range( len(img_list) ):
if i == j:
continue
score_total += cv2.compareHist(img_list[i][2], img_list[j][2], cv2.HISTCMP_BHATTACHARYYA)
img_list[i][3] = score_total
pitch_yaw_img_list[p] = sorted(img_list, key=operator.itemgetter(3), reverse=True)
return idx, pitch_yaw_img_list
#override
def get_data_name (self, data):
return "Bunch of images"
#override
def __init__(self, yaws_sample_list ):
self.yaws_sample_list = yaws_sample_list
self.yaws_sample_list_len = len(yaws_sample_list)
def __init__(self, pitch_yaw_sample_list ):
self.pitch_yaw_sample_list = pitch_yaw_sample_list
self.pitch_yaw_sample_list_len = len(pitch_yaw_sample_list)
self.yaws_sample_list_idxs = [ i for i in range(self.yaws_sample_list_len) if self.yaws_sample_list[i] is not None ]
self.result = [ None for _ in range(self.yaws_sample_list_len) ]
self.pitch_yaw_sample_list_idxs = [ i for i in range(self.pitch_yaw_sample_list_len) if self.pitch_yaw_sample_list[i] is not None ]
self.result = [ None for _ in range(self.pitch_yaw_sample_list_len) ]
super().__init__('FinalHistDissimSubprocessor', FinalHistDissimSubprocessor.Cli)
#override
@ -640,7 +647,7 @@ class FinalHistDissimSubprocessor(Subprocessor):
}
#override
def on_clients_initialized(self):
io.progress_bar ("Sort by hist-dissim", self.yaws_sample_list_len)
io.progress_bar ("Sort by hist-dissim", len(self.pitch_yaw_sample_list_idxs) )
#override
def on_clients_finalized(self):
@ -648,15 +655,15 @@ class FinalHistDissimSubprocessor(Subprocessor):
#override
def get_data(self, host_dict):
if len (self.yaws_sample_list_idxs) > 0:
idx = self.yaws_sample_list_idxs.pop(0)
if len (self.pitch_yaw_sample_list_idxs) > 0:
idx = self.pitch_yaw_sample_list_idxs.pop(0)
return idx, self.yaws_sample_list[idx]
return idx, self.pitch_yaw_sample_list[idx]
return None
#override
def on_data_return (self, host_dict, data):
self.yaws_sample_list_idxs.insert(0, data[0])
self.pitch_yaw_sample_list_idxs.insert(0, data[0])
#override
def on_result (self, host_dict, data, result):
@ -721,15 +728,65 @@ def sort_final(input_path, include_by_blur=True):
yaws_sample_list[g] = img_list
yaws_sample_list = FinalHistDissimSubprocessor(yaws_sample_list).run()
for g in io.progress_bar_generator ( range (grads), "Fetching best"):
yaw_pitch_sample_list = [None]*grads
pitch_grads = imgs_per_grad
for g in io.progress_bar_generator ( range (grads), "Sort by pitch"):
img_list = yaws_sample_list[g]
if img_list is None:
continue
final_img_list += img_list[0:imgs_per_grad]
trash_img_list += img_list[imgs_per_grad:]
pitch_sample_list = [None]*pitch_grads
grads_space = np.linspace (-1.0,1.0, pitch_grads )
for pg in range (pitch_grads):
pitch = grads_space[pg]
next_pitch = grads_space[pg+1] if pg < pitch_grads-1 else pitch
pitch_samples = []
for img in img_list:
s_pitch = img[4]
if (pg == 0 and s_pitch < next_pitch) or \
(pg < pitch_grads-1 and s_pitch >= pitch and s_pitch < next_pitch) or \
(pg == pitch_grads-1 and s_pitch >= pitch):
pitch_samples += [ img ]
if len(pitch_samples) > 0:
pitch_sample_list[pg] = pitch_samples
yaw_pitch_sample_list[g] = pitch_sample_list
yaw_pitch_sample_list = FinalHistDissimSubprocessor(yaw_pitch_sample_list).run()
for g in io.progress_bar_generator (range (grads), "Fetching the best"):
pitch_sample_list = yaw_pitch_sample_list[g]
if pitch_sample_list is None:
continue
n = imgs_per_grad
while n > 0:
n_prev = n
for pg in range(pitch_grads):
img_list = pitch_sample_list[pg]
if img_list is None:
continue
final_img_list += [ img_list.pop(0) ]
if len(img_list) == 0:
pitch_sample_list[pg] = None
n -= 1
if n == 0:
break
if n_prev == n:
break
for pg in range(pitch_grads):
img_list = pitch_sample_list[pg]
if img_list is None:
continue
trash_img_list += img_list
return final_img_list, trash_img_list