| """ | |
| Convert an LSUN lmdb database into a directory of images. | |
| """ | |
| import argparse | |
| import io | |
| import os | |
| from PIL import Image | |
| import lmdb | |
| import numpy as np | |
| def read_images(lmdb_path, image_size): | |
| env = lmdb.open(lmdb_path, map_size=1099511627776, max_readers=100, readonly=True) | |
| with env.begin(write=False) as transaction: | |
| cursor = transaction.cursor() | |
| for _, webp_data in cursor: | |
| img = Image.open(io.BytesIO(webp_data)) | |
| width, height = img.size | |
| scale = image_size / min(width, height) | |
| img = img.resize( | |
| (int(round(scale * width)), int(round(scale * height))), | |
| resample=Image.BOX, | |
| ) | |
| arr = np.array(img) | |
| h, w, _ = arr.shape | |
| h_off = (h - image_size) // 2 | |
| w_off = (w - image_size) // 2 | |
| arr = arr[h_off : h_off + image_size, w_off : w_off + image_size] | |
| yield arr | |
| def dump_images(out_dir, images, prefix): | |
| if not os.path.exists(out_dir): | |
| os.mkdir(out_dir) | |
| for i, img in enumerate(images): | |
| Image.fromarray(img).save(os.path.join(out_dir, f"{prefix}_{i:07d}.png")) | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--image-size", help="new image size", type=int, default=256) | |
| parser.add_argument("--prefix", help="class name", type=str, default="bedroom") | |
| parser.add_argument("lmdb_path", help="path to an LSUN lmdb database") | |
| parser.add_argument("out_dir", help="path to output directory") | |
| args = parser.parse_args() | |
| images = read_images(args.lmdb_path, args.image_size) | |
| dump_images(args.out_dir, images, args.prefix) | |
| if __name__ == "__main__": | |
| main() | |