crosstyan/OpenGait
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Support new Dataset: GREW

Browse Source
This commit is contained in:
Noah
2022-03-21 18:49:09 +08:00
parent 9c1aed2c87
commit 0071833d3f
5 changed files with 322 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
misc/partitions/GREW.json
+1
View File
File diff suppressed because one or more lines are too long
misc/pretreatment.py
+15 -6
View File
@@ -14,7 +14,7 @@ import numpy as np
from tqdm import tqdm
def imgs2pickle(img_groups: Tuple, output_path: Path, img_size: int = 64, verbose: bool = False) -> None:
def imgs2pickle(img_groups: Tuple, output_path: Path, img_size: int = 64, verbose: bool = False, dataset='CASIAB') -> None:
"""Reads a group of images and saves the data in pickle format.
Args:
@@ -31,6 +31,10 @@ def imgs2pickle(img_groups: Tuple, output_path: Path, img_size: int = 64, verbos
logging.debug(f'Reading sid {sinfo[0]}, seq {sinfo[1]}, view {sinfo[2]} from {img_file}')
img = cv2.imread(str(img_file), cv2.IMREAD_GRAYSCALE)
if dataset == 'GREW':
to_pickle.append(img.astype('uint8'))
continue
if img.sum() <= 10000:
if verbose:
@@ -76,6 +80,8 @@ def imgs2pickle(img_groups: Tuple, output_path: Path, img_size: int = 64, verbos
if to_pickle:
to_pickle = np.asarray(to_pickle)
dst_path = os.path.join(output_path, *sinfo)
# print(img_paths[0].as_posix().split('/'),img_paths[0].as_posix().split('/')[-5])
# dst_path = os.path.join(output_path, img_paths[0].as_posix().split('/')[-5], *sinfo) if dataset == 'GREW' else dst
os.makedirs(dst_path, exist_ok=True)
pkl_path = os.path.join(dst_path, f'{sinfo[2]}.pkl')
if verbose:
@@ -89,7 +95,7 @@ def imgs2pickle(img_groups: Tuple, output_path: Path, img_size: int = 64, verbos
def pretreat(input_path: Path, output_path: Path, img_size: int = 64, workers: int = 4, verbose: bool = False) -> None:
def pretreat(input_path: Path, output_path: Path, img_size: int = 64, workers: int = 4, verbose: bool = False, dataset: str = 'CASIAB') -> None:
"""Reads a dataset and saves the data in pickle format.
Args:
@@ -103,6 +109,8 @@ def pretreat(input_path: Path, output_path: Path, img_size: int = 64, workers: i
logging.info(f'Listing {input_path}')
total_files = 0
for img_path in input_path.rglob('*.png'):
if 'gei.png' in img_path.as_posix():
continue
if verbose:
logging.debug(f'Adding {img_path}')
*_, sid, seq, view, _ = img_path.as_posix().split('/')
@@ -115,18 +123,19 @@ def pretreat(input_path: Path, output_path: Path, img_size: int = 64, workers: i
with mp.Pool(workers) as pool:
logging.info(f'Start pretreating {input_path}')
for _ in pool.imap_unordered(partial(imgs2pickle, output_path=output_path, img_size=img_size, verbose=verbose), img_groups.items()):
for _ in pool.imap_unordered(partial(imgs2pickle, output_path=output_path, img_size=img_size, verbose=verbose, dataset=dataset), img_groups.items()):
progress.update(1)
logging.info('Done')
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='OpenGait dataset pretreatment module.')
parser.add_argument('-r', '--input_path', default='', type=str, help='Root path of raw dataset.')
parser.add_argument('-i', '--input_path', default='', type=str, help='Root path of raw dataset.')
parser.add_argument('-o', '--output_path', default='', type=str, help='Output path of pickled dataset.')
parser.add_argument('-l', '--log_file', default='./pretreatment.log', type=str, help='Log file path. Default: ./pretreatment.log')
parser.add_argument('-n', '--n_workers', default=4, type=int, help='Number of thread workers. Default: 4')
parser.add_argument('-i', '--img_size', default=64, type=int, help='Image resizing size. Default 64')
parser.add_argument('-r', '--img_size', default=64, type=int, help='Image resizing size. Default 64')
parser.add_argument('-d', '--dataset', default='CASIAB', type=str, help='Dataset for pretreatment.')
parser.add_argument('-v', '--verbose', default=False, action='store_true', help='Display debug info.')
args = parser.parse_args()
@@ -138,4 +147,4 @@ if __name__ == '__main__':
for k, v in args.__dict__.items():
logging.debug(f'{k}: {v}')
pretreat(input_path=Path(args.input_path), output_path=Path(args.output_path), img_size=args.img_size, workers=args.n_workers, verbose=args.verbose)
pretreat(input_path=Path(args.input_path), output_path=Path(args.output_path), img_size=args.img_size, workers=args.n_workers, verbose=args.verbose, dataset=args.dataset)
misc/rearrange_GREW.py
+89
View File
@@ -0,0 +1,89 @@
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('.png'):
os.symlink(os.path.join(src, subfile),
os.path.join(dst, 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('.png'):
os.symlink(os.path.join(src, subfile),
os.path.join(dst, 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('.png'):
os.symlink(os.path.join(src, subfile),
os.path.join(dst, 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)