handles other pose models better

Author: davidpagnon

Sports2D/Sports2D.py

@@ -138,9 +138,9 @@ def base_params(config_dict):
     if video_dir == '': video_dir = os.getcwd()
     video_files = config_dict.get('project').get('video_files')
     if isinstance(video_files, str):
-        video_files = [Path(video_files)]
+        video_files = [Path(os.getcwd()) / video_files]
     else: 
-        video_files = [Path(v) for v in video_files]
+        video_files = [Path(os.getcwd()) / v for v in video_files]
     result_dir = Path(config_dict.get('project').get('result_dir')).resolve()
     if result_dir == '': result_dir = os.getcwd()
 

Sports2D/compute_angles.py

@@ -402,7 +402,7 @@ def compute_angles_fun(config_dict, video_file):
 
     # Find csv position files in video_dir, search pose_model and video_file.stem
     logging.info(f'Retrieving csv position files in {result_dir}...')
-    csv_paths = list(result_dir.glob(f'*{video_file.stem}_{pose_model}_*points*.csv'))
+    csv_paths = list(result_dir.glob(f'*{video_file.stem}_*points*.csv'))
     logging.info(f'{len(csv_paths)} persons found.')
 
     # Compute angles
@@ -476,9 +476,9 @@ def compute_angles_fun(config_dict, video_file):
     # Add angles to vid and img
     if show_angles_img or show_angles_vid:
         video_base = Path(video_dir / video_file)
-        img_pose = result_dir / (video_base.stem + '_' + pose_model + '_img')
-        video_pose = result_dir / (video_base.stem + '_' + pose_model + '.mp4')
-        video_pose2 = result_dir / (video_base.stem + '_' + pose_model + '2.mp4')
+        img_pose = result_dir / (video_base.stem + '_img')
+        video_pose = result_dir / (video_base.stem + '.mp4')
+        video_pose2 = result_dir / (video_base.stem + '2.mp4')
 
         if show_angles_vid:
             logging.info(f'Saving video in {str(video_pose)}.')
@@ -514,7 +514,7 @@ def compute_angles_fun(config_dict, video_file):
                     frame = overlay_angles(frame, df_angles_list_frame)
                     if show_angles_img:
                         if frame_nb==0: img_pose.mkdir(parents=True, exist_ok=True)
-                        cv2.imwrite(str(img_pose / (video_base.stem + '_' + pose_model + '.' + str(frame_nb).zfill(5)+'.png')), frame)
+                        cv2.imwrite(str(img_pose / (video_base.stem + '_' + '.' + str(frame_nb).zfill(5)+'.png')), frame)
                     if show_angles_vid:
                         writer.write(frame)
                     frame_nb+=1

Sports2D/detect_pose.py

@@ -286,7 +286,6 @@ def json_to_csv(json_path, frame_rate, pose_model, interp_gap_smaller_than, filt
     model = eval(pose_model)
     keypoints_ids = [node.id for _, _, node in RenderTree(model) if node.id!=None]
     keypoints_names = [node.name for _, _, node in RenderTree(model) if node.id!=None]
-    keypoints_names_rearranged = [y for x,y in sorted(zip(keypoints_ids,keypoints_names))]
     keypoints_nb = len(keypoints_ids)
 
     # Retrieve coordinates
@@ -300,7 +299,8 @@ def json_to_csv(json_path, frame_rate, pose_model, interp_gap_smaller_than, filt
             keypt = []
             # Retrieve coords for this frame 
             for ppl in range(len(json_file['people'])):  # for each detected person
-                keypt += [np.asarray(json_file['people'][ppl]['pose_keypoints_2d']).reshape(-1,3)]
+                keypt_all = np.asarray(json_file['people'][ppl]['pose_keypoints_2d']).reshape(-1,3)[keypoints_ids]
+                keypt += [keypt_all]
             keypt = np.array(keypt)
             # Make sure keypt is as large as the number of persons that need to be detected
             if len(keypt) < nb_persons_to_detect:
@@ -321,7 +321,7 @@ def json_to_csv(json_path, frame_rate, pose_model, interp_gap_smaller_than, filt
         # Prepare csv header
         scorer = ['DavidPagnon']*(keypoints_nb*3+1)
         individuals = [f'person{i}']*(keypoints_nb*3+1)
-        bodyparts = [[p]*3 for p in keypoints_names_rearranged]
+        bodyparts = [[p]*3 for p in keypoints_names]
         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))
@@ -604,7 +604,7 @@ def detect_pose_fun(config_dict, video_file):
 
     if pose_algo == 'OPENPOSE':
         pose_model = config_dict.get('pose').get('OPENPOSE').get('openpose_model')
-        json_path = result_dir / '_'.join((video_file_stem,pose_model,'json'))
+        json_path = result_dir / '_'.join((video_file_stem,'json'))
 
         # Pose detection skipped if load existing json files
         if load_pose and len(list(json_path.glob('*')))>0:
@@ -633,8 +633,8 @@ def detect_pose_fun(config_dict, video_file):
             elif platform == "linux" or platform=="linux2":
                 run_openpose_linux(video_path, json_path, pose_model)
             os.chdir(root_dir)
-
-     
+    
+    
     elif pose_algo == 'BLAZEPOSE':
         pose_model = pose_algo
         json_path = result_dir / '_'.join((video_file_stem,pose_algo,'json'))