Skip to content
Snippets Groups Projects
Select Git revision
  • master default protected
  • use_flight_altitude
  • distribute
3 results
Blame Permalink
esrgan.py 4.80 KiB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137 






from collections import OrderedDict
from functools import partial
import tensorflow as tf
from util import esrgan_utils as utils

""" Keras Models for ESRGAN
    Classes:
      RRDBNet: Generator of ESRGAN. (Residual in Residual Network)
      VGGArch: VGG28 Architecture making the Discriminator ESRGAN
"""


class RRDBNet(tf.keras.Model):
  """ Residual in Residual Network consisting of:
      - Convolution Layers
      - Residual in Residual Block as the trunk of the model
      - Pixel Shuffler layers (tf.nn.depth_to_space)
      - Upscaling Convolutional Layers

      Args:
        out_channel: number of channels of the fake output image.
        num_features (default: 32): number of filters to use in the convolutional layers.
        trunk_size (default: 3): number of Residual in Residual Blocks to form the trunk.
        growth_channel (default: 32): number of filters to use in the internal convolutional layers.
        use_bias (default: True): boolean to indicate if bias is to be used in the conv layers.
  """

  def __init__(
          self,
          out_channel,
          num_features=32,
          trunk_size=11,
          growth_channel=32,
          use_bias=True,
          first_call=True):
    super(RRDBNet, self).__init__()
    self.rrdb_block = partial(utils.RRDB, growth_channel, first_call=first_call)
    conv = partial(
        tf.keras.layers.Conv2D,
        kernel_size=[3, 3],
        strides=[1, 1],
        padding="same",
        use_bias=use_bias)
    conv_transpose = partial(
        tf.keras.layers.Conv2DTranspose,
        kernel_size=[3, 3],
        strides=[2, 2],
        padding="same",
        use_bias=use_bias)
    self.conv_first = conv(filters=num_features)
    self.rdb_trunk = tf.keras.Sequential(
        [self.rrdb_block() for _ in range(trunk_size)])
    self.conv_trunk = conv(filters=num_features)
    # Upsample
    self.upsample1 = conv_transpose(num_features)
    self.upsample2 = conv_transpose(num_features)
    self.conv_last_1 = conv(num_features)
    self.conv_last_2 = conv(out_channel)
    self.lrelu = tf.keras.layers.LeakyReLU(alpha=0.2)

  # @tf.function(
  #    input_signature=[
  #        tf.TensorSpec(shape=[None, None, None, 3],
  #                      dtype=tf.float32)])
  def call(self, inputs):
    return self.unsigned_call(inputs)

  def unsigned_call(self, input_):
    feature = self.lrelu(self.conv_first(input_))
    trunk = self.conv_trunk(self.rdb_trunk(feature))
    feature = trunk + feature
    feature = self.lrelu(
            self.upsample1(feature))
    feature = self.lrelu(
          self.upsample2(feature))
    feature = self.lrelu(self.conv_last_1(feature))
    out = self.conv_last_2(feature)
    return out


class VGGArch(tf.keras.Model):
  """ Keras Model for VGG28 Architecture needed to form
      the discriminator of the architecture.
      Args:
        output_shape (default: 1): output_shape of the generator
        num_features (default: 64): number of features to be used in the convolutional layers
                                    a factor of 2**i will be multiplied as per the need
        use_bias (default: True): Boolean to indicate whether to use biases for convolution layers
  """

  def __init__(
          self,
          batch_size=8,
          output_shape=1,
          num_features=64,
          use_bias=False):

    super(VGGArch, self).__init__()
    conv = partial(
        tf.keras.layers.Conv2D,
        kernel_size=[3, 3], use_bias=use_bias, padding="same")
    batch_norm = partial(tf.keras.layers.BatchNormalization)
    def no_batch_norm(x): return x
    self._lrelu = tf.keras.layers.LeakyReLU(alpha=0.2)
    self._dense_1 = tf.keras.layers.Dense(100)
    self._dense_2 = tf.keras.layers.Dense(output_shape)
    self._conv_layers = OrderedDict()
    self._batch_norm = OrderedDict()
    self._conv_layers["conv_0_0"] = conv(filters=num_features, strides=1)
    self._conv_layers["conv_0_1"] = conv(filters=num_features, strides=2)
    self._batch_norm["bn_0_1"] = batch_norm()
    for i in range(1, 4):
      for j in range(1, 3):
        self._conv_layers["conv_%d_%d" % (i, j)] = conv(
            filters=num_features * (2**i), strides=j)
        self._batch_norm["bn_%d_%d" % (i, j)] = batch_norm()

  def call(self, inputs):
    return self.unsigned_call(inputs)
  def unsigned_call(self, input_):

    features = self._lrelu(self._conv_layers["conv_0_0"](input_))
    features = self._lrelu(
        self._batch_norm["bn_0_1"](
            self._conv_layers["conv_0_1"](features)))
    # VGG Trunk
    for i in range(1, 4):
      for j in range(1, 3):
        features = self._lrelu(
            self._batch_norm["bn_%d_%d" % (i, j)](
                self._conv_layers["conv_%d_%d" % (i, j)](features)))

    flattened = tf.keras.layers.Flatten()(features)
    dense = self._lrelu(self._dense_1(flattened))
    out = self._dense_2(dense)
    return out