diff --git a/facelib/LandmarksExtractor.py b/facelib/LandmarksExtractor.py index 810c81a..b61298c 100644 --- a/facelib/LandmarksExtractor.py +++ b/facelib/LandmarksExtractor.py @@ -3,7 +3,8 @@ import numpy as np import os import cv2 from pathlib import Path - +from facelib import FaceType +from facelib import LandmarksProcessor class LandmarksExtractor(object): def __init__ (self, keras): @@ -23,7 +24,7 @@ class LandmarksExtractor(object): del self.keras_model return False #pass exception between __enter__ and __exit__ to outter level - def extract_from_bgr (self, input_image, rects): + def extract_from_bgr (self, input_image, rects, second_pass_extractor=None): input_image = input_image[:,:,::-1].copy() (h, w, ch) = input_image.shape @@ -44,10 +45,29 @@ class LandmarksExtractor(object): landmarks.append ( ( (left, top, right, bottom),pts_img ) ) except Exception as e: print ("extract_from_bgr: ", traceback.format_exc() ) - landmarks.append ( ( (left, top, right, bottom), [ (0,0) for _ in range(68) ] ) ) + landmarks.append ( ( (left, top, right, bottom), None ) ) + + if second_pass_extractor is not None: + for i in range(len(landmarks)): + rect, lmrks = landmarks[i] + if lmrks is None: + continue + + image_to_face_mat = LandmarksProcessor.get_transform_mat (lmrks, 256, FaceType.FULL) + face_image = cv2.warpAffine(input_image, image_to_face_mat, (256, 256), cv2.INTER_CUBIC) + + rects2 = second_pass_extractor.extract_from_bgr(face_image) + if len(rects2) != 1: #dont do second pass if more than 1 face detected in cropped image + continue + + rect2 = rects2[0] + + lmrks2 = self.extract_from_bgr (face_image, [rect2] )[0][1] + source_lmrks2 = LandmarksProcessor.transform_points (lmrks2, image_to_face_mat, True) + landmarks[i] = (rect, source_lmrks2) return landmarks - + def transform(self, point, center, scale, resolution): pt = np.array ( [point[0], point[1], 1.0] ) h = 200.0 * scale diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index e9b78bd..b21216f 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -52,7 +52,7 @@ class ExtractSubprocessor(Subprocessor): nnlib.import_all (device_config) self.e = facelib.S3FDExtractor() else: - raise ValueError ("Wrond detector type.") + raise ValueError ("Wrong detector type.") if self.e is not None: self.e.__enter__() @@ -61,6 +61,11 @@ class ExtractSubprocessor(Subprocessor): nnlib.import_all (device_config) self.e = facelib.LandmarksExtractor(nnlib.keras) self.e.__enter__() + if device_config.gpu_vram_gb[0] >= 2: + self.second_pass_e = facelib.S3FDExtractor() + self.second_pass_e.__enter__() + else: + self.second_pass_e = None elif self.type == 'final': pass @@ -76,7 +81,7 @@ class ExtractSubprocessor(Subprocessor): filename_path_str = str(filename_path) if self.cached_image[0] == filename_path_str: - image = self.cached_image[1] + image = self.cached_image[1] #cached image for manual extractor else: image = cv2_imread( filename_path_str ) @@ -102,8 +107,14 @@ class ExtractSubprocessor(Subprocessor): image = image[0:h-hm,0:w-wm,:] self.cached_image = ( filename_path_str, image ) + src_dflimg = None + h, w, ch = image.shape + if h == w: + #extracting from already extracted jpg image? + if filename_path.suffix == '.jpg': + src_dflimg = DFLJPG.load ( str(filename_path) ) + if self.type == 'rects': - h, w, ch = image.shape if min(w,h) < 128: self.log_err ( 'Image is too small %s : [%d, %d]' % ( str(filename_path), w, h ) ) rects = [] @@ -116,18 +127,13 @@ class ExtractSubprocessor(Subprocessor): rects = data[1] if rects is None: landmarks = None - else: - landmarks = self.e.extract_from_bgr (image, rects) + else: + landmarks = self.e.extract_from_bgr (image, rects, self.second_pass_e if src_dflimg is None else None) return [str(filename_path), landmarks] elif self.type == 'final': - src_dflimg = None - (h,w,c) = image.shape - if h == w: - #extracting from already extracted jpg image? - if filename_path.suffix == '.jpg': - src_dflimg = DFLJPG.load ( str(filename_path) ) + result = [] faces = data[1] @@ -139,7 +145,10 @@ class ExtractSubprocessor(Subprocessor): face_idx = 0 for face in faces: rect = np.array(face[0]) - image_landmarks = np.array(face[1]) + image_landmarks = face[1] + if image_landmarks is None: + continue + image_landmarks = np.array(image_landmarks) if self.face_type == FaceType.MARK_ONLY: face_image = image diff --git a/nnlib/device.py b/nnlib/device.py index 67f1901..a0d9e4f 100644 --- a/nnlib/device.py +++ b/nnlib/device.py @@ -64,6 +64,7 @@ class device: self.cpu_only = (len(self.gpu_idxs) == 0) + if not self.cpu_only: self.gpu_names = [] self.gpu_compute_caps = [] @@ -73,7 +74,11 @@ class device: self.gpu_compute_caps += [ device.getDeviceComputeCapability(gpu_idx) ] self.gpu_vram_gb += [ device.getDeviceVRAMTotalGb(gpu_idx) ] self.cpu_only = (len(self.gpu_idxs) == 0) - + else: + self.gpu_names = ['CPU'] + self.gpu_compute_caps = [99] + self.gpu_vram_gb = [0] + if self.cpu_only: self.backend = "tensorflow-cpu"