--- a +++ b/internal/labeling_script.py @@ -0,0 +1,57 @@ +import cv2 +import pandas as pd +import numpy as np +dataset = pd.read_csv('assets/extracted_landmarks_test.csv') + +def is_float(num): + try: + float(num) + return True + except ValueError: + return False + +dataset['landmarks'] = dataset['landmarks'].apply(lambda arr: np.array([float(n) for n in arr.split() if is_float(n)])) + +video_dir = 'assets/dataset_videos/batch_test' # Directory where your video files are stored + +dataset['Label'] = None +for index, row in dataset.iterrows(): + video = row['video'] + frame = row['frame'] + + # Construct the path to the video file + video_path = video_dir + video # Update the file extension if necessary + + # Open the video file + cap = cv2.VideoCapture(video_path) + + # Get the frames per second (FPS) of the video + fps = cap.get(cv2.CAP_PROP_FPS) + + print(f"{video} at group {row['group']}. Frame {frame}") + # Read and discard frames until the desired frame index + cap.set(cv2.CAP_PROP_POS_FRAMES, frame) + + # Display ~10 frames + for _ in range(10): + ret, frame = cap.read() + if ret: + cv2.imshow('Video', frame) + + # Ask for user input + key = cv2.waitKey(0) & 0xFF + if key == ord('g'): + dataset.at[index, 'Label'] = 'good' # Update 'Label' to 'good' in the corresponding row + print("Label set to 'good'") + elif key == ord('b'): + dataset.at[index, 'Label'] = 'bad' # Update 'Label' to 'bad' in the corresponding row + print("Label set to 'bad'") + + # Release the video capture + cap.release() + + # Close the video display window + cv2.destroyAllWindows() + +dataset.to_csv('assets/labeled_dataset_test.csv', index=False) +print(dataset.head())