Spaces:
Running
Running
| import os | |
| import tensorflow as tf | |
| # Define Training variable | |
| BUFFER_SIZE = 400 | |
| BATCH_SIZE = 32 | |
| IMG_WIDTH = 256 | |
| IMG_HEIGHT = 256 | |
| AUTOTUNE = tf.data.AUTOTUNE | |
| def load_images(image_file): | |
| image = tf.io.read_file(image_file) | |
| image = tf.image.decode_jpeg(image) | |
| image = tf.cast(image, tf.float32) | |
| return image | |
| def resize(content_image, style_image, height, width): | |
| content_image = tf.image.resize(content_image, [height, width], | |
| method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) | |
| if style_image is not None: | |
| style_image = tf.image.resize(style_image, [height, width], | |
| method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) | |
| return content_image, style_image | |
| def random_crop(content_image, style_image): | |
| stacked_image = tf.stack([content_image, style_image], axis=0) | |
| cropped_image = tf.image.random_crop( | |
| stacked_image, size=[2, IMG_HEIGHT, IMG_WIDTH, 3]) | |
| return cropped_image[0], cropped_image[1] | |
| def normalize(content_image, style_image): | |
| content_image = (content_image / 127.5) - 1 | |
| if style_image is not None: | |
| style_image = (style_image / 127.5) - 1 | |
| return content_image, style_image | |
| def random_jitter(content_image, style_image): | |
| # resizing to 286 x 286 x 3 | |
| content_image, style_image = resize(content_image, style_image, 286, 286) | |
| # randomly cropping to 256 x 256 x 3 | |
| content_image, style_image = random_crop(content_image, style_image) | |
| if tf.random.uniform(()) > 0.5: | |
| # random mirroring | |
| content_image = tf.image.flip_left_right(content_image) | |
| style_image = tf.image.flip_left_right(style_image) | |
| return content_image, style_image | |
| def preprocess_train_image(content_path, style_path): | |
| content_image = load_images(content_path) | |
| style_image = load_images(style_path) | |
| content_image, style_image = random_jitter(content_image, style_image) | |
| content_image, style_image = normalize(content_image, style_image) | |
| return content_image, style_image | |
| def preprocess_test_image(content_path, style_path=None): | |
| content_image = load_images(content_path) | |
| if style_path is None: | |
| style_image = None | |
| else: | |
| style_image = load_images(style_path) | |
| content_image, style_image = resize(content_image, style_image, | |
| IMG_HEIGHT, IMG_WIDTH) | |
| content_image, style_image = normalize(content_image, style_image) | |
| if style_image is None: | |
| return content_image | |
| else: | |
| return content_image, style_image | |
| def create_image_loader(path): | |
| images = os.listdir(path) | |
| images = [os.path.join(path, p) for p in images] | |
| images.sort() | |
| # split the images in train and test | |
| total_images = len(images) | |
| train = images[: int(0.9 * total_images)] | |
| test = images[int(0.9 * total_images):] | |
| # Build the tf.data datasets. | |
| train_ds = tf.data.Dataset.from_tensor_slices(train) | |
| test_ds = tf.data.Dataset.from_tensor_slices(test) | |
| return train_ds, test_ds | |
| def data_loader(content_path="../data/face", style_path="../data/comics"): | |
| train_content_ds, test_content_ds = create_image_loader(content_path) | |
| train_style_ds, test_style_ds = create_image_loader(style_path) | |
| # Zipping the style and content datasets. | |
| train_ds = ( | |
| tf.data.Dataset.zip((train_content_ds, train_style_ds)) | |
| .map(preprocess_train_image) | |
| .shuffle(BUFFER_SIZE) | |
| .batch(BATCH_SIZE) | |
| .prefetch(AUTOTUNE) | |
| ) | |
| test_ds = ( | |
| tf.data.Dataset.zip((test_content_ds, test_style_ds)) | |
| .map(preprocess_test_image) | |
| .batch(BATCH_SIZE) | |
| .prefetch(AUTOTUNE) | |
| ) | |
| return train_ds, test_ds | |