├── README.md ├── prepare_data.py ├── test_model.py ├── thumbnail.jpg ├── train_model.py └── utils.py /README.md: -------------------------------------------------------------------------------- 1 | # emotion-recognition-python-scikit-learn-mediapipe 2 | 3 |

4 | 5 | Watch the video 6 |
Watch on YouTube: Emotion recognition with Python, Scikit Learn and Mediapipe ! 7 | 8 |

9 | 10 | ## data 11 | 12 | You can download the same dataset I used in this tutorial [here](https://www.patreon.com/posts/97173398). 13 | -------------------------------------------------------------------------------- /prepare_data.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import cv2 4 | import numpy as np 5 | 6 | from utils import get_face_landmarks 7 | 8 | 9 | data_dir = './data' 10 | 11 | output = [] 12 | for emotion_indx, emotion in enumerate(sorted(os.listdir(data_dir))): 13 | for image_path_ in os.listdir(os.path.join(data_dir, emotion)): 14 | image_path = os.path.join(data_dir, emotion, image_path_) 15 | 16 | image = cv2.imread(image_path) 17 | 18 | face_landmarks = get_face_landmarks(image) 19 | 20 | if len(face_landmarks) == 1404: 21 | face_landmarks.append(int(emotion_indx)) 22 | output.append(face_landmarks) 23 | 24 | np.savetxt('data.txt', np.asarray(output)) 25 | -------------------------------------------------------------------------------- /test_model.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | 3 | import cv2 4 | 5 | from utils import get_face_landmarks 6 | 7 | 8 | emotions = ['HAPPY', 'SAD', 'SURPRISED'] 9 | 10 | with open('./model', 'rb') as f: 11 | model = pickle.load(f) 12 | 13 | cap = cv2.VideoCapture(2) 14 | 15 | ret, frame = cap.read() 16 | 17 | while ret: 18 | ret, frame = cap.read() 19 | 20 | face_landmarks = get_face_landmarks(frame, draw=True, static_image_mode=False) 21 | 22 | output = model.predict([face_landmarks]) 23 | 24 | cv2.putText(frame, 25 | emotions[int(output[0])], 26 | (10, frame.shape[0] - 1), 27 | cv2.FONT_HERSHEY_SIMPLEX, 28 | 3, 29 | (0, 255, 0), 30 | 5) 31 | 32 | cv2.imshow('frame', frame) 33 | 34 | cv2.waitKey(25) 35 | 36 | 37 | cap.release() 38 | cv2.destroyAllWindows() -------------------------------------------------------------------------------- /thumbnail.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/computervisioneng/emotion-recognition-python-scikit-learn-mediapipe/b59c6915e4e2614aa42757df3c058cb51f3c9ab9/thumbnail.jpg -------------------------------------------------------------------------------- /train_model.py: -------------------------------------------------------------------------------- 1 | import pickle 2 | 3 | import numpy as np 4 | from sklearn.model_selection import train_test_split 5 | from sklearn.ensemble import RandomForestClassifier 6 | from sklearn.metrics import accuracy_score, confusion_matrix 7 | 8 | # Load data from the text file 9 | data_file = "data.txt" 10 | data = np.loadtxt(data_file) 11 | 12 | # Split data into features (X) and labels (y) 13 | X = data[:, :-1] # Features are all columns except the last one 14 | y = data[:, -1] # Labels are the last column 15 | 16 | # Split the data into training and testing sets 17 | X_train, X_test, y_train, y_test = train_test_split(X, 18 | y, 19 | test_size=0.2, 20 | random_state=42, 21 | shuffle=True, 22 | stratify=y) 23 | 24 | # Initialize the Random Forest Classifier 25 | rf_classifier = RandomForestClassifier() 26 | 27 | # Train the classifier on the training data 28 | rf_classifier.fit(X_train, y_train) 29 | 30 | # Make predictions on the test data 31 | y_pred = rf_classifier.predict(X_test) 32 | 33 | # Evaluate the accuracy of the model 34 | accuracy = accuracy_score(y_test, y_pred) 35 | print(f"Accuracy: {accuracy * 100:.2f}%") 36 | print(confusion_matrix(y_test, y_pred)) 37 | 38 | with open('./model', 'wb') as f: 39 | pickle.dump(rf_classifier, f) 40 | 41 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import mediapipe as mp 3 | 4 | 5 | def get_face_landmarks(image, draw=False, static_image_mode=True): 6 | 7 | # Read the input image 8 | image_input_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 9 | 10 | face_mesh = mp.solutions.face_mesh.FaceMesh(static_image_mode=static_image_mode, 11 | max_num_faces=1, 12 | min_detection_confidence=0.5) 13 | image_rows, image_cols, _ = image.shape 14 | results = face_mesh.process(image_input_rgb) 15 | 16 | image_landmarks = [] 17 | 18 | if results.multi_face_landmarks: 19 | 20 | if draw: 21 | 22 | mp_drawing = mp.solutions.drawing_utils 23 | mp_drawing_styles = mp.solutions.drawing_styles 24 | drawing_spec = mp_drawing.DrawingSpec(thickness=2, circle_radius=1) 25 | 26 | mp_drawing.draw_landmarks( 27 | image=image, 28 | landmark_list=results.multi_face_landmarks[0], 29 | connections=mp.solutions.face_mesh.FACEMESH_CONTOURS, 30 | landmark_drawing_spec=drawing_spec, 31 | connection_drawing_spec=drawing_spec) 32 | 33 | ls_single_face = results.multi_face_landmarks[0].landmark 34 | xs_ = [] 35 | ys_ = [] 36 | zs_ = [] 37 | for idx in ls_single_face: 38 | xs_.append(idx.x) 39 | ys_.append(idx.y) 40 | zs_.append(idx.z) 41 | for j in range(len(xs_)): 42 | image_landmarks.append(xs_[j] - min(xs_)) 43 | image_landmarks.append(ys_[j] - min(ys_)) 44 | image_landmarks.append(zs_[j] - min(zs_)) 45 | 46 | return image_landmarks --------------------------------------------------------------------------------