github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-15 09:33: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

26
mainscripts/Sorter.py
View file

@ -510,19 +510,27 @@ class FinalLoaderSubprocessor(SubprocessorBase):
#override #override
def onClientProcessData(self, data): def onClientProcessData(self, data):
filepath = Path(data[0]) filepath = Path(data[0])
try:
if filepath.suffix != '.png': if filepath.suffix != '.png':
print ("%s is not a png file required for sort_final" % (filepath.name) ) raise Exception ("%s is not a png file required for sort_final" % (filepath.name) )
return [ 1, [str(filepath)] ]
dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True) dflpng = DFLPNG.load (str(filepath), print_on_no_embedded_data=True)
if dflpng is None: if dflpng is None:
return [ 1, [str(filepath)] ] raise Exception ("")
bgr = cv2.imread(str(filepath)) 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 = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
gray_masked = ( gray * LandmarksProcessor.get_image_hull_mask (bgr, dflpng.get_landmarks() )[:,:,0] ).astype(np.uint8) gray_masked = ( gray * LandmarksProcessor.get_image_hull_mask (bgr, dflpng.get_landmarks() )[:,:,0] ).astype(np.uint8)
sharpness = estimate_sharpness(gray_masked) sharpness = estimate_sharpness(gray_masked)
hist = cv2.calcHist([gray], [0], None, [256], [0, 256]) hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
except Exception as e:
print (e)
return [ 1, [str(filepath)] ]
return [ 0, [str(filepath), sharpness, hist, dflpng.get_yaw_value() ] ] return [ 0, [str(filepath), sharpness, hist, dflpng.get_yaw_value() ] ]
@ -551,7 +559,6 @@ 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)
@ -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: