Reconstruct the directory structure

2022-04-12 13:44:07 +08:00
parent 213b3a658f
commit 28f50410a5
32 changed files with 162 additions and 45 deletions
@@ -0,0 +1,123 @@
+import os
+import cv2
+import numpy as np
+import argparse
+import pickle
+from tqdm import tqdm
+
+parser = argparse.ArgumentParser(description='Test')
+parser.add_argument('--input_train_path', default='', type=str,
+                    help='Root path of train.')
+parser.add_argument('--input_gallery_path', default='', type=str,
+                    help='Root path of gallery.')
+parser.add_argument('--input_probe_path', default='', type=str,
+                    help='Root path of probe.')
+parser.add_argument('--output_path', default='', type=str,
+                    help='Root path for output.')
+
+opt = parser.parse_args()
+
+OUTPUT_PATH = opt.output_path
+print('Pretreatment Start.\n'
+      'Input train path: {}\n'
+      'Input gallery path: {}\n'
+      'Input probe path: {}\n'
+      'Output path: {}\n'.format(
+          opt.input_train_path, opt.input_gallery_path, opt.input_probe_path, OUTPUT_PATH))
+
+INPUT_PATH = opt.input_train_path
+print("Walk the input train path")
+id_list = os.listdir(INPUT_PATH)
+id_list.sort()
+
+for _id in tqdm(id_list):
+    seq_type = os.listdir(os.path.join(INPUT_PATH, _id))
+    seq_type.sort()
+    for _seq_type in seq_type:
+        out_dir = os.path.join(OUTPUT_PATH, _id, _seq_type, "default")
+        count_frame = 0
+        all_imgs = []
+        frame_list = sorted(os.listdir(
+            os.path.join(INPUT_PATH, _id, _seq_type)))
+        for _frame_name in frame_list:
+            frame_path = os.path.join(
+                INPUT_PATH, _id, _seq_type, _frame_name)
+            img = cv2.imread(frame_path, cv2.IMREAD_GRAYSCALE)
+            if img is not None:
+                # Save the img
+                all_imgs.append(img)
+                count_frame += 1
+
+        all_imgs = np.asarray(all_imgs)
+
+        if count_frame > 0:
+            os.makedirs(out_dir, exist_ok=True)
+            all_imgs_pkl = os.path.join(out_dir, '{}.pkl'.format(_seq_type))
+            pickle.dump(all_imgs, open(all_imgs_pkl, 'wb'))
+
+        # Warn if the sequence contains less than 5 frames
+        if count_frame < 5:
+            print('Seq:{}-{}, less than 5 valid data.'.format(_id, _seq_type))
+
+print("Walk the input gallery path")
+INPUT_PATH = opt.input_gallery_path
+id_list = os.listdir(INPUT_PATH)
+id_list.sort()
+for _id in tqdm(id_list):
+    seq_type = os.listdir(os.path.join(INPUT_PATH, _id))
+    seq_type.sort()
+    for _seq_type in seq_type:
+        out_dir = os.path.join(OUTPUT_PATH, _id, _seq_type, "default")
+        count_frame = 0
+        all_imgs = []
+        frame_list = sorted(os.listdir(
+            os.path.join(INPUT_PATH, _id, _seq_type)))
+        for _frame_name in frame_list:
+            frame_path = os.path.join(
+                INPUT_PATH, _id, _seq_type, _frame_name)
+            img = cv2.imread(frame_path, cv2.IMREAD_GRAYSCALE)
+            if img is not None:
+                # Save the img
+                all_imgs.append(img)
+                count_frame += 1
+
+        all_imgs = np.asarray(all_imgs)
+
+        if count_frame > 0:
+            os.makedirs(out_dir, exist_ok=True)
+            all_imgs_pkl = os.path.join(out_dir, '{}.pkl'.format(_seq_type))
+            pickle.dump(all_imgs, open(all_imgs_pkl, 'wb'))
+
+        # Warn if the sequence contains less than 5 frames
+        if count_frame < 5:
+            print('Seq:{}-{}, less than 5 valid data.'.format(_id, _seq_type))
+    print("Finish {}".format(_id))
+
+print("Walk the input probe path")
+INPUT_PATH = opt.input_probe_path
+seq_type = os.listdir(INPUT_PATH)
+seq_type.sort()
+
+_id = "probe"
+for _seq_type in tqdm(seq_type):
+    out_dir = os.path.join(OUTPUT_PATH, _id, _seq_type, "default")
+    count_frame = 0
+    all_imgs = []
+    frame_list = sorted(os.listdir(
+        os.path.join(INPUT_PATH, _seq_type)))
+    for _frame_name in frame_list:
+        frame_path = os.path.join(
+            INPUT_PATH, _seq_type, _frame_name)
+        img = cv2.imread(frame_path, cv2.IMREAD_GRAYSCALE)
+        if img is not None:
+            # Save the img
+            all_imgs.append(img)
+            count_frame += 1
+    all_imgs = np.asarray(all_imgs)
+    if count_frame > 0:
+        os.makedirs(out_dir, exist_ok=True)
+        all_imgs_pkl = os.path.join(out_dir, '{}.pkl'.format(_seq_type))
+        pickle.dump(all_imgs, open(all_imgs_pkl, 'wb'))
+    # Warn if the sequence contains less than 5 frames
+    if count_frame < 5:
+        print('Seq:{}-{}, less than 5 valid data.'.format(_id, _seq_type))