mpkpts_extract_keypoints.py

"""
This file is used to extract the keypoints from the videos in the dataset by leveraging the KeypointsExtractor class.
"""

from pathlib import Path
import argparse

import numpy as np

from src.keypointsextractor import KeypointsExtractor
from src import constants

if __name__ == "__main__":
    
    parser = argparse.ArgumentParser(description='Subsample videos.')
    parser.add_argument('--input', type=str, required=True, help='Path to input directory.')
    parser.add_argument('--output', type=str, required=True, help='Path to output directory.')
    args = parser.parse_args()

    # Define paths
    path_input = Path(args.input)
    path_output = Path(args.output)

    # Create output folder if it doesn't exist
    path_output.mkdir(parents=True, exist_ok=True)

    # Loop through each label folder
    for target in ["0", "1"]:

        # Create the target directory if it doesn't exist
        (path_output / target).mkdir(parents=True, exist_ok=True)

        # Loop through each video in the input folder
        for input_file in sorted((path_input / target).glob("*")):

            # Skip if the input is not a file or if it is not a video
            if not (input_file.is_file() and input_file.suffix.lower() in constants.VIDEO_EXTENSIONS):
                continue
            
            # Define the output file
            output_file = (path_output / input_file.relative_to(path_input)).with_suffix(".npy")

            # If the keypoints file already exists, skip it
            if output_file.is_file():
                print(f"Skipping {input_file}")
                continue

            # Print the current file
            print(f"Processing {input_file}")
            
            # Extract keypoints from the video
            feature_extractor = KeypointsExtractor(str(input_file), show_image=True)
            keypoints = feature_extractor.extract_keypoints_from_video()

            # Save keypoints as npy file if keypoints is not None
            if keypoints is not None:
                np.save(output_file, keypoints)
            else:
                print("No hand detected for the whole video, or multiple hands detected.")

    print("Done extracting keypoints.")