import tensorflow as tf import matplotlib.pyplot as plt from tensorflow import keras from tensorflow.keras import layers import gradio as gr # Define EDSR custom model class EDSRModel(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 # Residual Block def ResBlock(inputs): x = layers.Conv2D(64, 3, padding="same", activation="relu")(inputs) x = layers.Conv2D(64, 3, padding="same")(x) x = layers.Add()([inputs, x]) return x # Upsampling Block 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): # Flexible Inputs to input_layer input_layer = layers.Input(shape=(None, None, 3)) # Scaling Pixel Values x = layers.Rescaling(scale=1.0 / 255)(input_layer) x = x_new = layers.Conv2D(num_filters, 3, padding="same")(x) # 16 residual blocks for _ in range(num_of_residual_blocks): x_new = ResBlock(x_new) x_new = layers.Conv2D(num_filters, 3, padding="same")(x_new) x = layers.Add()([x, x_new]) x = Upsampling(x) x = layers.Conv2D(3, 3, padding="same")(x) output_layer = layers.Rescaling(scale=255)(x) return EDSRModel(input_layer, output_layer) # Define PSNR metric 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 custom_objects = {"EDSRModel":EDSRModel} with keras.utils.custom_object_scope(custom_objects): new_model = keras.models.load_model("./trained.h5", custom_objects={'PSNR':PSNR}) def process_image(img): lowres = tf.convert_to_tensor(img, dtype=tf.uint8) lowres = tf.image.random_crop(lowres, (150, 150, 3)) preds = new_model.predict_step(lowres) preds = preds.numpy() lowres = lowres.numpy() return (lowres, preds) image = gr.inputs.Image() image_out = gr.outputs.Image() markdown_part = """ Model Link - https://huggingface.co/keras-io/EDSR """ examples = [["examples/1.png"]] gr.Interface( process_image, title="EDSR - Enhanced Deep Residual Networks for Single Image Super-Resolution", description="SuperResolution", inputs = image, examples = examples, outputs = gr.Gallery(label="Outputs, First image is low res, next one is High Res",visible=True), article = markdown_part, interpretation='default', allow_flagging='never', cache_examples=True ).launch(debug=True)