upscaling pipeline
This commit is contained in:
+23
-160
@@ -1,172 +1,35 @@
|
||||
import tensorflow as tf
|
||||
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseWithSkip, ConvNeXtWithSkip
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
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):
|
||||
class Upscaler(nn.Module):
|
||||
def __init__(self, low_res=32, kernel_size=3):
|
||||
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(256, 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(128, 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(64, 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(32, 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'),
|
||||
])
|
||||
self.seq = nn.Sequential(
|
||||
nn.Conv2d(3, 256, kernel_size=kernel_size, padding='same'),
|
||||
nn.ReLU(),
|
||||
nn.BatchNorm2d(256),
|
||||
nn.Conv2d(256, 128, kernel_size=kernel_size, padding='same'),
|
||||
nn.ReLU(),
|
||||
nn.BatchNorm2d(128),
|
||||
nn.Conv2d(128, 64, kernel_size=kernel_size, padding='same'),
|
||||
nn.ReLU(),
|
||||
nn.BatchNorm2d(64),
|
||||
nn.Conv2d(64, 32, kernel_size=kernel_size, padding='same'),
|
||||
nn.ReLU(),
|
||||
nn.BatchNorm2d(32),
|
||||
nn.Conv2d(32, 3, kernel_size=kernel_size, padding='same'),
|
||||
nn.Tanh()
|
||||
)
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
def forward(self, inputs) -> torch.Tensor:
|
||||
img, mask = inputs
|
||||
|
||||
x = self.seq(img, training=training)
|
||||
x = self.seq(img)
|
||||
|
||||
x = x * (1 - mask[:,:,:,tf.newaxis])
|
||||
x = x * (1 - mask.unsqueeze(1))
|
||||
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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user