173 lines
5.9 KiB
Python
173 lines
5.9 KiB
Python
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):
|
|
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):
|
|
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, low_res=32, high_res=64):
|
|
super(UpscalerBasic, self).__init__()
|
|
|
|
self.low_res = low_res
|
|
self.high_res = high_res
|
|
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(low_res * (low_res//2) * 3),
|
|
])
|
|
|
|
self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3))
|
|
|
|
def call(self, inputs, training):
|
|
img, mask = inputs
|
|
|
|
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=(self.high_res, self.high_res//2), method=tf.image.ResizeMethod.LANCZOS5)
|
|
|
|
return x
|
|
|
|
|
|
|
|
|
|
class UpscalerBasicNew(tf.keras.Model):
|
|
|
|
def __init__(self, low_res=32):
|
|
super(UpscalerBasicNew, self).__init__()
|
|
|
|
self.low_res = low_res
|
|
self.flatten = tf.keras.layers.Flatten(name='upscale_enhancer_flattent')
|
|
|
|
self.seq = tf.keras.Sequential([
|
|
tf.keras.layers.Dense(4096, activation='relu', name='upscale_enhancer_dense_1'),
|
|
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'),
|
|
tf.keras.layers.Dense(2048, activation='relu', name='upscale_enhancer_dense_2'),
|
|
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'),
|
|
tf.keras.layers.Dense(2048, activation='relu', name='upscale_enhancer_dense_3'),
|
|
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'),
|
|
tf.keras.layers.Dense(low_res * (low_res//2) * 3, activation='tanh', name='upscale_enhancer_dense_final'),
|
|
])
|
|
|
|
self.reshape = tf.keras.layers.Reshape((low_res, (low_res//2), 3), name='upscale_enhancer_reshape')
|
|
|
|
def call(self, inputs, training):
|
|
img, mask = inputs
|
|
|
|
x = self.flatten(img)
|
|
x = self.seq(x, training=training)
|
|
x = self.reshape(x)
|
|
|
|
x = x * (1 - mask[:,:,:,tf.newaxis])
|
|
x = x + img
|
|
|
|
return x
|
|
|
|
|
|
class UpscalerBasicConv(tf.keras.Model):
|
|
|
|
def __init__(self, low_res=32, kernel_size=3):
|
|
super(UpscalerBasicConv, self).__init__()
|
|
|
|
self.low_res = low_res
|
|
|
|
self.seq = tf.keras.Sequential([
|
|
tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_1'),
|
|
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'),
|
|
tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_2'),
|
|
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'),
|
|
tf.keras.layers.Conv2D(16, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_3'),
|
|
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'),
|
|
tf.keras.layers.Conv2D(16, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_4'),
|
|
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_4'),
|
|
tf.keras.layers.Dense(3, activation='tanh', name='upscale_enhancer_dense_final'),
|
|
])
|
|
|
|
|
|
def call(self, inputs, training):
|
|
img, mask = inputs
|
|
|
|
x = self.seq(img, training=training)
|
|
|
|
x = x * (1 - mask[:,:,:,tf.newaxis])
|
|
x = x + img
|
|
|
|
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):
|
|
# x, mask = inputs
|
|
|
|
# x = tf.transpose(x, [0,2,3,1])
|
|
# x = self.decoder(x)
|
|
|
|
# x = x * mask[:,:,:,tf.newaxis]
|
|
|
|
# return x
|
|
|