| |
|
|
| import pickle |
| from tqdm import tqdm |
| import os |
| import numpy as np |
| from PIL import Image |
| import argparse |
| import lmdb |
| from torchvision import transforms |
|
|
|
|
| MAX_SIZE = 1e12 |
|
|
|
|
| def load_and_resize(root, path, imscale): |
|
|
| transf_list = [] |
| transf_list.append(transforms.Resize(imscale)) |
| transf_list.append(transforms.CenterCrop(imscale)) |
| transform = transforms.Compose(transf_list) |
|
|
| img = Image.open(os.path.join(root, path[0], path[1], path[2], path[3], path)).convert('RGB') |
| img = transform(img) |
|
|
| return img |
|
|
|
|
| def main(args): |
|
|
| parts = {} |
| datasets = {} |
| imname2pos = {'train': {}, 'val': {}, 'test': {}} |
| for split in ['train', 'val', 'test']: |
| datasets[split] = pickle.load(open(os.path.join(args.save_dir, args.suff + 'recipe1m_' + split + '.pkl'), 'rb')) |
|
|
| parts[split] = lmdb.open(os.path.join(args.save_dir, 'lmdb_'+split), map_size=int(MAX_SIZE)) |
| with parts[split].begin() as txn: |
| present_entries = [key for key, _ in txn.cursor()] |
| j = 0 |
| for i, entry in tqdm(enumerate(datasets[split])): |
| impaths = entry['images'][0:5] |
|
|
| for n, p in enumerate(impaths): |
| if n == args.maxnumims: |
| break |
| if p.encode() not in present_entries: |
| im = load_and_resize(os.path.join(args.root, 'images', split), p, args.imscale) |
| im = np.array(im).astype(np.uint8) |
| with parts[split].begin(write=True) as txn: |
| txn.put(p.encode(), im) |
| imname2pos[split][p] = j |
| j += 1 |
| pickle.dump(imname2pos, open(os.path.join(args.save_dir, 'imname2pos.pkl'), 'wb')) |
|
|
|
|
| def test(args): |
|
|
| imname2pos = pickle.load(open(os.path.join(args.save_dir, 'imname2pos.pkl'), 'rb')) |
| paths = imname2pos['val'] |
|
|
| for k, v in paths.items(): |
| path = k |
| break |
| image_file = lmdb.open(os.path.join(args.save_dir, 'lmdb_' + 'val'), max_readers=1, readonly=True, |
| lock=False, readahead=False, meminit=False) |
| with image_file.begin(write=False) as txn: |
| image = txn.get(path.encode()) |
| image = np.fromstring(image, dtype=np.uint8) |
| image = np.reshape(image, (args.imscale, args.imscale, 3)) |
| image = Image.fromarray(image.astype('uint8'), 'RGB') |
| print (np.shape(image)) |
|
|
|
|
| if __name__ == "__main__": |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument('--root', type=str, default='path/to/recipe1m', |
| help='path to the recipe1m dataset') |
| parser.add_argument('--save_dir', type=str, default='../data', |
| help='path where the lmdbs will be saved') |
| parser.add_argument('--imscale', type=int, default=256, |
| help='size of images (will be rescaled and center cropped)') |
| parser.add_argument('--maxnumims', type=int, default=5, |
| help='maximum number of images to allow for each sample') |
| parser.add_argument('--suff', type=str, default='', |
| help='id of the vocabulary to use') |
| parser.add_argument('--test_only', dest='test_only', action='store_true') |
| parser.set_defaults(test_only=False) |
| args = parser.parse_args() |
|
|
| if not args.test_only: |
| main(args) |
| test(args) |
|
|
