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

Sorter: improved sort by blur, and sort by best faces

Browse source
This commit is contained in:
Colombo 2020-06-23 21:02:38 +04:00
parent 669935c9e6
commit 29b1050637
2 changed files with 11 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

9
core/imagelib/estimate_sharpness.py
View file

@ -269,9 +269,10 @@ def get_block_contrast(block):
def estimate_sharpness(image):
height, width = image.shape[:2]
if image.ndim == 3:
image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
if image.shape[2] > 1:
image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
else:
image = image[...,0]
return compute(image)

7
mainscripts/Sorter.py
View file

@ -33,6 +33,10 @@ class BlurEstimatorSubprocessor(Subprocessor):
return [ str(filepath), 0 ]
else:
image = cv2_imread( str(filepath) )
face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks())
image = (image*face_mask).astype(np.uint8)
return [ str(filepath), estimate_sharpness(image) ]
@ -448,7 +452,8 @@ class FinalLoaderSubprocessor(Subprocessor):
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)
face_mask = LandmarksProcessor.get_image_hull_mask (gray.shape, dflimg.get_landmarks())
sharpness = estimate_sharpness( (gray*face_mask).astype(np.uint8) )
pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll ( dflimg.get_landmarks(), size=dflimg.get_shape()[1] )