mirror of
https://github.com/KichangKim/DeepDanbooru.git
synced 2026-10-27 04:19:58 +08:00
118 lines
3.5 KiB
Python
118 lines
3.5 KiB
Python
import random
|
|
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
import tensorflow_io as tfio
|
|
|
|
import deepdanbooru as dd
|
|
|
|
|
|
class DatasetWrapper:
|
|
"""
|
|
Wrapper class for data pipelining/augmentation.
|
|
"""
|
|
|
|
def __init__(
|
|
self, inputs, tags, width, height, scale_range, rotation_range, shift_range
|
|
):
|
|
self.inputs = inputs
|
|
self.width = width
|
|
self.height = height
|
|
self.scale_range = scale_range
|
|
self.rotation_range = rotation_range
|
|
self.shift_range = shift_range
|
|
self.tag_all_array = np.array(tags)
|
|
|
|
def get_dataset(self, minibatch_size):
|
|
dataset = tf.data.Dataset.from_tensor_slices(self.inputs)
|
|
dataset = dataset.map(
|
|
self.map_load_image, num_parallel_calls=tf.data.AUTOTUNE
|
|
)
|
|
dataset = dataset.ignore_errors()
|
|
dataset = dataset.map(
|
|
self.map_transform_image_and_label,
|
|
num_parallel_calls=tf.data.AUTOTUNE,
|
|
)
|
|
dataset = dataset.batch(minibatch_size)
|
|
dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
|
|
# dataset = dataset.apply(
|
|
# tf.data.experimental.prefetch_to_device('/device:GPU:0'))
|
|
|
|
return dataset
|
|
|
|
def map_load_image(self, image_path, tag_string):
|
|
image_raw = tf.io.read_file(image_path)
|
|
try:
|
|
image = tf.io.decode_png(image_raw, channels=3)
|
|
except:
|
|
image = tfio.image.decode_webp(image_raw)
|
|
image = tfio.experimental.color.rgba_to_rgb(image)
|
|
|
|
if self.scale_range:
|
|
pre_scale = self.scale_range[1]
|
|
else:
|
|
pre_scale = 1.0
|
|
|
|
size = (int(self.height * pre_scale), int(self.width * pre_scale))
|
|
|
|
image = tf.image.resize(
|
|
image,
|
|
size=size,
|
|
method=tf.image.ResizeMethod.AREA,
|
|
preserve_aspect_ratio=True,
|
|
)
|
|
|
|
return (image, tag_string)
|
|
|
|
def map_transform_image_and_label(self, image, tag_string):
|
|
return tf.py_function(
|
|
self.map_transform_image_and_label_py,
|
|
(image, tag_string),
|
|
(tf.float32, tf.float32),
|
|
)
|
|
|
|
def map_transform_image_and_label_py(self, image, tag_string):
|
|
# transform image
|
|
image = image.numpy()
|
|
|
|
if self.scale_range:
|
|
scale = random.uniform(self.scale_range[0], self.scale_range[1]) * (
|
|
1.0 / self.scale_range[1]
|
|
)
|
|
else:
|
|
scale = None
|
|
|
|
if self.rotation_range:
|
|
rotation = random.uniform(self.rotation_range[0], self.rotation_range[1])
|
|
else:
|
|
rotation = None
|
|
|
|
if self.shift_range:
|
|
shift_x = random.uniform(self.shift_range[0], self.shift_range[1])
|
|
shift_y = random.uniform(self.shift_range[0], self.shift_range[1])
|
|
shift = (shift_x, shift_y)
|
|
else:
|
|
shift = None
|
|
|
|
image = dd.image.transform_and_pad_image(
|
|
image=image,
|
|
target_width=self.width,
|
|
target_height=self.height,
|
|
rotation=rotation,
|
|
scale=scale,
|
|
shift=shift,
|
|
)
|
|
|
|
image = image / 255.0 # normalize to 0~1
|
|
# image = image.astype(np.float32)
|
|
|
|
# transform tag
|
|
tag_string = tag_string.numpy().decode()
|
|
tag_array = np.array(tag_string.split(" "))
|
|
|
|
labels = np.where(np.isin(self.tag_all_array, tag_array), 1, 0).astype(
|
|
np.float32
|
|
)
|
|
|
|
return (image, labels)
|