Add time in csv

davidpagnon · May 23, 2023 · 4fe6ea7 · 4fe6ea7
1 parent 04eb529
commit 4fe6ea7
Show file tree

Hide file tree

Showing 6 changed files with 46 additions and 39 deletions.
diff --git a/Sports2D/Sports2D.ipynb b/Sports2D/Sports2D.ipynb
diff --git a/Sports2D/Sports2D.py b/Sports2D/Sports2D.py
@@ -145,7 +145,7 @@ def base_params(config_dict):
     try:
         1/frame_rate
     except ZeroDivisionError:
-        print(f'Frame rate of {str(video_dir / video_file)} could not be retrieved: check that it exists at the correct path')
+        print('Frame rate could not be retrieved: check that your video exists at the correct path')
         raise
     video.release()
 

diff --git a/Sports2D/Utilities/Blazepose_runsave.py b/Sports2D/Utilities/Blazepose_runsave.py
@@ -49,7 +49,7 @@
 
 
 ## FUNCTIONS
-def save_to_csv_or_h5(kpt_list, output_folder, video_name, to_csv, to_h5):
+def save_to_csv_or_h5(kpt_list, fps, output_folder, video_name, to_csv, to_h5):
     '''
     Saves blazepose keypoint coordinates to csv or h5 file, 
     in the DeepLabCut format.
@@ -66,15 +66,16 @@ def save_to_csv_or_h5(kpt_list, output_folder, video_name, to_csv, to_h5):
     '''
 
     # Prepare dataframe file
-    scorer = ['DavidPagnon']*len(mp_pose.PoseLandmark)*3
-    individuals = ['person']*len(mp_pose.PoseLandmark)*3
+    scorer = ['DavidPagnon']*(len(mp_pose.PoseLandmark)*3+1)
+    individuals = ['person']*(len(mp_pose.PoseLandmark)*3+1)
     bodyparts = [[p]*3 for p in ['Nose', 'LEyeInner', 'LEye', 'LEyeOuter', 'REyeInner', 'REye', 'REyeOuter', 'LEar', 'REar', 'LMouth', 'RMouth', 'LShoulder', 'RShoulder', 'LElbow', 'RElbow', 'LWrist', 'RWrist', 'LPinky', 'RPinky', 'LIndex', 'RIndex', 'LThumb', 'RThumb', 'LHip', 'RHip', 'LKnee', 'RKnee', 'LAnkle', 'RAnkle', 'LHeel', 'RHeel', 'LBigToe', 'RBigToe']]
-    bodyparts = [item for sublist in bodyparts for item in sublist]
-    coords = ['x', 'y', 'likelihood']*len(mp_pose.PoseLandmark)
-
+    bodyparts = ['Times'] + [item for sublist in bodyparts for item in sublist]
+    coords = ['seconds'] + ['x', 'y', 'likelihood']*len(mp_pose.PoseLandmark)
     tuples = list(zip(scorer, individuals, bodyparts, coords))
     index_csv = pd.MultiIndex.from_tuples(tuples, names=['scorer', 'individuals', 'bodyparts', 'coords'])
-    df = pd.DataFrame(np.array(kpt_list).T, index=index_csv).T
+
+    time = np.expand_dims( np.arange(0,len(kpt_list)/fps, 1/fps), axis=0 )
+    df = pd.DataFrame(np.concatenate(( time, np.array(kpt_list).T)), index=index_csv).T
 
     if to_csv:
         csv_file = os.path.join(output_folder, video_name+'_BLAZEPOSE_points.csv')
@@ -229,7 +230,7 @@ def blazepose_detec_func(**args):
 
     # Save coordinates
     if to_csv or to_h5:
-        save_to_csv_or_h5(kpt_list, output_folder, video_name, to_csv, to_h5)
+        save_to_csv_or_h5(kpt_list, fps, output_folder, video_name, to_csv, to_h5)
 
     if to_json:
         save_to_json(kpt_list, output_folder, video_name)

