| import argparse |
| import os.path as osp |
| from functools import partial |
|
|
| import mmcv |
| import numpy as np |
| from PIL import Image |
| from scipy.io import loadmat |
|
|
| AUG_LEN = 10582 |
|
|
|
|
| def convert_mat(mat_file, in_dir, out_dir): |
| data = loadmat(osp.join(in_dir, mat_file)) |
| mask = data['GTcls'][0]['Segmentation'][0].astype(np.uint8) |
| seg_filename = osp.join(out_dir, mat_file.replace('.mat', '.png')) |
| Image.fromarray(mask).save(seg_filename, 'PNG') |
|
|
|
|
| def generate_aug_list(merged_list, excluded_list): |
| return list(set(merged_list) - set(excluded_list)) |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser( |
| description='Convert PASCAL VOC annotations to mmsegmentation format') |
| parser.add_argument('devkit_path', help='pascal voc devkit path') |
| parser.add_argument('aug_path', help='pascal voc aug path') |
| 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() |
| devkit_path = args.devkit_path |
| aug_path = args.aug_path |
| nproc = args.nproc |
| if args.out_dir is None: |
| out_dir = osp.join(devkit_path, 'VOC2012', 'SegmentationClassAug') |
| else: |
| out_dir = args.out_dir |
| mmcv.mkdir_or_exist(out_dir) |
| in_dir = osp.join(aug_path, 'dataset', 'cls') |
|
|
| mmcv.track_parallel_progress( |
| partial(convert_mat, in_dir=in_dir, out_dir=out_dir), |
| list(mmcv.scandir(in_dir, suffix='.mat')), |
| nproc=nproc) |
|
|
| full_aug_list = [] |
| with open(osp.join(aug_path, 'dataset', 'train.txt')) as f: |
| full_aug_list += [line.strip() for line in f] |
| with open(osp.join(aug_path, 'dataset', 'val.txt')) as f: |
| full_aug_list += [line.strip() for line in f] |
|
|
| with open( |
| osp.join(devkit_path, 'VOC2012/ImageSets/Segmentation', |
| 'train.txt')) as f: |
| ori_train_list = [line.strip() for line in f] |
| with open( |
| osp.join(devkit_path, 'VOC2012/ImageSets/Segmentation', |
| 'val.txt')) as f: |
| val_list = [line.strip() for line in f] |
|
|
| aug_train_list = generate_aug_list(ori_train_list + full_aug_list, |
| val_list) |
| assert len(aug_train_list) == AUG_LEN, 'len(aug_train_list) != {}'.format( |
| AUG_LEN) |
|
|
| with open( |
| osp.join(devkit_path, 'VOC2012/ImageSets/Segmentation', |
| 'trainaug.txt'), 'w') as f: |
| f.writelines(line + '\n' for line in aug_train_list) |
|
|
| aug_list = generate_aug_list(full_aug_list, ori_train_list + val_list) |
| assert len(aug_list) == AUG_LEN - len( |
| ori_train_list), 'len(aug_list) != {}'.format(AUG_LEN - |
| len(ori_train_list)) |
| with open( |
| osp.join(devkit_path, 'VOC2012/ImageSets/Segmentation', 'aug.txt'), |
| 'w') as f: |
| f.writelines(line + '\n' for line in aug_list) |
|
|
| print('Done!') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
