diff --git a/mainscripts/Sorter.py b/mainscripts/Sorter.py index 8ff0ee9..e9332da 100644 --- a/mainscripts/Sorter.py +++ b/mainscripts/Sorter.py @@ -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