darkliang
124 lines
4.3 KiB
Python
124 lines
4.3 KiB
Python
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))
|