Spaces:
Runtime error
Runtime error
| import PIL | |
| import os | |
| import numpy as np | |
| import argparse | |
| import lmdb | |
| import sys | |
| import cv2 | |
| sys.path.append(".") | |
| from swapae.data.image_folder import make_dataset | |
| def create_lmdb_from_images(opt): | |
| paths = sorted(make_dataset(opt.input)) | |
| output_dir = opt.output | |
| print('Extracting images to "%s"' % output_dir) | |
| if not os.path.isdir(output_dir): | |
| os.makedirs(output_dir) | |
| # initialize lmdb | |
| output_dir = opt.output | |
| lmdb_env = lmdb.open(output_dir, map_size=1099511627776) | |
| with lmdb_env.begin(write=True) as txn: | |
| for idx, image_path in enumerate(paths): | |
| if idx % 10 == 0: | |
| print('%d\r' % idx, end='', flush=True) | |
| img = PIL.Image.open(image_path).convert('RGB') | |
| img = np.asarray(img) | |
| img = cv2.imencode('.png', img)[1].tostring() | |
| txn.put(image_path.encode('ascii'), img) | |
| def create_lmdb_from_tfrecords(opt): | |
| # initialize tensorflow | |
| assert opt.stylegan_codebase_path is not None | |
| sys.path.append(opt.stylegan_codebase_path) | |
| import dnnlib.tflib as tflib | |
| from training import dataset | |
| import tensorflow as tf | |
| tfrecord_dir = opt.input | |
| print('Loading dataset "%s"' % tfrecord_dir) | |
| tflib.init_tf({'gpu_options.allow_growth': True}) | |
| dset = dataset.TFRecordDataset(tfrecord_dir, max_label_size=0, repeat=False, shuffle_mb=0) | |
| tflib.init_uninitialized_vars() | |
| # initialize lmdb | |
| output_dir = opt.output | |
| lmdb_env = lmdb.open(output_dir, map_size=1099511627776) | |
| print('Extracting images to "%s"' % output_dir) | |
| if not os.path.isdir(output_dir): | |
| os.makedirs(output_dir) | |
| idx = 0 | |
| with lmdb_env.begin(write=True) as txn: | |
| while True: | |
| idx += 1 | |
| if idx % 10 == 0: | |
| print('%d\r' % idx, end='', flush=True) | |
| try: | |
| images, _labels = dset.get_minibatch_np(1) | |
| except tf.errors.OutOfRangeError: | |
| break | |
| if images.shape[1] == 1: | |
| img = PIL.Image.fromarray(images[0][0], 'L').convert('RGB') | |
| else: | |
| img = PIL.Image.fromarray(images[0].transpose(1, 2, 0), 'RGB') | |
| img = np.asarray(img) | |
| img = cv2.imencode('.png', img)[1].tostring() | |
| imagekey = "%08d" % idx | |
| txn.put(imagekey.encode('ascii'), img) | |
| if __name__ == "__main__": | |
| os.environ["CUDA_VISIBLE_DEVICES"] = "-1" | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("mode", | |
| choices=("create_lmdb_from_images", | |
| "create_lmdb_from_tfrecords", | |
| )) | |
| parser.add_argument("--input", help="input path") | |
| parser.add_argument("--output", help="input path") | |
| parser.add_argument("--stylegan_codebase_path", | |
| help="path to stylegan codebase. Path to git clone https://github.com/NVlabs/stylegan.git") | |
| opt = parser.parse_args() | |
| if opt.mode == "create_lmdb_from_images": | |
| create_lmdb_from_images(opt) | |
| elif opt.mode == "create_lmdb_from_tfrecords": | |
| create_lmdb_from_tfrecords(opt) | |
| print("Finished") | |