From 0b62603920e750b8d71a18fba8abc3f80a4160be Mon Sep 17 00:00:00 2001 From: Robbe De Waele Date: Fri, 17 Mar 2023 22:39:58 +0000 Subject: [PATCH] Resolve WES-78 "Implement pose normalization" --- src/keypoint_extractor.py | 70 +++++++++++++++++++++++++++++++++--- visualize_data.ipynb | 75 +++++++++++++++++++++++++++++++++++++-- 2 files changed, 138 insertions(+), 7 deletions(-) diff --git a/src/keypoint_extractor.py b/src/keypoint_extractor.py index 23ad12a..ac24154 100644 --- a/src/keypoint_extractor.py +++ b/src/keypoint_extractor.py @@ -56,6 +56,7 @@ class KeypointExtractor: df = pd.DataFrame(np.load(self.cache_folder + "/" + video + ".npy", allow_pickle=True), columns=self.columns) if normalize: df = self.normalize_hands(df, norm_algorithm=normalize) + df = self.normalize_pose_bohacek(df) return df # open video @@ -114,8 +115,10 @@ class KeypointExtractor: # save keypoints to cache np.save(self.cache_folder + "/" + video + ".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) if draw: return keypoints_df, output_frames @@ -164,7 +167,7 @@ class KeypointExtractor: def normalize_hands(self, dataframe: pd.DataFrame, norm_algorithm: str="minmax") -> pd.DataFrame: - """normalize_hand this function normalizes the hand keypoints of a dataframe + """normalize_hands this function normalizes the hand keypoints of a dataframe :param dataframe: the dataframe to normalize :type dataframe: pd.DataFrame @@ -190,7 +193,7 @@ class KeypointExtractor: return dataframe def normalize_hand_minmax(self, dataframe: pd.DataFrame, hand: str) -> pd.DataFrame: - """normalize_hand_helper this function normalizes the hand keypoints of a dataframe with respect to the minimum and maximum coordinates + """normalize_hand_minmax 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 @@ -233,7 +236,7 @@ class KeypointExtractor: return dataframe def normalize_hand_bohacek(self, dataframe: pd.DataFrame, hand: str) -> pd.DataFrame: - """normalize_hand_helper this function normalizes the hand keypoints of a dataframe using the bohacek normalization algorithm + """normalize_hand_bohacek this function normalizes the hand keypoints of a dataframe using the Bohacek-normalization algorithm :param dataframe: the dataframe to normalize :type dataframe: pd.DataFrame @@ -268,7 +271,7 @@ class KeypointExtractor: delta_y[~mask] = (0.1 * height)[~mask] delta_x[~mask] = (delta_y + ((height - width) / 2))[~mask] - # Set the starting and ending point of the normalization bounding box + # set the starting and ending point of the normalization bounding box starting_x, starting_y = min_x - delta_x, min_y - delta_y ending_x, ending_y = max_x + delta_x, max_y + delta_y @@ -291,4 +294,63 @@ class KeypointExtractor: # 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: + """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 + """ + # get the columns that belong to the pose + pose_columns = np.array([i for i in range(66)]) + + # get the x, y coordinates of the pose keypoints + 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) + + # calculate shoulder distance + 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 + 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) + + # 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 + bbox_width, bbox_height = ending_x - starting_x, ending_y - starting_y + + # repeat the center coordinates and bounding box dimensions to match the shape of pose_coords + 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_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 + # normalize the pose keypoints based on the bounding box + 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 \ No newline at end of file diff --git a/visualize_data.ipynb b/visualize_data.ipynb index fa65d61..c90b744 100644 --- a/visualize_data.ipynb +++ b/visualize_data.ipynb @@ -121,7 +121,7 @@ " hand_coords = frame_df.values.reshape(21, 2)\n", " \n", " x_coords = hand_coords[:, ::2] #Even indices\n", - " y_coords = hand_coords[:, 1::2] #Uneven indices\n", + " y_coords = -hand_coords[:, 1::2] #Uneven indices (negative because pixels start from the top left)\n", " \n", " #Plot the keypoints\n", " plt.scatter(x_coords, y_coords)\n", @@ -167,6 +167,8 @@ "source": [ "#Plot the NORMALIZED hand keypoints (using minxmax)\n", "df = keypoints_extractor.extract_keypoints_from_video(video_name, normalize=\"minmax\")\n", + "plt.xlim(-0.5,0.5)\n", + "plt.ylim(-0.5,0.5)\n", "plot_hand_keypoints(df, hand, frame)" ] }, @@ -178,15 +180,82 @@ "source": [ "#Plot the NORMALIZED hand keypoints (using bohacek)\n", "df = keypoints_extractor.extract_keypoints_from_video(video_name, normalize=\"bohacek\")\n", + "plt.xlim(-0.5,0.5)\n", + "plt.ylim(-0.5,0.5)\n", "plot_hand_keypoints(df, hand, frame)" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Pose keypoint visualization" + ] + }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], - "source": [] + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "def plot_pose_keypoints(dataframe, frame):\n", + " pose_columns = np.array([i for i in range(32)])\n", + " \n", + " # get the x, y coordinates of the relevant pose keypoints\n", + " frame_df = dataframe.iloc[frame:frame+1, pose_columns]\n", + " pose_coords = frame_df.values.reshape(16, 2)\n", + " \n", + " x_coords = pose_coords[:, ::2] #Even indices\n", + " y_coords = -pose_coords[:, 1::2] #Uneven indices (negative because pixels start from the top left)\n", + " \n", + " #Plot the keypoints\n", + " plt.scatter(x_coords, y_coords)\n", + " return frame_df.style" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "#Set video, hand and frame to display\n", + "video_name = '69547.mp4'\n", + "frame = 2\n", + "%reload_ext autoreload" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from src.keypoint_extractor import KeypointExtractor\n", + "import numpy as np\n", + "\n", + "#Extract keypoints from requested video\n", + "keypoints_extractor = KeypointExtractor(\"data/videos/\")\n", + "\n", + "#Plot the hand keypoints\n", + "df = keypoints_extractor.extract_keypoints_from_video(video_name)\n", + "df.head()\n", + "plot_pose_keypoints(df, frame)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "norm_df = keypoints_extractor.extract_keypoints_from_video(video_name, normalize=\"bohacek\")\n", + "plt.xlim(-0.5,0.5)\n", + "plt.ylim(-0.5,0.5)\n", + "plot_pose_keypoints(norm_df, frame)" + ] }, { "cell_type": "code", @@ -219,7 +288,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.16" + "version": "3.9.13" }, "vscode": { "interpreter": {