Add an Example: pix2pix

examples/pix2pix/input_pipeline.py

@@ -0,0 +1,105 @@
+import tensorflow as tf
+import tensorflow_datasets as tfds
+import os
+
+_URL = 'https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/facades.tar.gz'
+
+path_to_zip = tf.keras.utils.get_file('facades.tar.gz',
+                                      origin=_URL,
+                                      extract=True)
+
+PATH = os.path.join(os.path.dirname(path_to_zip), 'facades/')
+
+BUFFER_SIZE = 400
+BATCH_SIZE = 1
+IMG_WIDTH = 256
+IMG_HEIGHT = 256
+
+
+def load(image_file):
+  image = tf.io.read_file(image_file)
+  image = tf.image.decode_jpeg(image)
+
+  w = tf.shape(image)[1]
+
+  w = w // 2
+  real_image = image[:, :w, :]
+  input_image = image[:, w:, :]
+
+  input_image = tf.cast(input_image, tf.float32)
+  real_image = tf.cast(real_image, tf.float32)
+
+  return input_image, real_image
+
+
+def resize(input_image, real_image, height, width):
+  input_image = tf.image.resize(input_image, [height, width],
+                                method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
+  real_image = tf.image.resize(real_image, [height, width],
+                               method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
+
+  return input_image, real_image
+
+
+def random_crop(input_image, real_image):
+  stacked_image = tf.stack([input_image, real_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(input_image, real_image):
+  input_image = (input_image / 127.5) - 1
+  real_image = (real_image / 127.5) - 1
+
+  return input_image, real_image
+
+
+@tf.function()
+def random_jitter(input_image, real_image):
+  # resizing to 286 x 286 x 3
+  input_image, real_image = resize(input_image, real_image, 286, 286)
+
+  # randomly cropping to 256 x 256 x 3
+  input_image, real_image = random_crop(input_image, real_image)
+
+  if tf.random.uniform(()) > 0.5:
+    # random mirroring
+    input_image = tf.image.flip_left_right(input_image)
+    real_image = tf.image.flip_left_right(real_image)
+
+  return input_image, real_image
+
+
+def load_image_train(image_file):
+  input_image, real_image = load(image_file)
+  input_image, real_image = random_jitter(input_image, real_image)
+  input_image, real_image = normalize(input_image, real_image)
+
+  return input_image, real_image
+
+
+def load_image_test(image_file):
+  input_image, real_image = load(image_file)
+  input_image, real_image = resize(input_image, real_image,
+                                   IMG_HEIGHT, IMG_WIDTH)
+  input_image, real_image = normalize(input_image, real_image)
+
+  return input_image, real_image
+
+
+# Input Pipeline
+def get_dataset():
+  train_dataset = tf.data.Dataset.list_files(PATH + 'train/*.jpg')
+  train_dataset = train_dataset.map(load_image_train,
+                                    num_parallel_calls=tf.data.experimental.AUTOTUNE)
+  train_dataset = train_dataset.shuffle(BUFFER_SIZE)
+  train_dataset = train_dataset.batch(BATCH_SIZE)
+  train_dataset = tfds.as_numpy(train_dataset)
+
+  test_dataset = tf.data.Dataset.list_files(PATH + 'test/*.jpg')
+  test_dataset = test_dataset.map(load_image_test)
+  test_dataset = test_dataset.batch(BATCH_SIZE)
+  test_dataset = tfds.as_numpy(test_dataset)
+  return train_dataset, test_dataset

examples/pix2pix/models.py

@@ -0,0 +1,84 @@
+import jax
+import flax
+
+import jax.numpy as jnp
+
+OUTPUT_CHANNELS = 3
+
+
+class DownSample(flax.nn.Module):
+  def apply(self, x, features, size, apply_batchnorm=True):
+    x = flax.nn.Conv(x, features=features, kernel_size=(size, size),
+                     strides=(2, 2), padding='SAME', bias=False)
+    if apply_batchnorm:
+      x = flax.nn.BatchNorm(x)
+    x = flax.nn.leaky_relu(x)
+    return x
+
+
+class UpSample(flax.nn.Module):
+  def apply(self, x, features, size, apply_dropout=True):
+    x = flax.nn.ConvTranspose(x, features=features,
+                              kernel_size=(size, size), strides=(2, 2),
+                              padding='SAME', bias=False)
+    x = flax.nn.BatchNorm(x)
+    if apply_dropout:
+      x = flax.nn.dropout(x, 0.5)
+    x = flax.nn.relu(x)
+    return x
+
+
+down_list = [[64, 4, False],
+             [128, 4],
+             [256, 4],
+             [512, 4],
+             [512, 4],
+             [512, 4],
+             [512, 4],
+             [512, 4]]
+
+up_list = [[512, 4, True],
+           [512, 4, True],
+           [512, 4, True],
+           [512, 4],
+           [256, 4],
+           [128, 4],
+           [64, 4]]
+
+
+class Generator(flax.nn.Module):
+  def apply(self, x):
+    skips = []
+    for down in down_list:
+      x = DownSample(x, *down)
+      skips.append(x)
+
+    skips = list(reversed(skips[:-1]))
+    for up, skip in zip(up_list, skips):
+      x = UpSample(x, *up)
+      x = jnp.concatenate((x, skip))
+
+    x = flax.nn.ConvTranspose(x, features=OUTPUT_CHANNELS,
+                              kernel_size=(4, 4), strides=(2, 2),
+                              padding='SAME')
+    x = flax.nn.tanh(x)
+    return x
+
+
+class Discriminator(flax.nn.Module):
+  def apply(self, x):
+    x = DownSample(x, 64, 4, False)
+    x = DownSample(x, 128, 4)
+    x = DownSample(x, 256, 4)
+
+    x = jnp.pad(x, 1)  # padding with zeros
+
+    x = flax.nn.Conv(x, 512, kernel_size=(4, 4), strides=(1, 1), bias=False)
+    x = flax.nn.BatchNorm(x)
+    x = flax.nn.leaky_relu(x)
+
+    x = jnp.pad(x, 1)
+
+    x = flax.nn.Conv(x, 1, kernel_size=(4, 4), strides=(1, 1))
+
+    return x