train.py

import glob
import os

import cv2
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
from PIL import Image
from skimage.metrics import peak_signal_noise_ratio, structural_similarity
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.utils import plot_model

AUTOTUNE = tf.data.experimental.AUTOTUNE

print(tf.__version__)

# Download DIV2K from TF Datasets
div2k_data = tfds.image.Div2k(config="bicubic_x4")
div2k_data.download_and_prepare()

# Taking train and validation data from div2k_data object
train_div2k = div2k_data.as_dataset(split="train", as_supervised=True)
val_div2k = div2k_data.as_dataset(split="validation", as_supervised=True)


train_hr_dir = './data/train/HR'
train_lr_dir = './data/train/LR'
val_hr_dir = './data/validation/HR'
val_lr_dir = './data/validation/LR'


def load_custom_images(hr_dir, lr_dir):
    def image_generator():
        for filename in os.listdir(hr_dir):
            hr_image_path = os.path.join(hr_dir, filename)
            lr_image_path = os.path.join(lr_dir, filename)

            hr_image = Image.open(hr_image_path)
            lr_image = Image.open(lr_image_path)
            if (hr_image.mode != "RGB"):
                hr_image = hr_image.convert("RGB")

            if (lr_image.mode != "RGB"):
                lr_image = lr_image.convert("RGB")
            yield (tf.keras.preprocessing.image.img_to_array(lr_image),
                   tf.keras.preprocessing.image.img_to_array(hr_image))

    return tf.data.Dataset.from_generator(image_generator,
                                          output_types=(tf.uint8, tf.uint8),
                                          output_shapes=((None, None, 3), (None, None, 3)))


# Load custom images as a dataset
custom_train = load_custom_images(train_hr_dir, train_lr_dir)
custom_val = load_custom_images(val_hr_dir, val_lr_dir)


# # Load and decode images
# def load_and_decode_image(filename):
#     try:
#         file_name = tf.strings.split(filename, sep='/')[-1]
#         print("file_name",file_name)
#         image = tf.io.read_file(filename)
#         image = tf.image.decode_png(image, channels=3)
#         return tf.cast(image, tf.uint8)
#     except tf.errors.InvalidArgumentError as e:
#         print(f"Error decoding image file: {e}")
#         return None

# custom_train_hr = custom_train_hr.map(load_and_decode_image)
# custom_train_lr = custom_train_lr.map(load_and_decode_image)
# custom_val_hr = custom_val_hr.map(load_and_decode_image)
# custom_val_lr = custom_val_lr.map(load_and_decode_image)

# # Zip HR and LR images together to form input-output pairs
# custom_train = tf.data.Dataset.zip((custom_train_lr, custom_train_hr))
# custom_val = tf.data.Dataset.zip((custom_val_lr, custom_val_hr))

# # Concatenate custom dataset with div2k dataset
train = train_div2k.concatenate(custom_train)
val = val_div2k.concatenate(custom_val)

# custom_train2 = custom_train.repeat().take(180)


# Get the number of images in the training and validation datasets
# num_train_images = custom_train.cardinality().numpy()
# num_val_images = custom_val.cardinality().numpy()
# # print("Number of training samples: %d" % tf.data.experimental.cardinality(train))
# print("Number of validation samples: %d" %
#       custom_train2.cardinality().numpy())
# print("train", len(train))
# print("Number of custom training samples: %d" %
#       tf.data.experimental.cardinality(custom_train))
# print("Number of custom validation samples: %d" %
#       tf.data.experimental.cardinality(custom_val))
# print("Number of DIV2K training samples: %d" %
#       tf.data.experimental.cardinality(train_div2k))
# print("Number of DIV2K validation samples: %d" %
#       tf.data.experimental.cardinality(val_div2k))
# Optional: cache the dataset for faster access
train_cache = train
val_cache = val

# print("traincache",train_cache)
# print("validiation cache",val_cache)


def flip_left_right(lowres_img, highres_img):
    """Flips Images to left and right."""
    # Outputs random values from a uniform distribution in between 0 to 1
    rn = tf.random.uniform(shape=(), maxval=1)
    # If rn is less than 0.5 it returns original lowres_img and highres_img
    # If rn is greater than 0.5 it returns flipped image
    return tf.cond(
        rn < 0.5,
        lambda: (lowres_img, highres_img),
        lambda: (
            tf.image.flip_left_right(lowres_img),
            tf.image.flip_left_right(highres_img),
        ),
    )


def random_rotate(lowres_img, highres_img):
    """Rotates Images by 90 degrees."""
    # Outputs random values from uniform distribution in between 0 to 4

    rn = tf.random.uniform(shape=(), maxval=4, dtype=tf.int32)
    # Here rn signifies number of times the image(s) are rotated by 90 degrees
    return tf.image.rot90(lowres_img, rn), tf.image.rot90(highres_img, rn)


def random_crop(lowres_img, highres_img, hr_crop_size=96, scale=4):
    """Crop images.

    low resolution images: 24x24
    high resolution images: 96x96
    """
    lowres_crop_size = hr_crop_size // scale  # 96//4=24
    lowres_img_shape = tf.shape(lowres_img)[:2]  # (height,width)

    lowres_width = tf.random.uniform(
        shape=(), maxval=lowres_img_shape[1] - lowres_crop_size + 1, dtype=tf.int32
    )
    lowres_height = tf.random.uniform(
        shape=(), maxval=lowres_img_shape[0] - lowres_crop_size + 1, dtype=tf.int32
    )

    highres_width = lowres_width * scale
    highres_height = lowres_height * scale

    lowres_img_cropped = lowres_img[
        lowres_height: lowres_height + lowres_crop_size,
        lowres_width: lowres_width + lowres_crop_size,
    ]  # 24x24
    highres_img_cropped = highres_img[
        highres_height: highres_height + hr_crop_size,
        highres_width: highres_width + hr_crop_size,
    ]  # 96x96
    print("highres_img_cropped", highres_img_cropped.shape)
    return lowres_img_cropped, highres_img_cropped


