EasyMocap/apps/preprocess/merge_or_split.py
2022-08-21 23:05:59 +08:00

64 lines
2.4 KiB
Python

'''
@ Date: 2022-07-25 21:56:58
@ Author: Qing Shuai
@ Mail: s_q@zju.edu.cn
@ LastEditors: Qing Shuai
@ LastEditTime: 2022-07-25 21:59:29
@ FilePath: /EasyMocapPublic/apps/preprocess/merge_or_split.py
'''
import shutil
from easymocap.mytools.debug_utils import log, mkdir, mywarn, run_cmd
import os
from os.path import join
from tqdm import tqdm
def merge_directories(root, out):
mkdir(join(out, 'images', 'merge'))
sequence = sorted(os.listdir(root))
log('>>> Totally {} sub-folders'.format(len(sequence)))
records = []
for seq in tqdm(sequence, 'check sequence'):
subs = sorted(os.listdir(join(root, seq, 'images')))
for sub in subs:
imgnames = sorted(os.listdir(join(root, seq, 'images', sub)))
for imgname in imgnames:
records.append((seq, sub, imgname))
log('>>> Totally {} records'.format(len(records)))
for (seq, sub, imgname) in tqdm(records):
srcname = join(root, seq, 'images', sub, imgname)
dstname = join(out, 'images', 'merge', '{}+{}+{}'.format(seq, sub, imgname))
if not os.path.exists(dstname):
cmd = 'ln -s {} {}'.format(srcname, dstname)
# run_cmd(cmd)
shutil.copyfile(srcname, dstname)
with open(join(out, 'log.txt'), 'w') as f:
for (seq, sub, imgname) in tqdm(records, 'writing'):
f.write('{},{},{}\r\n'.format(seq, sub, imgname))
def split_directories(root, out):
with open(join(out, 'log.txt'), 'r') as f:
records = f.readlines()
for record in tqdm(records):
seq, sub, imgname = record.strip().split(',')
imgname = imgname.replace('.jpg', '.json')
srcname = join(out, 'annots', 'merge', '{}+{}+{}'.format(seq, sub, imgname))
dstname = join(root, seq, 'annots', sub, imgname)
if not os.path.exists(os.path.dirname(dstname)):
os.makedirs(os.path.dirname(dstname))
shutil.copyfile(srcname, dstname)
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('root', type=str)
parser.add_argument('out', type=str)
parser.add_argument('--merge', action='store_true')
parser.add_argument('--split', action='store_true')
args = parser.parse_args()
if args.merge:
merge_directories(args.root, args.out)
if args.split:
split_directories(args.root, args.out)