Spaces:
Runtime error
Runtime error
File size: 3,178 Bytes
1b2a9b1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
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")
|