| |
| import os |
| import json |
| import argparse |
| from PIL import Image |
| import numpy as np |
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--ann', default='datasets/cc3m/Train_GCC-training.tsv') |
| parser.add_argument('--save_image_path', default='datasets/cc3m/training/') |
| parser.add_argument('--cat_info', default='datasets/lvis/lvis_v1_val.json') |
| parser.add_argument('--out_path', default='datasets/cc3m/train_image_info.json') |
| parser.add_argument('--not_download_image', action='store_true') |
| args = parser.parse_args() |
| categories = json.load(open(args.cat_info, 'r'))['categories'] |
| images = [] |
| if not os.path.exists(args.save_image_path): |
| os.makedirs(args.save_image_path) |
| f = open(args.ann) |
| for i, line in enumerate(f): |
| cap, path = line[:-1].split('\t') |
| print(i, cap, path) |
| if not args.not_download_image: |
| os.system( |
| 'wget {} -O {}/{}.jpg'.format( |
| path, args.save_image_path, i + 1)) |
| try: |
| img = Image.open( |
| open('{}/{}.jpg'.format(args.save_image_path, i + 1), "rb")) |
| img = np.asarray(img.convert("RGB")) |
| h, w = img.shape[:2] |
| except: |
| continue |
| image_info = { |
| 'id': i + 1, |
| 'file_name': '{}.jpg'.format(i + 1), |
| 'height': h, |
| 'width': w, |
| 'captions': [cap], |
| } |
| images.append(image_info) |
| data = {'categories': categories, 'images': images, 'annotations': []} |
| for k, v in data.items(): |
| print(k, len(v)) |
| print('Saving to', args.out_path) |
| json.dump(data, open(args.out_path, 'w')) |
|
|
