diff --git a/facelib/LandmarksExtractor.py b/facelib/LandmarksExtractor.py index 8a8a38c..9f5d6bd 100644 --- a/facelib/LandmarksExtractor.py +++ b/facelib/LandmarksExtractor.py @@ -28,20 +28,22 @@ class LandmarksExtractor(object): landmarks = [] for (left, top, right, bottom) in rects: + try: + center = np.array( [ (left + right) / 2.0, (top + bottom) / 2.0] ) + center[1] -= (bottom - top) * 0.12 + scale = (right - left + bottom - top) / 195.0 - center = np.array( [ (left + right) / 2.0, (top + bottom) / 2.0] ) - center[1] -= (bottom - top) * 0.12 - scale = (right - left + bottom - top) / 195.0 - - image = self.crop(input_image, center, scale).astype(np.float32) - image = np.expand_dims(image, 0) - - predicted = self.keras_model.predict (image).transpose (0,3,1,2) - - pts_img = self.get_pts_from_predict ( predicted[-1], center, scale) - pts_img = [ ( int(pt[0]), int(pt[1]) ) for pt in pts_img ] - landmarks.append ( ( (left, top, right, bottom),pts_img ) ) - + image = self.crop(input_image, center, scale).astype(np.float32) + image = np.expand_dims(image, 0) + + predicted = self.keras_model.predict (image).transpose (0,3,1,2) + + pts_img = self.get_pts_from_predict ( predicted[-1], center, scale) + pts_img = [ ( int(pt[0]), int(pt[1]) ) for pt in pts_img ] + landmarks.append ( ( (left, top, right, bottom),pts_img ) ) + except: + landmarks.append ( ( (left, top, right, bottom), [ (0,0) for _ in range(68) ] ) ) + return landmarks def transform(self, point, center, scale, resolution): @@ -58,6 +60,7 @@ class LandmarksExtractor(object): def crop(self, image, center, scale, resolution=256.0): ul = self.transform([1, 1], center, scale, resolution).astype( np.int ) br = self.transform([resolution, resolution], center, scale, resolution).astype( np.int ) + if image.ndim > 2: newDim = np.array([br[1] - ul[1], br[0] - ul[0], image.shape[2]], dtype=np.int32) newImg = np.zeros(newDim, dtype=np.uint8) @@ -71,6 +74,7 @@ class LandmarksExtractor(object): oldX = np.array([max(1, ul[0] + 1), min(br[0], wd)], dtype=np.int32) oldY = np.array([max(1, ul[1] + 1), min(br[1], ht)], dtype=np.int32) newImg[newY[0] - 1:newY[1], newX[0] - 1:newX[1] ] = image[oldY[0] - 1:oldY[1], oldX[0] - 1:oldX[1], :] + newImg = cv2.resize(newImg, dsize=(int(resolution), int(resolution)), interpolation=cv2.INTER_LINEAR) return newImg