github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-15 01:23:44 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix sort by final

Browse source
This commit is contained in:
iperov 2019-01-07 17:27:10 +04:00
parent ff3a208e80
commit cf41e1b962
1 changed files with 33 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

48
mainscripts/Sorter.py
View file

@ -509,20 +509,28 @@ class FinalLoaderSubprocessor(SubprocessorBase):
#override #override
def onClientProcessData(self, data): def onClientProcessData(self, data):
filepath = Path(data[0]) filepath = Path(data[0])
if filepath.suffix != '.png':
print ("%s is not a png file required for sort_final" % (filepath.name) ) try:
if filepath.suffix != '.png':
raise Exception ("%s is not a png file required for sort_final" % (filepath.name) )
dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True)
if dflpng is None:
raise Exception ("")
bgr = cv2.imread(str(filepath))
if bgr is None:
raise Exception ("Unable to load %s" % (filepath.name) )
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
gray_masked = ( gray * LandmarksProcessor.get_image_hull_mask (bgr, dflpng.get_landmarks() )[:,:,0] ).astype(np.uint8)
sharpness = estimate_sharpness(gray_masked)
hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
except Exception as e:
print (e)
return [ 1, [str(filepath)] ] return [ 1, [str(filepath)] ]
dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True)
if dflpng is None:
return [ 1, [str(filepath)] ]
bgr = cv2.imread(str(filepath))
gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
gray_masked = ( gray * LandmarksProcessor.get_image_hull_mask (bgr, dflpng.get_landmarks() )[:,:,0] ).astype(np.uint8)
sharpness = estimate_sharpness(gray_masked)
hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
return [ 0, [str(filepath), sharpness, hist, dflpng.get_yaw_value() ] ] return [ 0, [str(filepath), sharpness, hist, dflpng.get_yaw_value() ] ]
@ -551,8 +559,7 @@ def sort_final(input_path):
grads = 128 grads = 128
imgs_per_grad = 15 imgs_per_grad = 15
sharpned_imgs_per_grad = imgs_per_grad*10
grads_space = np.linspace (-255,255,grads) grads_space = np.linspace (-255,255,grads)
yaws_sample_list = [None]*grads yaws_sample_list = [None]*grads
@ -570,6 +577,17 @@ def sort_final(input_path):
if len(yaw_samples) > 0: if len(yaw_samples) > 0:
yaws_sample_list[g] = yaw_samples yaws_sample_list[g] = yaw_samples
total_lack = 0
for g in tqdm ( range (grads), desc="" ):
img_list = yaws_sample_list[g]
img_list_len = len(img_list) if img_list is not None else 0
lack = imgs_per_grad - img_list_len
total_lack += max(lack, 0)
imgs_per_grad += total_lack // grads
sharpned_imgs_per_grad = imgs_per_grad*10
for g in tqdm ( range (grads), desc="Sort by blur" ): for g in tqdm ( range (grads), desc="Sort by blur" ):
img_list = yaws_sample_list[g] img_list = yaws_sample_list[g]
if img_list is None: if img_list is None: