custom upscaling

models/upscaler.py

@@ -0,0 +1,107 @@
+import tensorflow as tf
+from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseWithSkip, ConvNeXtWithSkip
+
+
+class Bottleneck(tf.keras.Model):
+    
+    def __init__(self, dim, kernel_size=7):
+        super(Bottleneck, self).__init__()
+
+        self.bottleneck = tf.keras.Sequential(
+                [Block(dim=dim, drop_path=0, kernel_size=kernel_size,
+                       layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
+            )
+
+    def call(self, inputs, training=None):
+        x = inputs
+
+        x = self.bottleneck(x, training=training)
+        
+        return x
+
+
+
+class Upscaler(tf.keras.Model):
+    
+    def __init__(self, bottleneck_size):
+        super(Upscaler, self).__init__()
+
+        depths = [27, 27]
+        dims = [512, 512]
+        reverse_dims = [16] + dims[:-1]
+        self.encoder = ConvNeXtWithSkip(depths=depths, dims=dims, kernel_size=3)
+
+        self.bottleneck = Bottleneck(dims[-1], kernel_size=3)
+
+        self.decoder = ConvNeXtReverseWithSkip(depths=depths, dims=reverse_dims, kernel_size=3)
+
+    def call(self, inputs, training=None):
+        img, mask = inputs
+
+        img = tf.transpose(img, [0,2,3,1])
+
+        x, skip_features = self.encoder(img, training=training)
+
+        x = self.bottleneck(x, training=training)
+
+        x = self.decoder([x, skip_features], training=training)
+
+        x = x * (1 - mask[:,:,:,tf.newaxis])
+        x = x + img
+        x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
+
+        return x
+
+
+class UpscalerBasic(tf.keras.Model):
+    
+    def __init__(self, bottleneck_size, size=32):
+        super(UpscalerBasic, self).__init__()
+
+        self.flatten = tf.keras.layers.Flatten()
+
+        self.seq = tf.keras.Sequential([
+                    tf.keras.layers.Dense(1024, activation='relu'),
+                    tf.keras.layers.BatchNormalization(),
+                    tf.keras.layers.Dense(1024, activation='relu'),
+                    tf.keras.layers.BatchNormalization(),
+                    tf.keras.layers.Dense(1024, activation='relu'),
+                    tf.keras.layers.BatchNormalization(),
+                    tf.keras.layers.Dense(size * (size//2) * 3),
+        ])
+
+        self.reshape = tf.keras.layers.Reshape((size, (size//2), 3))
+
+    def call(self, inputs, training=None):
+        img, mask = inputs
+
+        if img.shape[1] == 3:
+            img = tf.transpose(img, [0,2,3,1])
+
+        x = self.flatten(img)
+        x = self.seq(x, training=training)
+        x = self.reshape(x)
+
+        x = x * (1 - mask[:,:,:,tf.newaxis])
+        x = x + img
+        x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
+
+        return x
+
+# class Upscaler(tf.keras.Model):
+    
+#     def __init__(self, bottleneck_size):
+#         super(Upscaler, self).__init__()
+
+#         self.decoder = ConvNeXtReverse(depths=list(reversed([9, 9, 9])), dims=list(reversed([bottleneck_size, 128, 32, 16])))
+
+#     def call(self, inputs, training=None):
+#         x, mask = inputs
+
+#         x = tf.transpose(x, [0,2,3,1])
+#         x = self.decoder(x)
+
+#         x = x * mask[:,:,:,tf.newaxis]
+
+#         return x
+