create_datase.py

import mediapipe as mp
import cv2
import os
import matplotlib.pyplot as plt
import pickle

mp_hands = mp.solutions.hands
mp_drawing = mp.solutions.drawing_utils
mp_drawing_styles = mp.solutions.drawing_styles

hands = mp_hands.Hands(static_image_mode=False, min_detection_confidence=0.3)

DATA_DIR = './data'
data = []
labels = [] 

if not os.path.exists(DATA_DIR):
    print(f"Data directory not existis at {DATA_DIR}")
    
for dir in os.listdir(DATA_DIR):
     for frame_path in os.listdir(os.path.join(DATA_DIR, dir)):
        data_aux = []
        img = cv2.imread(os.path.join(DATA_DIR, dir, frame_path))
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        result = hands.process(img_rgb)
        
        if not result.multi_hand_landmarks:
            print(f"No hands detected in {frame_path}")
            continue
        
        for hand_landmarks in result.multi_hand_landmarks:
            # mp_drawing.draw_landmarks(
            #     img_rgb,
            #     hand_landmarks,
            #     mp_hands.HAND_CONNECTIONS,
            #     mp_drawing_styles.get_default_hand_landmarks_style(),
            #     mp_drawing_styles.get_default_hand_connections_style())
            
            for i in range(len(hand_landmarks.landmark)):
               x = hand_landmarks.landmark[i].x
               y = hand_landmarks.landmark[i].y 
               z =hand_landmarks.landmark[i].z
               data_aux.append(x)
               data_aux.append(y)
        
        data.append(data_aux)
        labels.append(dir)
        
        # plt.figure()
        # plt.imshow(img_rgb)
        
# plt.show()

f = open('data.pickle', 'wb')
pickle.dump({'data': data, 'labels': labels}, f)
f.close()