Dev
This commit is contained in:
@@ -10,10 +10,10 @@ import pandas as pd
|
||||
|
||||
|
||||
class KeypointExtractor:
|
||||
def __init__(self, video_folder: str, cache_folder: str = "cache"):
|
||||
def __init__(self, cache_folder: str = "cache"):
|
||||
self.mp_drawing = mp.solutions.drawing_utils
|
||||
self.mp_holistic = mp.solutions.holistic
|
||||
self.video_folder = video_folder
|
||||
# self.video_folder = video_folder
|
||||
self.cache_folder = cache_folder
|
||||
|
||||
# we will store the keypoints of each frame as a row in the dataframe. The columns are the keypoints: Pose (33), Left Hand (21), Right Hand (21). Each keypoint has 3 values: x, y
|
||||
@@ -40,10 +40,12 @@ class KeypointExtractor:
|
||||
:rtype: pd.DataFrame
|
||||
"""
|
||||
|
||||
video_name = video.split("/")[-1].split(".")[0]
|
||||
|
||||
if not draw:
|
||||
# check if video exists
|
||||
if not os.path.exists(self.video_folder + video):
|
||||
logging.error("Video does not exist at path: " + self.video_folder + video)
|
||||
if not os.path.exists(video):
|
||||
logging.error("Video does not exist at path: " + video)
|
||||
return None
|
||||
|
||||
# check if cache exists
|
||||
@@ -51,22 +53,22 @@ class KeypointExtractor:
|
||||
os.makedirs(self.cache_folder)
|
||||
|
||||
# check if cache file exists and return
|
||||
if os.path.exists(self.cache_folder + "/" + video + ".npy"):
|
||||
if os.path.exists(self.cache_folder + "/" + video_name + ".npy"):
|
||||
# create dataframe from cache
|
||||
df = pd.DataFrame(np.load(self.cache_folder + "/" + video + ".npy", allow_pickle=True), columns=self.columns)
|
||||
df = pd.DataFrame(np.load(self.cache_folder + "/" + video_name + ".npy", allow_pickle=True), columns=self.columns)
|
||||
if normalize:
|
||||
df = self.normalize_hands(df, norm_algorithm=normalize)
|
||||
df = self.normalize_pose_bohacek(df)
|
||||
df, _ = self.normalize_pose_bohacek(df)
|
||||
return df
|
||||
|
||||
# open video
|
||||
cap = cv2.VideoCapture(self.video_folder + video)
|
||||
cap = cv2.VideoCapture(video)
|
||||
|
||||
keypoints_df = pd.DataFrame(columns=self.columns)
|
||||
|
||||
# extract frames from video so we extract 5 frames per second
|
||||
frame_rate = int(cap.get(cv2.CAP_PROP_FPS))
|
||||
frame_skip = frame_rate // 10
|
||||
frame_skip = (frame_rate // 10) -1
|
||||
|
||||
output_frames = []
|
||||
|
||||
@@ -113,12 +115,12 @@ class KeypointExtractor:
|
||||
cap.release()
|
||||
|
||||
# save keypoints to cache
|
||||
np.save(self.cache_folder + "/" + video + ".npy", keypoints_df.to_numpy())
|
||||
np.save(self.cache_folder + "/" + video_name + ".npy", keypoints_df.to_numpy())
|
||||
|
||||
# normalize hands and pose keypoints
|
||||
if normalize:
|
||||
keypoints_df = self.normalize_hands(keypoints_df, norm_algorithm=normalize)
|
||||
keypoints_df = self.normalize_pose_bohacek(keypoints_df)
|
||||
keypoints_df, _ = self.normalize_pose_bohacek(keypoints_df)
|
||||
|
||||
if draw:
|
||||
return keypoints_df, output_frames
|
||||
@@ -179,28 +181,28 @@ class KeypointExtractor:
|
||||
|
||||
if norm_algorithm == "minmax":
|
||||
# normalize left hand
|
||||
dataframe = self.normalize_hand_minmax(dataframe, "left_hand")
|
||||
dataframe, _= self.normalize_hand_minmax(dataframe, "left_hand")
|
||||
# normalize right hand
|
||||
dataframe = self.normalize_hand_minmax(dataframe, "right_hand")
|
||||
dataframe, _= self.normalize_hand_minmax(dataframe, "right_hand")
|
||||
elif norm_algorithm == "bohacek":
|
||||
# normalize left hand
|
||||
dataframe = self.normalize_hand_bohacek(dataframe, "left_hand")
|
||||
dataframe, _= self.normalize_hand_bohacek(dataframe, "left_hand")
|
||||
# normalize right hand
|
||||
dataframe = self.normalize_hand_bohacek(dataframe, "right_hand")
|
||||
dataframe, _= self.normalize_hand_bohacek(dataframe, "right_hand")
|
||||
else:
|
||||
return dataframe
|
||||
|
||||
return dataframe
|
||||
|
||||
def normalize_hand_minmax(self, dataframe: pd.DataFrame, hand: str) -> pd.DataFrame:
|
||||
"""normalize_hand_minmax this function normalizes the hand keypoints of a dataframe with respect to the minimum and maximum coordinates
|
||||
def normalize_hand_minmax(self, dataframe: pd.DataFrame, hand: str) -> Tuple[pd.DataFrame, pd.DataFrame]:
|
||||
"""normalize_hand_helper this function normalizes the hand keypoints of a dataframe with respect to the minimum and maximum coordinates
|
||||
|
||||
:param dataframe: the dataframe to normalize
|
||||
:type dataframe: pd.DataFrame
|
||||
:param hand: the hand to normalize
|
||||
:type hand: str
|
||||
:return: the normalized dataframe
|
||||
:rtype: pd.DataFrame
|
||||
:return: the normalized dataframe and the bounding boxes dataframe
|
||||
:rtype: Tuple[pd.DataFrame, pd.DataFrame]
|
||||
"""
|
||||
# get all columns that belong to the hand (left hand column 66 - 107, right hand column 108 - 149)
|
||||
hand_columns = np.array([i for i in range(66 + (42 if hand == "right_hand" else 0), 108 + (42 if hand == "right_hand" else 0))])
|
||||
@@ -226,24 +228,28 @@ class KeypointExtractor:
|
||||
bbox_dims = np.concatenate((np.tile(bbox_width, (1, 21, 1)), np.tile(bbox_height, (1, 21, 1))), axis=2)
|
||||
|
||||
if np.any(bbox_dims == 0):
|
||||
return dataframe
|
||||
return dataframe, None
|
||||
# normalize the hand keypoints based on the bounding box around the hand
|
||||
norm_hand_coords = (hand_coords - center_coords) / bbox_dims
|
||||
|
||||
# flatten the normalized hand keypoints array and replace the original hand keypoints with the normalized hand keypoints in the dataframe
|
||||
dataframe.iloc[:, hand_columns] = norm_hand_coords.reshape(-1, 42)
|
||||
|
||||
# merge starting and ending points of the bounding boxes in a dataframe
|
||||
bbox_array = np.hstack((min_x.reshape(-1, 1), min_y.reshape(-1, 1), max_x.reshape(-1, 1), max_y.reshape(-1, 1)))
|
||||
bbox = pd.DataFrame(bbox_array, columns=['starting_x', 'starting_y', 'ending_x', 'ending_y'])
|
||||
|
||||
return dataframe
|
||||
return dataframe, bbox
|
||||
|
||||
def normalize_hand_bohacek(self, dataframe: pd.DataFrame, hand: str) -> pd.DataFrame:
|
||||
"""normalize_hand_bohacek this function normalizes the hand keypoints of a dataframe using the Bohacek-normalization algorithm
|
||||
def normalize_hand_bohacek(self, dataframe: pd.DataFrame, hand: str) -> Tuple[pd.DataFrame, pd.DataFrame]:
|
||||
"""normalize_hand_helper this function normalizes the hand keypoints of a dataframe using the bohacek normalization algorithm
|
||||
|
||||
:param dataframe: the dataframe to normalize
|
||||
:type dataframe: pd.DataFrame
|
||||
:param hand: the hand to normalize
|
||||
:type hand: str
|
||||
:return: the normalized dataframe
|
||||
:rtype: pd.DataFrame
|
||||
:return: the normalized dataframe and the bounding boxes dataframe
|
||||
:rtype: Tuple[pd.DataFrame, pd.DataFrame]
|
||||
"""
|
||||
# get all columns that belong to the hand (left hand column 66 - 107, right hand column 108 - 149)
|
||||
hand_columns = np.array([i for i in range(66 + (42 if hand == "right_hand" else 0), 108 + (42 if hand == "right_hand" else 0))])
|
||||
@@ -287,22 +293,28 @@ class KeypointExtractor:
|
||||
bbox_dims = np.concatenate((np.tile(bbox_width, (1, 21, 1)), np.tile(bbox_height, (1, 21, 1))), axis=2)
|
||||
|
||||
if np.any(bbox_dims == 0):
|
||||
return dataframe
|
||||
return dataframe, None
|
||||
# normalize the hand keypoints based on the bounding box around the hand
|
||||
norm_hand_coords = (hand_coords - center_coords) / bbox_dims
|
||||
|
||||
# flatten the normalized hand keypoints array and replace the original hand keypoints with the normalized hand keypoints in the dataframe
|
||||
dataframe.iloc[:, hand_columns] = norm_hand_coords.reshape(-1, 42)
|
||||
|
||||
return dataframe
|
||||
|
||||
def normalize_pose_bohacek(self, dataframe: pd.DataFrame) -> pd.DataFrame:
|
||||
# merge starting and ending points of the bounding boxes in a dataframe
|
||||
bbox_array = np.hstack((starting_x.reshape(-1, 1), starting_y.reshape(-1, 1), ending_x.reshape(-1, 1), ending_y.reshape(-1, 1)))
|
||||
bbox = pd.DataFrame(bbox_array, columns=['starting_x', 'starting_y', 'ending_x', 'ending_y'])
|
||||
|
||||
return dataframe, bbox
|
||||
|
||||
def normalize_pose_bohacek(self, dataframe: pd.DataFrame, bbox_size: float = 4) -> Tuple[pd.DataFrame, pd.DataFrame]:
|
||||
"""normalize_pose_bohacek this function normalizes the pose keypoints of a dataframe using the Bohacek-normalization algorithm
|
||||
|
||||
:param dataframe: the dataframe to normalize
|
||||
:type dataframe: pd.DataFrame
|
||||
:return: the normalized dataframe
|
||||
:rtype: pd.DataFrame
|
||||
:param bbox_size: the width and height of the normalization bounding box expressed in head metrics, defaults to 4
|
||||
:type bbox_size: float, optional
|
||||
:return: the normalized dataframe and the bounding boxes dataframe
|
||||
:rtype: Tuple[pd.DataFrame, pd.DataFrame]
|
||||
"""
|
||||
# get the columns that belong to the pose
|
||||
pose_columns = np.array([i for i in range(66)])
|
||||
@@ -311,28 +323,22 @@ class KeypointExtractor:
|
||||
pose_coords = dataframe.iloc[:, pose_columns].values.reshape(-1, 33, 2)
|
||||
|
||||
# check in what frames shoulders are visible
|
||||
left_shoulder_present_mask = pose_coords[:, 11, 0]!=0
|
||||
right_shoulder_present_mask = pose_coords[:, 12, 0]!=0
|
||||
shoulders_present_mask = np.logical_and(left_shoulder_present_mask,right_shoulder_present_mask)
|
||||
left_shoulder_present_mask = pose_coords[:, 11, 0] != 0
|
||||
right_shoulder_present_mask = pose_coords[:, 12, 0] != 0
|
||||
shoulders_present_mask = np.logical_and(left_shoulder_present_mask, right_shoulder_present_mask)
|
||||
|
||||
# calculate shoulder distance
|
||||
left_shoulder, right_shoulder = pose_coords[shoulders_present_mask, 11,], pose_coords[shoulders_present_mask, 12,]
|
||||
left_shoulder, right_shoulder = pose_coords[shoulders_present_mask, 11], pose_coords[shoulders_present_mask, 12]
|
||||
shoulder_distance = ((left_shoulder[:, 0] - right_shoulder[:, 0])**2 + (left_shoulder[:, 1] - right_shoulder[:, 1])**2)**0.5
|
||||
head_metric = shoulder_distance
|
||||
|
||||
# center of shoulders and left eye are necessary to construct bounding box
|
||||
center_shoulders = right_shoulder + (left_shoulder - right_shoulder)/2
|
||||
center_shoulders = right_shoulder + (left_shoulder - right_shoulder) / 2
|
||||
left_eye = pose_coords[shoulders_present_mask, 2]
|
||||
|
||||
# set the starting and ending point of the normalization bounding box
|
||||
starting_x, starting_y = center_shoulders[:, 0] - 2*head_metric, left_eye[:, 1] - 0.5*head_metric
|
||||
ending_x, ending_y = center_shoulders[:, 0] + 2*head_metric, starting_y + 4*head_metric
|
||||
|
||||
# ensure that the starting and ending point of the bounding box are not out of the frame
|
||||
#starting_x = np.clip(starting_x, 0, None)
|
||||
#starting_y = np.clip(starting_y, 0 ,None)
|
||||
#ending_x = np.clip(ending_x, 0, None)
|
||||
#ending_y = np.clip(ending_y, 0 ,None)
|
||||
starting_x, starting_y = center_shoulders[:, 0] - (bbox_size / 2) * head_metric, left_eye[:, 1] - 0.5 * head_metric
|
||||
ending_x, ending_y = center_shoulders[:, 0] + (bbox_size / 2) * head_metric, starting_y + (bbox_size - 0.5) * head_metric
|
||||
|
||||
# calculate the center of the bounding box and the bounding box dimensions
|
||||
bbox_center_x, bbox_center_y = (starting_x + ending_x) / 2, (starting_y + ending_y) / 2
|
||||
@@ -342,15 +348,19 @@ class KeypointExtractor:
|
||||
bbox_center_x, bbox_center_y = bbox_center_x.reshape(-1, 1, 1), bbox_center_y.reshape(-1, 1, 1)
|
||||
center_coords = np.concatenate((np.tile(bbox_center_x, (1, 33, 1)), np.tile(bbox_center_y, (1, 33, 1))), axis=2)
|
||||
|
||||
bbox_width, bbox_height = bbox_width.reshape(-1, 1, 1), bbox_height.reshape(-1, 1 ,1)
|
||||
bbox_width, bbox_height = bbox_width.reshape(-1, 1, 1), bbox_height.reshape(-1, 1, 1)
|
||||
bbox_dims = np.concatenate((np.tile(bbox_width, (1, 33, 1)), np.tile(bbox_height, (1, 33, 1))), axis=2)
|
||||
|
||||
if np.any(bbox_dims == 0):
|
||||
return dataframe
|
||||
return dataframe, None
|
||||
# normalize the pose keypoints based on the bounding box
|
||||
norm_pose_coords= (pose_coords - center_coords) / bbox_dims
|
||||
norm_pose_coords = (pose_coords - center_coords) / bbox_dims
|
||||
|
||||
# flatten the normalized pose keypoints array and replace the original pose keypoints with the normalized pose keypoints in the dataframe
|
||||
dataframe.iloc[shoulders_present_mask, pose_columns] = norm_pose_coords.reshape(-1, 66)
|
||||
|
||||
return dataframe
|
||||
# merge starting and ending points of the bounding boxes in a dataframe
|
||||
bbox_array = np.hstack((starting_x.reshape(-1, 1), starting_y.reshape(-1, 1), ending_x.reshape(-1, 1), ending_y.reshape(-1, 1)))
|
||||
bbox = pd.DataFrame(bbox_array, columns=['starting_x', 'starting_y', 'ending_x', 'ending_y'])
|
||||
|
||||
return dataframe, bbox
|
||||
|
||||
Reference in New Issue
Block a user