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