create_data.py

#Filename:create_data.py



import os
import cv2
import numpy as np
import mediapipe as mp
from utils import get_face_landmarks

# Get the absolute path to the directory of the current script
script_dir = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(script_dir, 'archive', 'train')

if not os.path.isdir(DATA_DIR):
    print(f"Error: The directory '{DATA_DIR}' does not exist.")
    exit()

emotions = [emotion for emotion in os.listdir(DATA_DIR) if os.path.isdir(os.path.join(DATA_DIR, emotion))]
labels = [i for i in range(len(emotions))]
label_dict = {emotion: label for emotion, label in zip(emotions, labels)}

print(f"Found emotions: {list(label_dict.keys())}")

# Initialize FaceMesh model
face_mesh = mp.solutions.face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1, min_detection_confidence=0.5)

data = []
for emotion, label in label_dict.items():
    emotion_dir = os.path.join(DATA_DIR, emotion)
    image_paths = [os.path.join(emotion_dir, img) for img in os.listdir(emotion_dir)]
    print(f"Processing {len(image_paths)} images for emotion: {emotion}")
    for img_path in image_paths:
        image = cv2.imread(img_path)
        if image is not None:
            landmarks = get_face_landmarks(image, face_mesh)
            if landmarks:
                data.append([label] + landmarks)
        else:
            print(f"Warning: Could not read image {img_path}")

output_file = os.path.join(script_dir, 'data.txt')
np.savetxt(output_file, data)
print(f"'{output_file}' created successfully with {len(data)} entries.")

face_mesh.close()