|
import argparse |
|
import os.path as osp |
|
|
|
import mmcv |
|
from cityscapesscripts.preparation.json2labelImg import json2labelImg |
|
|
|
|
|
def convert_json_to_label(json_file): |
|
label_file = json_file.replace('_polygons.json', '_labelTrainIds.png') |
|
json2labelImg(json_file, label_file, 'trainIds') |
|
|
|
|
|
def parse_args(): |
|
parser = argparse.ArgumentParser( |
|
description='Convert Cityscapes annotations to TrainIds') |
|
parser.add_argument('cityscapes_path', help='cityscapes data path') |
|
parser.add_argument('--gt-dir', default='gtFine', type=str) |
|
parser.add_argument('-o', '--out-dir', help='output path') |
|
parser.add_argument( |
|
'--nproc', default=1, type=int, help='number of process') |
|
args = parser.parse_args() |
|
return args |
|
|
|
|
|
def main(): |
|
args = parse_args() |
|
cityscapes_path = args.cityscapes_path |
|
out_dir = args.out_dir if args.out_dir else cityscapes_path |
|
mmcv.mkdir_or_exist(out_dir) |
|
|
|
gt_dir = osp.join(cityscapes_path, args.gt_dir) |
|
|
|
poly_files = [] |
|
for poly in mmcv.scandir(gt_dir, '_polygons.json', recursive=True): |
|
poly_file = osp.join(gt_dir, poly) |
|
poly_files.append(poly_file) |
|
if args.nproc > 1: |
|
mmcv.track_parallel_progress(convert_json_to_label, poly_files, |
|
args.nproc) |
|
else: |
|
mmcv.track_progress(convert_json_to_label, poly_files) |
|
|
|
split_names = ['train', 'val', 'test'] |
|
|
|
for split in split_names: |
|
filenames = [] |
|
for poly in mmcv.scandir( |
|
osp.join(gt_dir, split), '_polygons.json', recursive=True): |
|
filenames.append(poly.replace('_gtFine_polygons.json', '')) |
|
with open(osp.join(out_dir, f'{split}.txt'), 'w') as f: |
|
f.writelines(f + '\n' for f in filenames) |
|
|
|
|
|
if __name__ == '__main__': |
|
main() |
|
|