Spaces:
Runtime error
Runtime error
| """ | |
| FacePosition class | |
| Author: HenryAreiza | |
| Date: 08/09/2023 | |
| """ | |
| import os | |
| import cv2 | |
| import pickle | |
| import numpy as np | |
| import mediapipe as mp | |
| class FacePosition: | |
| """ | |
| A class for controlling the cursor based on head movements. | |
| This class provides functionality for detecting a face using | |
| the MediaPipe library and controlling the cursor's movement accordingly. | |
| Attributes: | |
| movement (list): List of classes corresponding to the predicted movement. | |
| images (list): List of images associated to each class | |
| cursor_model: The machine learning model for gesture prediction. | |
| face_detection: The MediaPipe Face Detection component. | |
| """ | |
| def __init__(self): | |
| """ | |
| Initializes the FaceCursorController class. | |
| """ | |
| self.movement = ['Center', 'Up', 'Right/Up', 'Right', 'Right/Down', 'Down', 'Left/Down', 'Left', 'Left/Up'] | |
| self.images = [cv2.imread(os.path.join('media', str(i)+'.png')) for i in range(9)] | |
| # Load the cursor movement model | |
| with open('cursor_movement_model.pkl', 'rb') as f: | |
| self.cursor_model = pickle.load(f) | |
| # Initialize the MediaPipe Face Detection component | |
| self.face_detection = mp.solutions.face_detection.FaceDetection(min_detection_confidence=0.5) | |
| def predict(self, frame): | |
| """ | |
| Move the cursor based on head position. | |
| Args: | |
| reference (list): A list containing reference coordinates and size of the bounding box. | |
| keypoints (list): A list of keypoints representing face landmarks. | |
| Returns: | |
| result (list): The predicted class image and label. | |
| """ | |
| # Perform face detection | |
| results = self.face_detection.process(frame) | |
| # Read the reference and landmarks from the detected face | |
| if results.detections: | |
| for detection in results.detections: | |
| reference = [[detection.location_data.relative_bounding_box.xmin, | |
| detection.location_data.relative_bounding_box.ymin], | |
| [detection.location_data.relative_bounding_box.width, | |
| detection.location_data.relative_bounding_box.height]] | |
| keypoints = [] | |
| for key_point in detection.location_data.relative_keypoints: | |
| keypoints.append([key_point.x, key_point.y]) | |
| break | |
| # Transform the lists into numpy arrays | |
| reference = np.array(reference) | |
| keypoints = np.array(keypoints) | |
| # Remove off-set from keypoints | |
| keypoints = (keypoints - reference[0]) / reference[1] | |
| # Recognize the head position | |
| prediction = self.cursor_model.predict(keypoints.reshape((1, -1)))[0] | |
| return [self.images[prediction], self.movement[prediction]] | |
| else: | |
| return [self.images[0], self.movement[0]] | |