File size: 1,499 Bytes
032e687 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
import os
import json
import argparse
import math
def main(args):
left_sam_images = []
for dir_name in os.listdir(args.sa1b_path):
if dir_name[:3] != 'sa_':
continue
for file in os.listdir(os.path.join(args.sa1b_path, dir_name)):
if not file.endswith('jpg'):
continue
file_name = file.split('.')[0]
if not os.path.exists(os.path.join(args.sa1b_path, dir_name, file_name+'.json')):
continue
if os.path.exists(os.path.join(args.work_dir, dir_name, file_name+'.json')):
continue
left_sam_images.append(os.path.join(args.sa1b_path, dir_name, file))
if not os.path.exists(args.task_file_path):
os.makedirs(args.task_file_path)
mini_batch = math.ceil(len(left_sam_images) / args.num_splits)
for ii in range(args.num_splits):
with open(os.path.join(args.task_file_path, f"rank{ii}.json"), mode='w') as savef:
savef.write(json.dumps(left_sam_images[ii*mini_batch:(ii+1)*mini_batch]))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--num_splits", type=int, default=8)
parser.add_argument("--sa1b_path", type=str, default="./data/sam/")
parser.add_argument("--work_dir", type=str, default="./work_dirs/object_level")
parser.add_argument("--task_file_path", type=str, default="./work_dirs/object_level_task")
args = parser.parse_args()
main(args) |