def dataset_object(dataset_cache, training=True):
    ds = dataset_cache
    ds = ds.map(
        lambda lowres, highres: random_crop(lowres, highres, scale=4),
        num_parallel_calls=AUTOTUNE,
    )

    if training:
        ds = ds.map(random_rotate, num_parallel_calls=AUTOTUNE)
        ds = ds.map(flip_left_right, num_parallel_calls=AUTOTUNE)
    # Batching Data
    ds = ds.batch(16)

    if training:
        # Repeating Data, so that cardinality if dataset becomes infinte
        ds = ds.repeat()
    # prefetching allows later images to be prepared while the current image is being processed
    ds = ds.prefetch(buffer_size=AUTOTUNE)
    return ds


train_ds = dataset_object(train_cache, training=True)
val_ds = dataset_object(val_cache, training=False)

lowres, highres = next(iter(train_ds))

# High Resolution Images
plt.figure(figsize=(10, 10))

for i in range(1):
    ax = plt.subplot(3, 3, i + 1)
    plt.imshow(highres[i].numpy().astype("uint8"))
    plt.title(highres[i].shape)
    plt.axis("off")
    plt.show()

# # Low Resolution Images
plt.figure(figsize=(10, 10))
for i in range(1):
    ax = plt.subplot(3, 3, i + 1)
    plt.imshow(lowres[i].numpy().astype("uint8"))
    plt.title(lowres[i].shape)
    plt.axis("off")
    plt.show()


def PSNR(super_resolution, high_resolution):
    """Compute the peak signal-to-noise ratio, measures quality of image."""
    # Max value of pixel is 255
    psnr_value = tf.image.psnr(
        high_resolution, super_resolution, max_val=255)[0]
    return psnr_value


class DeepMH(tf.keras.Model):
    def train_step(self, data):
        # Unpack the data. Its structure depends on your model and
        # on what you pass to `fit()`.
        x, y = data

        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # Forward pass
            # Compute the loss value
            # (the loss function is configured in `compile()`)
            loss = self.compiled_loss(
                y, y_pred, regularization_losses=self.losses)

        # Compute gradients
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)
        # Update weights
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))
        # Update metrics (includes the metric that tracks the loss)
        self.compiled_metrics.update_state(y, y_pred)
        # Return a dict mapping metric names to current value
        return {m.name: m.result() for m in self.metrics}

    def predict_step(self, x):
        # Adding dummy dimension using tf.expand_dims and converting to float32 using tf.cast

        x = tf.cast(tf.expand_dims(x, axis=0), tf.float32)
        # Passing low resolution image to model
        super_resolution_img = self(x, training=False)
        # Clips the tensor from min(0) to max(255)
        super_resolution_img = tf.clip_by_value(super_resolution_img, 0, 255)
        # Rounds the values of a tensor to the nearest integer
        super_resolution_img = tf.round(super_resolution_img)
        # Removes dimensions of size 1 from the shape of a tensor and converting to uint8
        super_resolution_img = tf.squeeze(
            tf.cast(super_resolution_img, tf.uint8), axis=0
        )
        return super_resolution_img


def ResBlock(inputs, num_filters):
    l1 = layers.Conv2D(num_filters, 3, padding="same",
                       activation="relu")(inputs)
    l1 = layers.Activation('relu')(l1)
    l2 = layers.Conv2D(num_filters, 3, padding="same")(l1)
    x = layers.Add()([inputs, l2])
    return x


def Upsampling(inputs, factor=2, **kwargs):
    x = layers.Conv2D(64 * (factor ** 2), 3, padding="same", **kwargs)(inputs)
    x = tf.nn.depth_to_space(x, block_size=factor)
    x = layers.Conv2D(64 * (factor ** 2), 3, padding="same", **kwargs)(x)
    x = tf.nn.depth_to_space(x, block_size=factor)
    return x


def make_model(num_filters, num_of_residual_blocks):
    input_layer = layers.Input(shape=(None, None, 3))
    x = layers.Conv2D(num_filters, 3, padding="same")(input_layer)
    x = layers.Activation('relu')(x)

    for _ in range(num_of_residual_blocks):
        x = ResBlock(x, num_filters)

    skip_connection = layers.Conv2D(
        num_filters, 3, padding="same")(input_layer)
    x = layers.Add()([skip_connection, x])
    x = Upsampling(x)
    x = layers.Activation('relu')(x)

    output_layer = layers.Conv2D(3, 3, padding="same")(x)
    return DeepMH(input_layer, output_layer)


model = make_model(num_filters=256, num_of_residual_blocks=5)

# model.summary()
# plot_model(model, to_file ='super_res.png',show_shapes=True)

# Using adam optimizer with initial learning rate as 1e-4, changing learning rate after 5000 steps to 5e-5
optim_DeepMH = keras.optimizers.Adam(
    learning_rate=keras.optimizers.schedules.PiecewiseConstantDecay(
        boundaries=[5000], values=[1e-4, 5e-5]
    )
)

# Compiling model with loss as mean absolute error(L1 Loss) and metric as psnr
model.compile(optimizer=optim_DeepMH, loss="mae", metrics=[PSNR])
# Training for more epochs will improve results
model.fit(train_ds, epochs=20, steps_per_epoch=200, validation_data=val_ds)
model.save('DeepMH.h5')
model.save('DeepMH')