prepare_data.py

import argparse
from io import BytesIO
import multiprocessing
from functools import partial

from PIL import Image
import lmdb
from tqdm import tqdm
from torchvision import datasets
from torchvision.transforms import functional as trans_fn


def resize_and_convert(img, size, quality=100):
    img = trans_fn.resize(img, size, Image.LANCZOS)
    img = trans_fn.center_crop(img, size)
    buffer = BytesIO()
    img.save(buffer, format='jpeg', quality=quality)
    val = buffer.getvalue()

    return val


def resize_multiple(img, sizes=(8, 16, 32, 64, 128, 256, 512, 1024), quality=100):
    imgs = []

    for size in sizes:
        imgs.append(resize_and_convert(img, size, quality))

    return imgs


def resize_worker(img_file, sizes):
    i, file, label = img_file
    img = Image.open(file)
    img = img.convert('RGB')
    out = resize_multiple(img, sizes=sizes)

    return i, out, label


def prepare(transaction, dataset, n_worker, sizes=(8, 16, 32, 64, 128, 256, 512, 1024)):
    resize_fn = partial(resize_worker, sizes=sizes)

    files = sorted(dataset.imgs, key=lambda x: x[0])
    files = [(i, file, label) for i, (file, label) in enumerate(files)]
    num_classes = len(set(label for i, f, label in files))
    total = 0

    with multiprocessing.Pool(n_worker) as pool:
        for i, imgs, label in tqdm(pool.imap_unordered(resize_fn, files), total=len(files)):
            for size, img in zip(sizes, imgs):
                key = f'{size}-{str(i).zfill(5)}'.encode('utf-8')
                transaction.put(key, img)

            key = f'label-{str(i).zfill(5)}'.encode('utf-8')
            transaction.put(key, str(label).encode('utf-8'))

            total += 1

        transaction.put('length'.encode('utf-8'), str(total).encode('utf-8'))
        transaction.put('num_classes'.encode('utf-8'), str(num_classes).encode('utf-8'))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--out', type=str)
    parser.add_argument('--n_worker', type=int, default=8)
    parser.add_argument('path', type=str)
    parser.add_argument('--labels', action='store_true', help='Add labels too')
    parser.add_argument('--sizes', nargs='+', type=int, default=(8, 16, 32, 64, 128, 256, 512, 1024))
    parser.add_argument('--min_size', type=int, default=None, help='Minimum image size to not drop it')

    args = parser.parse_args()

    if args.min_size is None:
        valid_file = None
    else:
        valid_file = lambda x: min(Image.open(x).size) >= args.min_size

    imgset = datasets.ImageFolder(args.path, is_valid_file=valid_file)

    with lmdb.open(args.out, map_size=1024 ** 4, readahead=False) as env:
        with env.begin(write=True) as txn:
            prepare(txn, imgset, args.n_worker, args.sizes)