keras-io/EDSR

1import tensorflow as tf2import matplotlib.pyplot as plt3from tensorflow import keras4from tensorflow.keras import layers5import gradio as gr 6 7# Define EDSR custom model8 9class EDSRModel(tf.keras.Model):10    def train_step(self, data):11        # Unpack the data. Its structure depends on your model and12        # on what you pass to `fit()`.13        x, y = data14 15        with tf.GradientTape() as tape:16            y_pred = self(x, training=True)  # Forward pass17            # Compute the loss value18            # (the loss function is configured in `compile()`)19            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)20 21        # Compute gradients22        trainable_vars = self.trainable_variables23        gradients = tape.gradient(loss, trainable_vars)24        # Update weights25        self.optimizer.apply_gradients(zip(gradients, trainable_vars))26        # Update metrics (includes the metric that tracks the loss)27        self.compiled_metrics.update_state(y, y_pred)28        # Return a dict mapping metric names to current value29        return {m.name: m.result() for m in self.metrics}30 31    def predict_step(self, x):32        # Adding dummy dimension using tf.expand_dims and converting to float32 using tf.cast33        x = tf.cast(tf.expand_dims(x, axis=0), tf.float32)34        # Passing low resolution image to model35        super_resolution_img = self(x, training=False)36        # Clips the tensor from min(0) to max(255)37        super_resolution_img = tf.clip_by_value(super_resolution_img, 0, 255)38        # Rounds the values of a tensor to the nearest integer39        super_resolution_img = tf.round(super_resolution_img)40        # Removes dimensions of size 1 from the shape of a tensor and converting to uint841        super_resolution_img = tf.squeeze(42            tf.cast(super_resolution_img, tf.uint8), axis=043        )44        return super_resolution_img45 46 47# Residual Block48def ResBlock(inputs):49    x = layers.Conv2D(64, 3, padding="same", activation="relu")(inputs)50    x = layers.Conv2D(64, 3, padding="same")(x)51    x = layers.Add()([inputs, x])52    return x53 54 55# Upsampling Block56def Upsampling(inputs, factor=2, **kwargs):57    x = layers.Conv2D(64 * (factor ** 2), 3, padding="same", **kwargs)(inputs)58    x = tf.nn.depth_to_space(x, block_size=factor)59    x = layers.Conv2D(64 * (factor ** 2), 3, padding="same", **kwargs)(x)60    x = tf.nn.depth_to_space(x, block_size=factor)61    return x62 63 64def make_model(num_filters, num_of_residual_blocks):65    # Flexible Inputs to input_layer66    input_layer = layers.Input(shape=(None, None, 3))67    # Scaling Pixel Values68    x = layers.Rescaling(scale=1.0 / 255)(input_layer)69    x = x_new = layers.Conv2D(num_filters, 3, padding="same")(x)70 71    # 16 residual blocks72    for _ in range(num_of_residual_blocks):73        x_new = ResBlock(x_new)74 75    x_new = layers.Conv2D(num_filters, 3, padding="same")(x_new)76    x = layers.Add()([x, x_new])77 78    x = Upsampling(x)79    x = layers.Conv2D(3, 3, padding="same")(x)80 81    output_layer = layers.Rescaling(scale=255)(x)82    return EDSRModel(input_layer, output_layer)83 84 85# Define PSNR metric86 87def PSNR(super_resolution, high_resolution):88    """Compute the peak signal-to-noise ratio, measures quality of image."""89    # Max value of pixel is 25590    psnr_value = tf.image.psnr(high_resolution, super_resolution, max_val=255)[0]91    return psnr_value92 93custom_objects = {"EDSRModel":EDSRModel}94 95with keras.utils.custom_object_scope(custom_objects):96    new_model = keras.models.load_model("./trained.h5", custom_objects={'PSNR':PSNR})97 98 99def process_image(img):100    lowres = tf.convert_to_tensor(img, dtype=tf.uint8)101    lowres = tf.image.random_crop(lowres, (150, 150, 3))102    preds = new_model.predict_step(lowres)103    preds = preds.numpy()104    lowres = lowres.numpy()105    return (lowres, preds)106 107image = gr.inputs.Image()108#image_out = gr.outputs.Image()109 110markdown_part = """111 112Model Link - https://huggingface.co/keras-io/EDSR113 114"""115 116examples = [["./examples/1.png"]]117 118gr.Interface(119    process_image, 120    title="EDSR - Enhanced Deep Residual Networks for Single Image Super-Resolution",121    description="SuperResolution",122    inputs = image,123    examples = examples,124    outputs = gr.Gallery(label="Outputs, First image is low res, next one is High Res",visible=True).style(grid=[2], height="auto"),125    article = markdown_part,126    interpretation='default',127    allow_flagging='never'128            ).launch(debug=True)