diff --git a/core/imagelib/estimate_sharpness.py b/core/imagelib/estimate_sharpness.py index 8401c91..e4b3e2d 100644 --- a/core/imagelib/estimate_sharpness.py +++ b/core/imagelib/estimate_sharpness.py @@ -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) diff --git a/mainscripts/Sorter.py b/mainscripts/Sorter.py index f04409a..179fc53 100644 --- a/mainscripts/Sorter.py +++ b/mainscripts/Sorter.py @@ -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] )