Support skeleton (#155)

* pose * pose * pose * pose * 你的提交消息 * pose * pose * Delete train1.sh * pretreatment * configs * pose * reference * Update gaittr.py * naming * naming * Update transform.py * update for datasets * update README * update name and README * update * Update transform.py
2023-09-27 16:20:00 +08:00
parent 853bb1821d
commit 2c29afadf3
41 changed files with 4251 additions and 12 deletions
@@ -24,15 +24,24 @@ ls *.tgz | xargs -n1 tar xzvf

 After unpacking these compressed files, run this command:

-Step2 : To rearrange directory of GREW dataset, turning to id-type-view structure, Run 
+Step2-1 : To rearrange directory of GREW dataset(for silhouette), turning to id-type-view structure, Run 
 ```
 python datasets/GREW/rearrange_GREW.py --input_path Path_of_GREW-raw --output_path Path_of_GREW-rearranged
 ```  
+Step2-2 : To rearrange directory of GREW dataset(for pose), turning to id-type-view structure, Run 
+```
+python datasets/GREW/rearrange_GREW_pose.py --input_path Path_of_GREW-pose --output_path Path_of_GREW-pose-rearranged
+```  

-Step3: Transforming images to pickle file, run 
+Step3-1: Transforming images to pickle file, run 
 ```
 python datasets/pretreatment.py --input_path Path_of_GREW-rearranged --output_path Path_of_GREW-pkl --dataset GREW
 ```
+Step3-2: Transforming pose txts to pickle file, run 
+```
+python datasets/pretreatment.py --input_path Path_of_GREW-pose-rearranged --output_path Path_of_GREW-pose-pkl --pose --dataset GREW
+```
+
 Then you will see the structure like:

 - Processed
@@ -0,0 +1,92 @@
+import argparse
+import os
+import shutil
+from pathlib import Path
+
+from tqdm import tqdm
+
+TOTAL_Test = 24000
+TOTAL_Train = 20000
+
+def rearrange_train(train_path: Path, output_path: Path) -> None:
+    progress = tqdm(total=TOTAL_Train)
+    for sid in train_path.iterdir():
+        if not sid.is_dir():
+            continue
+        for sub_seq in sid.iterdir():
+            if not sub_seq.is_dir():
+                continue
+            for subfile in os.listdir(sub_seq):
+                src = os.path.join(train_path, sid.name, sub_seq.name)
+                dst = os.path.join(output_path, sid.name+'train', '00', sub_seq.name)
+                os.makedirs(dst,exist_ok=True)
+                if subfile not in os.listdir(dst) and subfile.endswith('_2d_pose.txt'):
+                    pose_subfile = 'pose_'+subfile
+                    os.symlink(os.path.join(src, subfile),
+                               os.path.join(dst, pose_subfile))
+        progress.update(1)
+
+def rearrange_test(test_path: Path, output_path: Path) -> None:
+    # for gallery
+    gallery = Path(os.path.join(test_path, 'gallery'))
+    probe = Path(os.path.join(test_path, 'probe'))
+    progress = tqdm(total=TOTAL_Test)
+    for sid in gallery.iterdir():
+        if not sid.is_dir():
+            continue
+        cnt = 1
+        for sub_seq in sid.iterdir():
+            if not sub_seq.is_dir():
+                continue
+            for subfile in sorted(os.listdir(sub_seq)):
+                src = os.path.join(gallery, sid.name, sub_seq.name)
+                dst = os.path.join(output_path, sid.name, '%02d'%cnt, sub_seq.name)
+                os.makedirs(dst,exist_ok=True)
+                if subfile not in os.listdir(dst) and subfile.endswith('_2d_pose.txt'):
+                    pose_subfile = 'pose_'+subfile
+                    os.symlink(os.path.join(src, subfile),
+                               os.path.join(dst, pose_subfile))
+            cnt += 1
+            progress.update(1)
+    # for probe
+    for sub_seq in probe.iterdir():
+        if not sub_seq.is_dir():
+            continue
+        for subfile in os.listdir(sub_seq):
+            src = os.path.join(probe, sub_seq.name)
+            dst = os.path.join(output_path, 'probe', '03', sub_seq.name)
+            os.makedirs(dst,exist_ok=True)
+            if subfile not in os.listdir(dst) and subfile.endswith('_2d_pose.txt'):
+                pose_subfile = 'pose_'+subfile
+                os.symlink(os.path.join(src, subfile),
+                            os.path.join(dst, pose_subfile))
+            progress.update(1)
+
+def rearrange_GREW(input_path: Path, output_path: Path) -> None:
+    os.makedirs(output_path, exist_ok=True)
+
+    for folder in input_path.iterdir():
+        if not folder.is_dir():
+            continue
+
+        print(f'Rearranging {folder}')
+        if folder.name == 'train':
+            rearrange_train(folder,output_path)
+        if folder.name == 'test':
+            rearrange_test(folder, output_path)
+        if folder.name == 'distractor':
+            pass
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='GREW rearrange tool')
+    parser.add_argument('-i', '--input_path', required=True, type=str,
+                        help='Root path of raw dataset.')
+    parser.add_argument('-o', '--output_path', default='GREW_rearranged', type=str,
+                        help='Root path for output.')
+
+    args = parser.parse_args()
+
+    input_path = Path(args.input_path).resolve()
+    output_path = Path(args.output_path).resolve()
+    rearrange_GREW(input_path, output_path)