| import argparse |
| from io import BytesIO |
| import multiprocessing |
| from functools import partial |
|
|
| import os |
| 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, resample, quality=100): |
| img = trans_fn.resize(img, size, resample) |
| 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=(128, 256, 512, 1024), resample=Image.LANCZOS, quality=100 |
| ): |
| imgs = [] |
|
|
| for size in sizes: |
| imgs.append(resize_and_convert(img, size, resample, quality)) |
|
|
| return imgs |
|
|
|
|
| def resize_worker(img_file, sizes, resample): |
| i, file = img_file |
| img = Image.open(file) |
| img = img.convert("RGB") |
| out = resize_multiple(img, sizes=sizes, resample=resample) |
|
|
| return i, out |
|
|
|
|
| def prepare( |
| env, dataset, n_worker, sizes=(128, 256, 512, 1024), resample=Image.LANCZOS |
| ): |
| resize_fn = partial(resize_worker, sizes=sizes, resample=resample) |
|
|
| files = sorted(dataset.imgs, key=lambda x: x[0]) |
| files = [(i, file) for i, (file, label) in enumerate(files)] |
| total = 0 |
|
|
| with multiprocessing.Pool(n_worker) as pool: |
| for i, imgs in tqdm(pool.imap_unordered(resize_fn, files)): |
| for size, img in zip(sizes, imgs): |
| key = f"{size}-{str(i).zfill(5)}".encode("utf-8") |
|
|
| with env.begin(write=True) as txn: |
| txn.put(key, img) |
|
|
| total += 1 |
|
|
| with env.begin(write=True) as txn: |
| txn.put("length".encode("utf-8"), str(total).encode("utf-8")) |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description="Preprocess images for model training") |
| parser.add_argument("--out", type=str, help="filename of the result lmdb dataset") |
| parser.add_argument( |
| "--size", |
| type=str, |
| default="128,256,512,1024", |
| help="resolutions of images for the dataset", |
| ) |
| parser.add_argument( |
| "--n_worker", |
| type=int, |
| default=8, |
| help="number of workers for preparing dataset", |
| ) |
| parser.add_argument( |
| "--resample", |
| type=str, |
| default="lanczos", |
| help="resampling methods for resizing images", |
| ) |
| parser.add_argument("path", type=str, help="path to the image dataset") |
|
|
| args = parser.parse_args() |
|
|
| if not os.path.exists(args.out): |
| os.makedirs(args.out) |
| |
| resample_map = {"lanczos": Image.LANCZOS, "bilinear": Image.BILINEAR} |
| resample = resample_map[args.resample] |
|
|
| sizes = [int(s.strip()) for s in args.size.split(",")] |
|
|
| print(f"Make dataset of image sizes:", ", ".join(str(s) for s in sizes)) |
|
|
| imgset = datasets.ImageFolder(args.path) |
|
|
| with lmdb.open(args.out, map_size=1024 ** 4, readahead=False) as env: |
| prepare(env, imgset, args.n_worker, sizes=sizes, resample=resample) |
|
|
