--- a +++ b/datasets/HID/pretreatment_HID.py @@ -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))