def create_from_lmdb(tfrecord_dir, lmdb_dir, ratio = None, max_imgs = None, shards_num = 5):
import lmdb
print("Loading dataset %s" % lmdb_dir)
bad_imgs = 0
with lmdb.open(lmdb_dir, readonly = True).begin(write = False) as txn:
if max_imgs is None:
max_imgs = txn.stat()["entries"]
with TFRecordExporter(tfrecord_dir, max_imgs, verbose = False, shards_num = shards_num) as tfr:
for idx, (_key, value) in tqdm(enumerate(txn.cursor()), total = max_imgs):
try:
img = PIL.Image.open(io.BytesIO(value))
img = misc.crop_max_rectangle(img, ratio)
img = misc.pad_min_square(img)
pow2size = 2 ** int(np.round(np.log2(img.size[0])))
img = img.resize((pow2size, pow2size), PIL.Image.ANTIALIAS)
img = np.asarray(img)
img = img.transpose([2, 0, 1]) # HWC => CHW
tfr.add_img(img)
except:
bad_imgs += 1
pass
if tfr.curr_imgnum == max_imgs:
break
if bad_imgs > 0:
print("Couldn't read {} out of {} images".format(bad_imgs, max_imgs))
def create_from_tfds(tfrecord_dir,
dataset_name,
ratio=None,
max_imgs=None,
shards_num=5):
import tensorflow_datasets as tfds
print("Loading dataset %s" % dataset_name)
ds = tfds.load(dataset_name,
split="train",
data_dir=f"{tfrecord_dir}/tfds")
with TFRecordExporter(tfrecord_dir, 0, shards_num=shards_num) as tfr:
for i, ex in tqdm(enumerate(tfds.as_numpy(ds))):
img = PIL.Image.fromarray(ex["image"])
img = misc.crop_max_rectangle(img, ratio)
img = misc.pad_min_square(img)
pow2size = 2**int(np.round(np.log2(img.size[0])))
img = img.resize((pow2size, pow2size), PIL.Image.ANTIALIAS)
img = np.asarray(img)
img = img.transpose([2, 0, 1]) # HWC => CHW
tfr.add_img(img)
if max_imgs is not None and i > max_imgs:
break
def create_from_imgs(tfrecord_dir,
img_dir,
format="png",
shuffle=False,
ratio=None,
max_imgs=None,
shards_num=5):
print("Loading images from %s" % img_dir)
img_filenames = sorted(
glob.glob(f"{img_dir}/**/*.{format}", recursive=True))
if len(img_filenames) == 0:
error("No input images found")
if max_imgs is None:
max_imgs = len(img_filenames)
# Check image shape
img = np.asarray(PIL.Image.open(img_filenames[0]).convert("RGB"))
resolution = img.shape[0]
channels = img.shape[2] if img.ndim == 3 else 1
if channels not in [1, 3]:
error("Input images must be stored as RGB or grayscale")
# if img.shape[1] != resolution:
# error("Input images must have the same width and height")
# if resolution != 2 ** int(np.floor(np.log2(resolution))):
# error("Input image resolution must be a power-of-two")
with TFRecordExporter(tfrecord_dir,
len(img_filenames),
shards_num=shards_num) as tfr:
order = tfr.choose_shuffled_order() if shuffle else np.arange(
len(img_filenames))
for idx in trange(max_imgs):
img = PIL.Image.open(img_filenames[order[idx]]).convert("RGB")
img = misc.crop_max_rectangle(img, ratio)
img = misc.pad_min_square(img)
pow2size = 2**int(np.round(np.log2(img.size[0])))
img = img.resize((pow2size, pow2size), PIL.Image.ANTIALIAS)
img = np.asarray(img)
if channels == 1:
img = img[np.newaxis, :, :] # HW => CHW
else:
img = img.transpose([2, 0, 1]) # HWC => CHW
tfr.add_img(img)
def create_from_tfrecords(dataset_dir,
tfrecords_dir,
ratio=None,
max_imgs=None):
import tensorflow_datasets as tfds
def parse_tfrecord_tf(record):
features = tf.io.parse_single_example(
record,
features={
"shape": tf.io.FixedLenFeature([3], tf.int64),
"data": tf.io.FixedLenFeature([], tf.string)
})
data = tf.io.decode_raw(features["data"], tf.uint8)
data = tf.reshape(data, features["shape"])
data = tf.transpose(data, [1, 2, 0])
return data
print("Loading dataset %s" % tfrecords_dir)
maxres_file = sorted(
glob.glob(os.path.join(tfrecords_dir, "*.tfrecords1of*")))[-1]
tffiles = sorted(glob.glob(re.sub("1of.*", "*", maxres_file)))
dataset = tf.data.TFRecordDataset(tffiles)
dataset = dataset.map(parse_tfrecord_tf, num_parallel_calls=4)
with DatasetExporter(dataset_dir, 0) as tfr:
for i, img in tqdm(enumerate(tfds.as_numpy(dataset)), total=max_imgs):
img = PIL.Image.fromarray(img)
img = misc.crop_max_rectangle(img, ratio)
img = misc.pad_min_square(img)
pow2size = 2**int(np.round(np.log2(img.size[0])))
img = img.resize((pow2size, pow2size), PIL.Image.ANTIALIAS)
img = np.asarray(img)
img = img.transpose([2, 0, 1]) # HWC => CHW
tfr.add_img(img)
if max_imgs is not None and i > max_imgs:
break