frame_based_pose_recognition/get_data.py at main · bolducmanuel/frame_based_pose_recognition · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
import cv2
import mediapipe as mp
import pandas as pd
import os
import numpy as np
import logging
import sys

logger = logging.getLogger(__name__)

def image_processed(file_path, pose):

    # reading the static image
    hand_img = cv2.imread(file_path)

    # Image processing
    # 1. Convert BGR to RGB
    img_rgb = cv2.cvtColor(hand_img, cv2.COLOR_BGR2RGB)

    # 2. Flip the img in Y-axis
    img_flip = cv2.flip(img_rgb, 1)


    # Results
    output = pose.process(img_flip)


    try:
        data = output.pose_landmarks.landmark
        #print(data)

        bad_data = False

        clean = []

        for i,landmark in enumerate(data):

            clean.append(float(landmark.x))
            clean.append(float(landmark.y))


            if (landmark.x > 1.0 or landmark.y > 1.0) or (landmark.x < 0. or landmark.y < 0.) : #ignore bad data

                bad_data = True

                break

        if bad_data:
            return
        else:
            return clean

    except:
        pass

def make_csv(pose):
    """
    Utility function for writing keypoint information in csv format for each frame of the dataset.

    TODO: frame notation of the dataset can be optimized by calling mediapipe pose only once at
    the beginning of the for loop, rather than in an individual function such as in image_processed.
    """

    mypath = 'DATASET'
    file_name = open('dataset.csv', 'a')

    num_classes = len(os.listdir(mypath))

    if num_classes < 2:
        logger.error(f"Not enough gestures trained. Please train at least two gestures.")
        sys.exit(1)


    for each_folder in os.listdir(mypath):

        good_frames = 0
        bad_frames = 0

        label = each_folder

        if '._' in each_folder:
            pass

        else:
            for each_number in os.listdir(mypath + '/' + each_folder):
                if '._' in each_number:
                    pass

                else:


                    file_loc = mypath + '/' + each_folder + '/' + each_number

                    data = image_processed(file_loc, pose)
                    try:
                        for id,i in enumerate(data):

                            file_name.write(str(i))
                            file_name.write(',')

                        file_name.write(label)
                        file_name.write('\n')

                        good_frames += 1

                    except:
                        bad_frames += 1
                        pass
                        # file_name.write('0')
                        # file_name.write(',')

                        # file_name.write('None')
                        # file_name.write('\n')

        print('gesture', label, ': Data was generated for', good_frames, 'frames, and ', bad_frames, 'frames were discarded')

    file_name.close()
    print('Data Created !!!')

if __name__ == "__main__":

    # accessing MediaPipe solutions
    mp_pose = mp.solutions.pose

        # Initialize Hands
    pose = mp_pose.Pose(smooth_landmarks=True)


    make_csv(pose)

    pose.close()