BigGait

2024-05-11 09:27:34 +08:00
parent dc8fc2ffc7
commit a9f711d2e6
15 changed files with 1836 additions and 0 deletions
@@ -0,0 +1,71 @@
+import os
+from time import time
+from multiprocessing import Pool
+from tqdm import tqdm
+import numpy as np
+import os
+import pickle
+import numpy as np
+import cv2
+from tqdm import tqdm
+
+SRC_0 = '../SUSTech1K-Released-2023_mask/'
+DST_0 = '../SUSTech1K-Released-2023_mask_256128pkl/'
+
+SRC = SRC_0             # Path_of_RGB_rearranged
+DST = DST_0             # Path_of_RGB_256128pkl_PadResized
+
+def resize_with_padding(img, target_size):
+    h, w, _ = img.shape
+    target_h, target_w = target_size
+    resized_img = cv2.resize(img, (int(w * target_h / h), target_h))
+    padded_img = np.zeros((target_h, target_w, 3), dtype=np.uint8)
+    x_offset = (target_w - resized_img.shape[1]) // 2
+    if x_offset < 0 :
+        x_offset = abs(x_offset)
+        padded_img = resized_img[:, x_offset:x_offset+target_w,:]
+    else: 
+        padded_img[:, x_offset:x_offset + resized_img.shape[1]] = resized_img
+    return padded_img
+
+def job(src, id):
+    for ty in sorted(os.listdir(os.path.join(src, id))):
+        for vi in sorted(os.listdir(os.path.join(src, id, ty))):
+            exist_file = os.path.join(DST, id, ty, vi, vi+"-aligned-rgbs.pkl")
+            if os.path.exists(exist_file):
+                print('Have Passed: ' + DST + '/' + id + '/' + ty)
+                continue
+            ratios = []
+            aligned_imgs = []
+            for img_file in sorted(os.listdir(os.path.join(src, id, ty, vi))):
+                img_path = os.path.join(src, id, ty, vi, img_file)
+                img = cv2.imread(img_path)
+                ratio = img.shape[1]/img.shape[0]
+                ratios.append(ratio)
+                aligned_img = np.transpose(cv2.cvtColor(resize_with_padding(img, (256, 128)), cv2.COLOR_BGR2RGB), (2, 0, 1))
+                aligned_imgs.append(aligned_img)
+                if len(aligned_imgs) > 0:
+                    output_path = os.path.join(DST, id, ty, vi)
+                    os.makedirs(output_path, exist_ok=True)
+                    pickle.dump(np.asarray(aligned_imgs), open(os.path.join(output_path, vi+"-aligned-rgbs.pkl"), "wb"))
+                    pickle.dump(np.asarray(ratios), open(os.path.join(output_path, vi+"-ratios.pkl"), "wb"))
+            print('Successfully saved: ' + DST + '/' + id + '/' + ty +  '/' + vi)
+                    
+if __name__ == '__main__':
+    a = time()
+    po = Pool(8)
+    src_path = SRC
+    
+    cnt = 0
+    need_data = sorted(os.listdir(src_path))
+    for id in tqdm(need_data[:]):
+        po.apply_async(job,(src_path, id,))
+        cnt = cnt + 1
+    
+    print('---START---')
+    po.close()
+    po.join()
+    print(cnt)
+
+    t = time() - a
+    print('---END---{}'.format(t))