SkeletonGait

2024-03-05 16:13:11 +08:00
parent 2019f9a525
commit a7e6a7886a
19 changed files with 2040 additions and 7 deletions
@@ -0,0 +1,77 @@
+import os
+import sys
+import argparse
+from tqdm import tqdm
+from glob import glob
+
+def get_args():
+    parser = argparse.ArgumentParser(description='Symlink silouette data and pose data into the same folder for SkeletonGait++ training.')
+    parser.add_argument('--heatmapt_data_path', type=str, required=True, help="path of heatmap data, must be the absolute path.")
+    parser.add_argument('--silouette_data_path', type=str, required=True, help="path of silouette data, must be the absolute path.")
+    parser.add_argument('--dataset_pkl_ext_name', type=str, default='.pkl', help="The extent name for .pkl files of silouettes data.")
+    parser.add_argument('--output_path', type=str, required=True, help="path of output data")
+    opt = parser.parse_args()
+    return opt
+
+def main():
+    opt = get_args()
+    heatmap_data_path = opt.heatmapt_data_path
+    silouette_data_path = opt.silouette_data_path
+    if not os.path.exists(heatmap_data_path):
+        print(f"heatmap data path {heatmap_data_path} does not exist.")
+        sys.exit(1)
+    if not os.path.exists(silouette_data_path):
+        print(f"silouette data path {silouette_data_path} does not exist.")
+        sys.exit(1)
+    
+    all_heatmap_files = sorted(glob(os.path.join(heatmap_data_path, "*/*/*/*.pkl")))
+    all_silouette_files = sorted(glob(os.path.join(silouette_data_path, f"*/*/*/*{opt.dataset_pkl_ext_name}")))
+    # print(len(all_heatmap_files), len(all_silouette_files))
+    # assert len(all_heatmap_files) == len(all_silouette_files), "The number of heatmap files and silouette files are not equal."
+    
+    if len(all_heatmap_files) >= len(all_silouette_files):
+        for heatmap_file in tqdm(all_heatmap_files):
+            tmp_list = heatmap_file.split('/')
+            sil_folder = os.path.join(silouette_data_path, *tmp_list[-4:-1])
+            if not os.path.exists(sil_folder):
+                print(f"silouette folder {sil_folder} does not exist.")
+                continue
+            else:
+                silouette_file = sorted(glob(os.path.join(sil_folder, f"*{opt.dataset_pkl_ext_name}")))[0]
+
+            output_file = os.path.join(opt.output_path, *tmp_list[-4:-1])
+            os.makedirs(output_file, exist_ok=True)
+            os.system(f"ln -s {silouette_file} {output_file}/1_sil.pkl")
+            os.system(f"ln -s {heatmap_file} {output_file}/0_heatmap.pkl")
+    else:
+        for silouette_file in tqdm(all_silouette_files):
+            tmp_list = silouette_file.split('/')
+            heatmap_folder = os.path.join(heatmap_data_path, *tmp_list[-4:-1])
+            if not os.path.exists(heatmap_folder):
+                print(f"heatmap folder {heatmap_folder} does not exist.")
+                continue
+            else:
+                heatmap_file = sorted(glob(os.path.join(heatmap_folder, "*.pkl")))[0]
+
+            output_file = os.path.join(opt.output_path, *tmp_list[-4:-1])
+            os.makedirs(output_file, exist_ok=True)
+            os.system(f"ln -s {silouette_file} {output_file}/1_sil.pkl")
+            os.system(f"ln -s {heatmap_file} {output_file}/0_heatmap.pkl")
+
+    print("Done! Output data is in ", opt.output_path)
+
+    # for tmp_file in tqdm(iter_files):
+    #     heatmap_file = all_heatmap_files[i]
+    #     silouette_file = all_silouette_files[i]
+    #     sil_tmp_list = silouette_file.split('/')
+    #     heatmap_tmp_list = heatmap_file.split('/')
+    #     if 
+
+    #     output_file = os.path.join(opt.output_path, *tmp_list[-4:-1])
+    #     os.makedirs(output_file, exist_ok=True)
+
+    #     os.system(f"ln -s {silouette_file} {output_file}/1_sil.pkl")
+    #     os.system(f"ln -s {heatmap_file} {output_file}/0_heatmap.pkl")
+
+if __name__ == "__main__":
+    main()