Spaces:
Running
Running
| import os | |
| import cv2 | |
| import dlib | |
| import mediapipe as mp | |
| def read_image(image_path): | |
| """ | |
| Read an image from the given path and convert it to RGB. | |
| """ | |
| image = cv2.imread(image_path) | |
| if image is None: | |
| raise FileNotFoundError(f"Image not found at path: {image_path}") | |
| return cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| class LandmarksDetector: | |
| def __init__(self, predictor_model_path=f'{os.path.dirname(os.path.abspath(__file__))}/utils/shape_predictor_68_face_landmarks.dat'): | |
| """ | |
| :param predictor_model_path: path to shape_predictor_68_face_landmarks.dat file | |
| """ | |
| self.detector = dlib.get_frontal_face_detector() # cnn_face_detection_model_v1 also can be used | |
| self.shape_predictor = dlib.shape_predictor(predictor_model_path) | |
| def get_landmarks(self, image): | |
| img = dlib.load_rgb_image(image) | |
| dets = self.detector(img, 1) | |
| dets = [dets[0]] | |
| for detection in dets: | |
| try: | |
| face_landmarks = [(item.x, item.y) for item in self.shape_predictor(img, detection).parts()] | |
| yield face_landmarks | |
| except: | |
| print("Exception in get_landmarks()!") | |
| class DlibLandmarkDetector: | |
| def __init__(self, predictor_model_path=f'{os.path.dirname(os.path.abspath(__file__))}/utils/shape_predictor_68_face_landmarks.dat'): | |
| """ | |
| :param predictor_model_path: path to shape_predictor_68_face_landmarks.dat file | |
| """ | |
| self.detector = dlib.get_frontal_face_detector() # cnn_face_detection_model_v1 also can be used | |
| self.predictor = dlib.shape_predictor(predictor_model_path) | |
| def get_landmarks(self, image_path): | |
| # image = dlib.load_rgb_image(image_path) | |
| image = read_image(image_path) | |
| height, width, _ = image.shape | |
| # Detect the faces in the image | |
| dets = self.detector(image, 1) # 1 indicates to upsample the image 1 time. Higher values may give better results | |
| # Raise an exception if no face is detected | |
| if len(dets) == 0: | |
| raise Exception("No face detected in the image at path: ", image_path) | |
| # Get the landmarks of the first face detected | |
| face_landmarks = [(item.x, item.y) for item in self.predictor(image, dets[0]).parts()] | |
| # Add corner and edge midpoints as landmarks to include the background | |
| corner_landmarks = [(1, 1), (1, height - 1), (width - 1, 1), (width - 1, height - 1)] | |
| edge_landmarks = [(1, (height - 1)//2), ((width - 1)//2, 1), ((width - 1)//2, height - 1), (width - 1, (height - 1)//2)] | |
| # Concatenate the landmarks | |
| landmarks = face_landmarks + corner_landmarks + edge_landmarks | |
| return landmarks, image | |
| def show_landmarked_image(self, image_path, landmarks): | |
| image = read_image(image_path) | |
| for landmark in landmarks: | |
| x, y = landmark | |
| cv2.circle(image, (x, y), 1, (255, 255, 0), -1) # image, (x, y), radius, color, thickness (-1 to fill) | |
| cv2.imshow('image', image) | |
| cv2.waitKey(0) | |
| cv2.destroyAllWindows() | |
| class MediaPipeLandmarkDetector: | |
| def __init__(self): | |
| self.face_mesh = mp.solutions.face_mesh.FaceMesh( | |
| static_image_mode=True, | |
| max_num_faces=1, | |
| min_detection_confidence=0.5) | |
| def get_landmarks(self, image_path): | |
| image = read_image(image_path) | |
| height, width, _ = image.shape | |
| image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| # Process the image | |
| results = self.face_mesh.process(image_rgb) | |
| # Raise an exception if no face is detected | |
| if results.multi_face_landmarks is None: | |
| raise Exception("No face detected in the image at path: ", image_path) | |
| # Extract the face landmarks | |
| face_landmarks = results.multi_face_landmarks[0] | |
| face_landmarks_normalized = [[landmark.x , landmark.y] for landmark in face_landmarks.landmark] | |
| # Add corner and edge midpoints as landmarks to include the background | |
| corner_landmarks = [(0, 0), (0, 1), (1, 0), (1, 1)] | |
| edge_landmarks = [(0, 0.5), (0.5, 0), (0.5, 1), (1, 0.5)] | |
| # Concatenate the corner and edge landmarks | |
| landmarks = corner_landmarks + edge_landmarks + face_landmarks_normalized | |
| # Multiply the landmarks with the image dimensions | |
| landmarks = [(int(x * width) - 1, int(y * height) - 1) for x, y in landmarks] | |
| landmarks = [(max(1, x), max(1, y)) for x, y in landmarks] | |
| return landmarks, image | |
| def show_landmarked_image(self, image_path, landmarks): | |
| image = cv2.imread(image_path) | |
| for landmark in landmarks: | |
| x, y = landmark | |
| cv2.circle(image, (x, y), 1, (255, 255, 0), -1) | |
| cv2.imshow('image', image) | |
| cv2.waitKey(0) | |
| cv2.destroyAllWindows() |