(dataset, splits, output_dir, generate_mask=False)
| 80 | |
| 81 | |
| 82 | def prepare_dataset(dataset, splits, output_dir, generate_mask=False): |
| 83 | ann_path = os.path.join(output_dir, 'anns', dataset) |
| 84 | mask_path = os.path.join(output_dir, 'masks', dataset) |
| 85 | if not os.path.exists(ann_path): |
| 86 | os.makedirs(ann_path) |
| 87 | if not os.path.exists(mask_path): |
| 88 | os.makedirs(mask_path) |
| 89 | |
| 90 | for split in splits: |
| 91 | dataset_array = [] |
| 92 | ref_ids = refer.getRefIds(split=split) |
| 93 | print('Processing split:{} - Len: {}'.format(split, np.alen(ref_ids))) |
| 94 | for i in tqdm(ref_ids): |
| 95 | ref_dict = {} |
| 96 | |
| 97 | refs = refer.Refs[i] |
| 98 | bboxs = refer.getRefBox(i) |
| 99 | sentences = refs['sentences'] |
| 100 | image_urls = refer.loadImgs(image_ids=refs['image_id'])[0] |
| 101 | cat = cat_process(refs['category_id']) |
| 102 | image_urls = image_urls['file_name'] |
| 103 | if dataset == 'refclef' and image_urls in [ |
| 104 | '19579.jpg', '17975.jpg', '19575.jpg' |
| 105 | ]: |
| 106 | continue |
| 107 | box_info = bbox_process(bboxs) |
| 108 | |
| 109 | ref_dict['bbox'] = box_info |
| 110 | ref_dict['cat'] = cat |
| 111 | ref_dict['segment_id'] = i |
| 112 | ref_dict['img_name'] = image_urls |
| 113 | |
| 114 | if generate_mask: |
| 115 | cv2.imwrite(os.path.join(mask_path, |
| 116 | str(i) + '.png'), |
| 117 | refer.getMask(refs)['mask'] * 255) |
| 118 | |
| 119 | sent_dict = [] |
| 120 | for i, sent in enumerate(sentences): |
| 121 | sent_dict.append({ |
| 122 | 'idx': i, |
| 123 | 'sent_id': sent['sent_id'], |
| 124 | 'sent': sent['sent'].strip() |
| 125 | }) |
| 126 | |
| 127 | ref_dict['sentences'] = sent_dict |
| 128 | ref_dict['sentences_num'] = len(sent_dict) |
| 129 | |
| 130 | dataset_array.append(ref_dict) |
| 131 | print('Dumping json file...') |
| 132 | with open(os.path.join(output_dir, 'anns', dataset, split + '.json'), |
| 133 | 'w') as f: |
| 134 | json.dump(dataset_array, f) |
| 135 | |
| 136 | |
| 137 | prepare_dataset(args.dataset, splits, args.output_dir, args.generate_mask) |
no test coverage detected