make_data.py

""" Create dataset """
import os
import uuid

import cv2 as cv

DATA_DIR = "data"
DIRECTIONS = ["up", "down", "right", "left", "neutral"]


def prepare_data_dir():
    for direction in DIRECTIONS:
        direction = os.path.join(DATA_DIR, direction)
        if not os.path.exists(direction):
            os.makedirs(direction)


def capture_hand():
    prepare_data_dir()
    cap = cv.VideoCapture(0)
    while True:
        ret, frame = cap.read()

        if not ret:
            continue

        flipped = cv.flip(frame, 1)
        cv.imshow("frame", flipped)

        key = cv.waitKey(1) & 0xFF
        filepath = ""
        fileid = uuid.uuid4().urn[9:]
        # it the user press "u", the image will be saved in data/up/uuid.jpg
        # same for "d", "r","l", and "n" for neutral
        # q to quit
        if key == ord("u"):
            filepath = os.path.join(DATA_DIR, "up", fileid + ".jpg")
        elif key == ord("d"):
            filepath = os.path.join(DATA_DIR, "down", fileid + ".jpg")
        elif key == ord("r"):
            filepath = os.path.join(DATA_DIR, "right", fileid + ".jpg")
        elif key == ord("l"):
            filepath = os.path.join(DATA_DIR, "left", fileid + ".jpg")
        elif key == ord("n"):
            filepath = os.path.join(DATA_DIR, "neutral", fileid + ".jpg")
        elif key == ord("q"):
            break

        if filepath:
            cv.imwrite(filepath, frame)
            print(f"Saved image to {filepath}")

    cap.release()
    cv.destroyAllWindows()


if __name__ == "__main__":
    capture_hand()