Upload New File

model.py

+from tensorflow import keras
+import tensorflow as tf
+
+
+class FastSRGAN(object):
+    """SRGAN for fast super resolution."""
+
+    def __init__(self, args):
+        """
+        Initializes the Mobile SRGAN class.
+        Args:
+            args: CLI arguments that dictate how to build the model.
+        Returns:
+            None
+        """
+        self.hr_height = args.hr_size
+        self.hr_width = args.hr_size
+        self.lr_height = self.hr_height // 4  # Low resolution height
+        self.lr_width = self.hr_width // 4  # Low resolution width
+        self.lr_shape = (self.lr_height, self.lr_width, 3)
+        self.hr_shape = (self.hr_height, self.hr_width, 3)
+        self.iterations = 0
+
+        # Number of inverted residual blocks in the mobilenet generator
+        self.n_residual_blocks = 6
+
+        # Define a learning rate decay schedule.
+        self.gen_schedule = keras.optimizers.schedules.ExponentialDecay(
+            args.lr,
+            decay_steps=100000,
+            decay_rate=0.1,
+            staircase=True
+        )
+
+        self.disc_schedule = keras.optimizers.schedules.ExponentialDecay(
+            args.lr * 5,  # TTUR - Two Time Scale Updates
+            decay_steps=100000,
+            decay_rate=0.1,
+            staircase=True
+        )
+
+        self.gen_optimizer = keras.optimizers.Adam(learning_rate=self.gen_schedule)
+        self.disc_optimizer = keras.optimizers.Adam(learning_rate=self.disc_schedule)
+
+        # We use a pre-trained VGG19 model to extract image features from the high resolution
+        # and the generated high resolution images and minimize the mse between them
+        self.vgg = self.build_vgg()
+        self.vgg.trainable = False
+
+        # Calculate output shape of D (PatchGAN)
+        patch = int(self.hr_height / 2 ** 4)
+        self.disc_patch = (patch, patch, 1)
+
+        # Number of filters in the first layer of G and D
+        self.gf = 32  # Realtime Image Enhancement GAN Galteri et al.
+        self.df = 32
+
+        # Build and compile the discriminator
+        self.discriminator = self.build_discriminator()
+
+        # Build and compile the generator for pretraining.
+        self.generator = self.build_generator()
+
+    @tf.function
+    def content_loss(self, hr, sr):
+        sr = keras.applications.vgg19.preprocess_input(((sr + 1.0) * 255) / 2.0)
+        hr = keras.applications.vgg19.preprocess_input(((hr + 1.0) * 255) / 2.0)
+        sr_features = self.vgg(sr) / 12.75
+        hr_features = self.vgg(hr) / 12.75
+        return tf.keras.losses.MeanSquaredError()(hr_features, sr_features)
+
+    def build_vgg(self):
+        """
+        Builds a pre-trained VGG19 model that outputs image features extracted at the
+        third block of the model
+        """
+        # Get the vgg network. Extract features from Block 5, last convolution.
+        vgg = keras.applications.VGG19(weights="imagenet", input_shape=self.hr_shape, include_top=False)
+        vgg.trainable = False
+        for layer in vgg.layers:
+            layer.trainable = False
+
+        # Create model and compile
+        model = keras.models.Model(inputs=vgg.input, outputs=vgg.get_layer("block5_conv4").output)
+
+        return model
+
+    def build_generator(self):
+        """Build the generator that will do the Super Resolution task.
+        Based on the Mobilenet design. Idea from Galteri et al."""
+
+        def _make_divisible(v, divisor, min_value=None):
+                if min_value is None:
+                    min_value = divisor
+                new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+                # Make sure that round down does not go down by more than 10%.
+                if new_v < 0.9 * v:
+                    new_v += divisor
+                return new_v
+
+        def residual_block(inputs, filters, block_id, expansion=6, stride=1, alpha=1.0):
+            """Inverted Residual block that uses depth wise convolutions for parameter efficiency.
+            Args:
+                inputs: The input feature map.
+                filters: Number of filters in each convolution in the block.
+                block_id: An integer specifier for the id of the block in the graph.
+                expansion: Channel expansion factor.
+                stride: The stride of the convolution.
+                alpha: Depth expansion factor.
+            Returns:
+                x: The output of the inverted residual block.
+            """
+            channel_axis = 1 if keras.backend.image_data_format() == 'channels_first' else -1
+
+            in_channels = keras.backend.int_shape(inputs)[channel_axis]
+            pointwise_conv_filters = int(filters * alpha)
+            pointwise_filters = _make_divisible(pointwise_conv_filters, 8)
+            x = inputs
+            prefix = 'block_{}_'.format(block_id)
+
+            if block_id:
+                # Expand
+                x = keras.layers.Conv2D(expansion * in_channels,
+                                        kernel_size=1,
+                                        padding='same',
+                                        use_bias=True,
+                                        activation=None,
+                                        name=prefix + 'expand')(x)
+                x = keras.layers.BatchNormalization(axis=channel_axis,
+                                                    epsilon=1e-3,
+                                                    momentum=0.999,
+                                                    name=prefix + 'expand_BN')(x)
+                x = keras.layers.Activation('relu', name=prefix + 'expand_relu')(x)
+            else:
+                prefix = 'expanded_conv_'
+
+            # Depthwise
+            x = keras.layers.DepthwiseConv2D(kernel_size=3,
+                                             strides=stride,
+                                             activation=None,
+                                             use_bias=True,
+                                             padding='same' if stride == 1 else 'valid',
+                                             name=prefix + 'depthwise')(x)
+            x = keras.layers.BatchNormalization(axis=channel_axis,
+                                                epsilon=1e-3,
+                                                momentum=0.999,
+                                                name=prefix + 'depthwise_BN')(x)
+
+            x = keras.layers.Activation('relu', name=prefix + 'depthwise_relu')(x)
+
+            # Project
+            x = keras.layers.Conv2D(pointwise_filters,
+                                    kernel_size=1,
+                                    padding='same',
+                                    use_bias=True,
+                                    activation=None,
+                                    name=prefix + 'project')(x)
+            x = keras.layers.BatchNormalization(axis=channel_axis,
+                                                epsilon=1e-3,
+                                                momentum=0.999,
+                                                name=prefix + 'project_BN')(x)
+
+            if in_channels == pointwise_filters and stride == 1:
+                return keras.layers.Add(name=prefix + 'add')([inputs, x])
+            return x
+
+        def deconv2d(layer_input):
+            """Upsampling layer to increase height and width of the input.
+            Uses PixelShuffle for upsampling.
+            Args:
+                layer_input: The input tensor to upsample.
+            Returns:
+                u: Upsampled input by a factor of 2.
+            """
+            u = keras.layers.UpSampling2D(size=2, interpolation='bilinear')(layer_input)
+            u = keras.layers.Conv2D(self.gf, kernel_size=3, strides=1, padding='same')(u)
+            u = keras.layers.PReLU(shared_axes=[1, 2])(u)
+            return u
+
+        # Low resolution image input
+        img_lr = keras.Input(shape=self.lr_shape)
+
+        # Pre-residual block
+        c1 = keras.layers.Conv2D(self.gf, kernel_size=3, strides=1, padding='same')(img_lr)
+        c1 = keras.layers.BatchNormalization()(c1)
+        c1 = keras.layers.PReLU(shared_axes=[1, 2])(c1)
+
+        # Propogate through residual blocks
+        r = residual_block(c1, self.gf, 0)
+        for idx in range(1, self.n_residual_blocks):
+            r = residual_block(r, self.gf, idx)
+
+        # Post-residual block
+        c2 = keras.layers.Conv2D(self.gf, kernel_size=3, strides=1, padding='same')(r)
+        c2 = keras.layers.BatchNormalization()(c2)
+        c2 = keras.layers.Add()([c2, c1])
+        
+        # Upsampling
+        u1 = deconv2d(c2)
+        u2 = deconv2d(u1)
+
+        # Generate high resolution output
+        gen_hr = keras.layers.Conv2D(3, kernel_size=3, strides=1, padding='same', activation='tanh')(u2)
+
+        return keras.models.Model(img_lr, gen_hr)
+
+    def build_discriminator(self):
+        """Builds a discriminator network based on the SRGAN design."""
+
+        def d_block(layer_input, filters, strides=1, bn=True):
+            """Discriminator layer block.
+            Args:
+                layer_input: Input feature map for the convolutional block.
+                filters: Number of filters in the convolution.
+                strides: The stride of the convolution.
+                bn: Whether to use batch norm or not.
+            """
+            d = keras.layers.Conv2D(filters, kernel_size=3, strides=strides, padding='same')(layer_input)
+            if bn:
+                d = keras.layers.BatchNormalization(momentum=0.8)(d)
+            d = keras.layers.LeakyReLU(alpha=0.2)(d)
+                
+            return d
+
+        # Input img
+        d0 = keras.layers.Input(shape=self.hr_shape)
+
+        d1 = d_block(d0, self.df, bn=False)
+        d2 = d_block(d1, self.df, strides=2)
+        d3 = d_block(d2, self.df)
+        d4 = d_block(d3, self.df, strides=2)
+        d5 = d_block(d4, self.df * 2)
+        d6 = d_block(d5, self.df * 2, strides=2)
+        d7 = d_block(d6, self.df * 2)
+        d8 = d_block(d7, self.df * 2, strides=2)
+
+        validity = keras.layers.Conv2D(1, kernel_size=1, strides=1, activation='sigmoid', padding='same')(d8)
+
+        return keras.models.Model(d0, validity)