diff --git a/Sports2D/compute_angles.py b/Sports2D/compute_angles.py
@@ -123,15 +123,17 @@ def display_figures_fun(df_list):
     - matplotlib window with tabbed figures for each angle
     '''
 
-    angle_names = df_list[0].columns.get_level_values(2)
+    angle_names = df_list[0].iloc[:,1:].columns.get_level_values(2)
+    time = df_list[0].iloc[:,0]
 
     pw = common.plotWindow()
     for id, angle in enumerate(angle_names): # angles
         f = plt.figure()
 
         plt.plot()
-        [plt.plot(df_list[0].index, df.iloc[:,id], label=['unfiltered' if i==0 else 'filtered' if i==1 else ''][0]) for i,df in enumerate(df_list)]
-        plt.ylabel(angle) # nom angle
+        [plt.plot(time, df.iloc[:,id+1], label=['unfiltered' if i==0 else 'filtered' if i==1 else ''][0]) for i,df in enumerate(df_list)]
+        plt.xlabel('Time (seconds)')
+        plt.ylabel(angle)
         plt.legend()
 
         pw.addPlot(angle, f)
@@ -187,7 +189,7 @@ def flip_left_right_direction(df_points):
     righ_orientation = df_points.iloc[:,df_points.columns.get_level_values(2)=='RBigToe'].iloc[:,0] - df_points.iloc[:,df_points.columns.get_level_values(2)=='RHeel'].iloc[:,0]
     left_orientation = df_points.iloc[:,df_points.columns.get_level_values(2)=='LBigToe'].iloc[:,0] - df_points.iloc[:,df_points.columns.get_level_values(2)=='LHeel'].iloc[:,0]
     orientation = righ_orientation + left_orientation
-    df_points.iloc[:,1::3] = df_points.iloc[:,1::3] * np.where(orientation>=0, 1, -1).reshape(-1,1)
+    df_points.iloc[:,2::3] = df_points.iloc[:,2::3] * np.where(orientation>=0, 1, -1).reshape(-1,1)
 
     return df_points
 
@@ -217,7 +219,6 @@ def joint_angles_series_from_points(df_points, angle_params):
     ang_series = np.where(ang_series>180,ang_series-360,ang_series)
     ang_series = np.where((ang_series==0) | (ang_series==90) | (ang_series==180), +0, ang_series)
 
-
     return ang_series
 
 
@@ -387,17 +388,18 @@ def compute_angles_fun(config_dict):
     df_angles_list = []
     for i, c in enumerate(csv_paths):
         # Prepare angle csv header
-        scorer = ['DavidPagnon']*angle_nb
-        individuals = [f'person{i}']*angle_nb
-        angs = joint_angles + segment_angles
-        coords = [joint_angle_dict.get(j)[1] for j in joint_angles] + [segment_angle_dict.get(s)[1] for s in segment_angles]
+        scorer = ['DavidPagnon']*(angle_nb+1)
+        individuals = [f'person{i}']*(angle_nb+1)
+        angs = ['Time'] + joint_angles + segment_angles
+        coords = ['seconds'] + [joint_angle_dict.get(j)[1] for j in joint_angles] + [segment_angle_dict.get(s)[1] for s in segment_angles]
         tuples = list(zip(scorer, individuals, angs, coords))
-        index_angs_csv = pd.MultiIndex.from_tuples(tuples, names=['scorer', 'individuals', 'joints', 'angles'])    
+        index_angs_csv = pd.MultiIndex.from_tuples(tuples, names=['scorer', 'individuals', 'angs', 'coords'])    
 
         # Compute angles for each person, for each angle, with each required keypoint position
         logging.info(f'Person {i}: Computing 2D joint and segment angles.')
         with open(c) as c_f:
             df_points = pd.read_csv(c_f, header=[0,1,2,3])
+            time = [np.array(df_points.iloc[:,1])]
 
             # Flip along x when feet oriented to the left
             df_points = flip_left_right_direction(df_points)
@@ -416,7 +418,7 @@ def compute_angles_fun(config_dict):
                 s_ang_series = segment_angles_series_from_points(df_points, angle_params, s)
                 segment_angle_series += [s_ang_series]
 
-            angle_series = joint_angle_series + segment_angle_series
+            angle_series = time + joint_angle_series + segment_angle_series
             df_angles = []
             df_angles += [pd.DataFrame(angle_series, index=index_angs_csv).T]
 
@@ -464,7 +466,7 @@ def compute_angles_fun(config_dict):
             writer = cv2.VideoWriter(str(video_pose2), fourcc, fps, (int(W), int(H)))
 
         # Preferentially from pose image files
-        frames_img = sorted(list(img_pose.glob('*.png')))
+        frames_img = list(img_pose.glob('*'))
         if len(frames_img)>0:
             for frame_nb in range(df_angles_list[0].shape[0]):
                 df_angles_list_frame = [df_angles_list[n].iloc[frame_nb,:] for n in range(len(df_angles_list))]
@@ -500,4 +502,4 @@ def compute_angles_fun(config_dict):
             os.rename(video_pose2,video_pose)
             if Path.exists(video_pose2): os.remove(video_pose2)
 
-    logging.info(f'Done.')
+    logging.info(f'Done.')
diff --git a/Sports2D/detect_pose.py b/Sports2D/detect_pose.py
@@ -91,14 +91,14 @@ def display_figures_fun(df_list):
         f = plt.figure()
 
         axX = plt.subplot(211)
-        [plt.plot(df_list[0].index, df.iloc[:,id*3], label=['unfiltered' if i==0 else 'filtered' if i==1 else ''][0]) for i,df in enumerate(df_list)]
+        [plt.plot(df.iloc[:,0], df.iloc[:,id*3+1], label=['unfiltered' if i==0 else 'filtered' if i==1 else ''][0]) for i,df in enumerate(df_list)]
         plt.setp(axX.get_xticklabels(), visible=False)
         axX.set_ylabel(keypoint+' X')
         plt.legend()
 
         axY = plt.subplot(212)
-        [plt.plot(df_list[0].index, df.iloc[:,id*3+1]) for df in df_list]
-        plt.setp(axY.get_xticklabels(), visible=False)
+        [plt.plot(df.iloc[:,0], df.iloc[:,id*3+2]) for df in df_list]
+        axY.set_xlabel('Time (seconds)')
         axY.set_ylabel(keypoint+' Y')
 
         pw.addPlot(keypoint, f)
@@ -263,7 +263,7 @@ def sort_people(keyptpre, keypt, nb_persons_to_detect):
     return keypt
 
 
-def json_to_csv(json_path, pose_model, interp_gap_smaller_than, filter_options, show_plots):
+def json_to_csv(json_path, frame_rate, pose_model, interp_gap_smaller_than, filter_options, show_plots):
     '''
     Converts frame-by-frame json coordinate files 
     to one csv files per detected person
