|
|
|
|
| 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)
|
|
|
