How to effectively infer artifact removal cnn models? - tensorflow

Question

So earlier, I trained the SNET model on images present here for the purpose of artifact removal. The hyperparameters with which I have trained the model are mentioned below:
So SNET has eight convolution-based heads that would output a certain size of patch in the image. The input would be the same patch that has jpeg artifacts artificially added into them. The predicted and the ground truth patch would be compared with one another, and then MSE loss between them is backpropagated.
The hyperparameters, that I used:
learning_rate = 0.0001, min_learning_rate = 0.000001 for exponential scheduler
optimizer = Adam
loss metric = MSE(mean squared error)
patch_size = 48 x 48
Training batch-size = 16
Evaluation metrics:
PSNR: This is used for quality measurement between the original and reconstructed image.
SSIM: A common metric is to quantify the difference in the values of each of the corresponding pixels between the sample and the reference images.
So while training the images, I was evaluating the model's performance at every step. So there are eight 48 x 48 patches of the same image, since there 8 heads in the model. The model was trained for 100 iterations. Outputs of random patches from 998 to 999 iterations are given below:
The first image is the artifact-added image, the second is the predicted one, third is the ground truth image.
After training the model, I had to test them on bigger images with a larger context(that has an object). So instead of resizing the images to 48 x 48, I divided them into patches of 48 x 48, therefore only tested on images that have width and height values that are multiples of 48. But the problem is that, though the image has high psnr and ssim values, there is a fine gap between patch to patch as shown below:
Is there a way to efficiently tackle this issue? Please suggest, open to any kind of feedback.
Below is the code for model that I used:
import tensorflow as tf
from PIL import Image
import cv2
import tensorflow as tf
import numpy as np
import os
def MSE(input,target):
#return tf.reduce_sum(tf.reduce_mean(tf.abs(input - target),axis=0))
return tf.reduce_mean(tf.abs(input - target))
initializer = tf.initializers.VarianceScaling()
def EncoderBlock(x, activation = tf.keras.layers.LeakyReLU(alpha=0.2), nf = 256):
x = tf.keras.layers.Conv2D(nf, 5, strides=1, padding='same', kernel_initializer=initializer, use_bias=True)(x)
x = activation(x)
x = tf.keras.layers.Conv2D(nf, 5, strides=1, padding='same', kernel_initializer=initializer, use_bias=True)(x)
x = activation(x)
return x
def DecoderBlock(x, activation = tf.keras.layers.LeakyReLU(alpha=0.2), nf = 256):
x = tf.keras.layers.Conv2D(nf, 5, strides=1, padding='same', kernel_initializer=initializer, use_bias=True)(x)
x = activation(x)
x = tf.keras.layers.Conv2D(3, 5, strides=1, padding='same', kernel_initializer=initializer, use_bias=True)(x)
x = activation(x)
return x
def ConvolutionalUnit(x, structure_type = 'classic', activation = tf.keras.layers.LeakyReLU(alpha=0.2), nf = 256):
residual = x
if structure_type == "classic":
x = tf.keras.layers.Conv2D(nf, 5, strides=1, padding='same', kernel_initializer=initializer, use_bias=True)(x)
x = activation(x)
x = tf.keras.layers.Add()([x, residual])
elif structure_type == "advanced":
x = tf.keras.layers.Conv2D(nf, 5, strides=1, padding='same', kernel_initializer=initializer, use_bias=True)(x)
x = activation(x)
x = tf.keras.layers.Conv2D(nf, 5, strides=1, padding='same', kernel_initializer=initializer, use_bias=True)(x)
x = tf.keras.layers.Lambda(lambda x: x * 0.1)(x)
x = tf.keras.layers.Add()([x, residual])
return x
def S_Net(channels = 3, num_metrics=8 , structure_type='advanced', nf = 256):
inputs = tf.keras.layers.Input(shape=[None, None, channels])
encoder = EncoderBlock(inputs, nf = nf)
convolution_units = []
decoders = []
for i in range(num_metrics):
convolution_units.append(ConvolutionalUnit( convolution_units[-1] if len(convolution_units)>0 else ConvolutionalUnit(encoder, nf=nf), structure_type = structure_type, nf=nf))
decoders.append(DecoderBlock(convolution_units[-1],nf=nf))
return tf.keras.Model(inputs=[inputs], outputs=decoders)