@@ -288,7 +288,7 @@ def json_to_csv(json_path, pose_model, interp_gap_smaller_than, filter_options,
 
     # Retrieve coordinates
     logging.info('Sorting people across frames.')
-    json_fnames = sorted(list(json_path.glob('*.json')))
+    json_fnames = list(json_path.glob('*.json'))
     nb_persons_to_detect = max([len(json.load(open(json_fname))['people']) for json_fname in json_fnames])
     Coords = [np.array([]).reshape(0,keypoints_nb*3)] * nb_persons_to_detect
     for json_fname in json_fnames:    # for each frame
@@ -316,18 +316,22 @@ def json_to_csv(json_path, pose_model, interp_gap_smaller_than, filter_options,
     # Inject coordinates in dataframes and save
     for i in range(nb_persons_to_detect): 
         # Prepare csv header
-        scorer = ['DavidPagnon']*keypoints_nb*3
-        individuals = [f'person{i}']*keypoints_nb*3
+        scorer = ['DavidPagnon']*(keypoints_nb*3+1)
+        individuals = [f'person{i}']*(keypoints_nb*3+1)
         bodyparts = [[p]*3 for p in keypoints_names_rearranged]
-        bodyparts = [item for sublist in bodyparts for item in sublist]
-        coords = ['x', 'y', 'likelihood']*keypoints_nb
+        bodyparts = ['Time']+[item for sublist in bodyparts for item in sublist]
+        coords = ['seconds']+['x', 'y', 'likelihood']*keypoints_nb
         tuples = list(zip(scorer, individuals, bodyparts, coords))
         index_csv = pd.MultiIndex.from_tuples(tuples, names=['scorer', 'individuals', 'bodyparts', 'coords'])
-
+
+        # Create dataframe
+        df_list=[]
+        time = np.expand_dims( np.arange(0,len(Coords[i])/frame_rate, 1/frame_rate), axis=0 )
+        time_coords = np.concatenate(( time, Coords[i].T ))
+        df_list += [pd.DataFrame(time_coords, index=index_csv).T]
+
         # Interpolate
         logging.info(f'Person {i}: Interpolating missing sequences if they are smaller than {interp_gap_smaller_than} frames.')
-        df_list=[]
-        df_list += [pd.DataFrame(Coords[i].T, index=index_csv).T]
         df_list[0] = df_list[0].apply(common.interpolate_zeros_nans, axis=0, args = [interp_gap_smaller_than, 'linear'])
 
         # Filter
@@ -483,9 +487,9 @@ def save_imgvid_reID(video_path, video_result_path, save_vid=1, save_img=1, *pos
     while(cap.isOpened()):
         ret, frame = cap.read()
         if ret == True:
-            X = [np.array(coord.iloc[f,1::3]) for coord in coords]
+            X = [np.array(coord.iloc[f,2::3]) for coord in coords]
             X = [np.where(x==0., np.nan, x) for x in X]
-            Y = [np.array(coord.iloc[f,2::3]) for coord in coords]
+            Y = [np.array(coord.iloc[f,3::3]) for coord in coords]
             Y = [np.where(y==0., np.nan, y) for y in Y]
 
             # Draw bounding box
@@ -584,7 +588,7 @@ def detect_pose_fun(config_dict):
             os.chdir(root_dir)
 
     # Sort people and save to csv, optionally display plot
-        json_to_csv(json_path, pose_model, interp_gap_smaller_than, filter_options, show_plots)
+        json_to_csv(json_path, frame_rate, pose_model, interp_gap_smaller_than, filter_options, show_plots)
 
     # Save images and files after reindentification
         if save_img and save_vid:

diff --git a/setup.cfg b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
 name = sports2d
-version = 0.2.0
+version = 0.2.1
 author = David Pagnon
 author_email = [email protected]
 description = Detect pose and compute 2D joint angles from a video.