reaarrange files
This commit is contained in:
+207
@@ -0,0 +1,207 @@
|
||||
import tensorflow as tf
|
||||
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
|
||||
|
||||
class Bottleneck(tf.keras.Model):
|
||||
|
||||
def __init__(self, dim, latent_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)]
|
||||
)
|
||||
|
||||
self.flatten = tf.keras.layers.Flatten()
|
||||
self.dense = tf.keras.layers.Dense(latent_dim)
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.bottleneck(x, training=training)
|
||||
x = self.flatten(x)
|
||||
x = self.dense(x, training=training)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class VariationalEncoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, latent_dim, depths, dims, kernel_size=7):
|
||||
super(VariationalEncoder, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
|
||||
self.flatten = tf.keras.layers.Flatten()
|
||||
self.mean = tf.keras.layers.Dense(latent_dim)
|
||||
self.log_var = tf.keras.layers.Dense(latent_dim)
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.flatten(x)
|
||||
mean = self.mean(x)
|
||||
log_var = self.log_var(x)
|
||||
|
||||
return mean, log_var
|
||||
|
||||
|
||||
class AEEncoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, latent_dim, depths, dims, bottleneck_size=32, kernel_size=7):
|
||||
super(AEEncoder, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
self.bottleneck = Bottleneck(bottleneck_size, latent_dim, kernel_size)
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
if len(inputs) == 2:
|
||||
x, _ = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.bottleneck(x, training=training)
|
||||
|
||||
if len(inputs) == 2:
|
||||
return x, inputs[1]
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
class Decoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False, upscale=False):
|
||||
super(Decoder, self).__init__()
|
||||
|
||||
self.dense = tf.keras.layers.Dense(tf.reduce_prod(input_size))
|
||||
self.reshape = tf.keras.layers.Reshape(input_size)
|
||||
self.upscale = upscale
|
||||
|
||||
if not is_label:
|
||||
self.decoder = ConvNeXtReverse(depths=list(reversed(depths)), dims=list(reversed(dims)), kernel_size=kernel_size)
|
||||
else:
|
||||
self.decoder = ConvNeXtReverseLabels(depths=list(reversed(depths)), dims=list(reversed(dims)), nr_of_classes=7)
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
if len(inputs) == 2:
|
||||
x, mask = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.dense(x)
|
||||
x = self.reshape(x)
|
||||
x = self.decoder(x, training=training)
|
||||
|
||||
if self.upscale:
|
||||
x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
|
||||
|
||||
|
||||
if len(inputs) == 2:
|
||||
if len(mask.shape) == 2:
|
||||
x = x * mask[tf.newaxis,:,:,tf.newaxis]
|
||||
else:
|
||||
x = x * mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
|
||||
class VariationalEncoderNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, depths, dims, kernel_size=7):
|
||||
super(VariationalEncoderNew, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
|
||||
self.mean_branch = tf.keras.Sequential(
|
||||
[Block(dim=64, drop_path=0, kernel_size=kernel_size,
|
||||
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
|
||||
)
|
||||
|
||||
self.var_branch = tf.keras.Sequential(
|
||||
[Block(dim=64, drop_path=0, kernel_size=kernel_size,
|
||||
layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)]
|
||||
)
|
||||
|
||||
self.flatten_mean = tf.keras.layers.Flatten()
|
||||
self.flatten_var = tf.keras.layers.Flatten()
|
||||
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
mean = self.mean_branch(x)
|
||||
log_var = self.var_branch(x)
|
||||
|
||||
mean = self.flatten_mean(mean)
|
||||
log_var = self.flatten_var(log_var)
|
||||
|
||||
return mean, log_var
|
||||
|
||||
|
||||
|
||||
class AEEncoderNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, depths, dims, kernel_size=7, activation='tanh'):
|
||||
super(AEEncoderNew, self).__init__()
|
||||
|
||||
self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
|
||||
self.act = tf.keras.layers.Activation(activation)
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
if len(inputs) == 2:
|
||||
x, _ = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.act(x)
|
||||
|
||||
if len(inputs) == 2:
|
||||
return x, inputs[1]
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
class DecoderNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, depths, dims, kernel_size=7, is_label=False):
|
||||
super(DecoderNew, self).__init__()
|
||||
|
||||
self.is_label = is_label
|
||||
|
||||
if not is_label:
|
||||
self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size)
|
||||
else:
|
||||
self.decoder = ConvNeXtReverseLabels(depths=depths, dims=dims, nr_of_classes=7)
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
if len(inputs) == 2:
|
||||
x, mask = inputs
|
||||
else:
|
||||
x = inputs
|
||||
|
||||
x = self.decoder(x, training=training)
|
||||
|
||||
if self.is_label:
|
||||
x = x * mask[:,:,:,tf.newaxis]
|
||||
|
||||
background = tf.ones((mask.shape[0],512,256,1)) * (1 - mask[:,:,:,tf.newaxis])
|
||||
x = tf.concat((background, x), axis=-1)
|
||||
|
||||
x = tf.math.softmax(x)
|
||||
|
||||
if len(inputs) == 2:
|
||||
if len(mask.shape) == 2:
|
||||
x = x * mask[tf.newaxis,:,:,tf.newaxis]
|
||||
else:
|
||||
x = x * mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,38 @@
|
||||
import tensorflow as tf
|
||||
from models.base import AEEncoder, Decoder, AEEncoderNew, DecoderNew
|
||||
|
||||
|
||||
class AutoEncoder(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(AutoEncoder, self).__init__()
|
||||
|
||||
self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size'])
|
||||
self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'], upscale=config['upscale_output'])
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.decoder(x, training=training)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class AutoEncoderNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(AutoEncoderNew, self).__init__()
|
||||
|
||||
self.encoder : AEEncoderNew = AEEncoderNew(config['encoder_depths'], config['encoder_dims'], config['kernel_size'], config['activation'])
|
||||
self.decoder : DecoderNew = DecoderNew(config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
x = inputs
|
||||
|
||||
x = self.encoder(x, training=training)
|
||||
x = self.decoder(x, training=training)
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,306 @@
|
||||
from unittest import skip
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
|
||||
GlobalAveragePooling2D, Layer,
|
||||
LayerNormalization, Flatten, Reshape, UpSampling2D)
|
||||
|
||||
def drop_path(inputs, drop_prob, is_training):
|
||||
# borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py
|
||||
if (not is_training) or (drop_prob == 0.):
|
||||
return inputs
|
||||
|
||||
# Compute keep_prob
|
||||
keep_prob = 1.0 - drop_prob
|
||||
|
||||
# Compute drop_connect tensor
|
||||
random_tensor = keep_prob
|
||||
shape = (tf.shape(inputs)[0],) + (1,) * \
|
||||
(len(tf.shape(inputs)) - 1)
|
||||
random_tensor += tf.random.uniform(shape, dtype=inputs.dtype)
|
||||
binary_tensor = tf.floor(random_tensor)
|
||||
output = tf.math.divide(inputs, keep_prob) * binary_tensor
|
||||
return output
|
||||
|
||||
|
||||
class DropPath(tf.keras.layers.Layer):
|
||||
# borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py
|
||||
def __init__(self, drop_prob=None):
|
||||
super().__init__()
|
||||
self.drop_prob = drop_prob
|
||||
|
||||
def call(self, x, training):
|
||||
return drop_path(x, self.drop_prob, training)
|
||||
|
||||
|
||||
class Block(Layer):
|
||||
r""" ConvNeXt Block. There are two equivalent implementations:
|
||||
(1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
|
||||
(2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
|
||||
We use (2) as we find it slightly faster in PyTorch
|
||||
|
||||
Args:
|
||||
dim (int): Number of input channels.
|
||||
drop_path (float): Stochastic depth rate. Default: 0.0
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
"""
|
||||
|
||||
def __init__(self, dim, kernel_size=7, drop_path=0., layer_scale_init_value=1e-6, prefix=''):
|
||||
super().__init__()
|
||||
self.dwconv = DepthwiseConv2D(
|
||||
kernel_size=kernel_size, padding='same') # depthwise conv
|
||||
self.norm = LayerNormalization(epsilon=1e-6)
|
||||
# pointwise/1x1 convs, implemented with linear layers
|
||||
self.pwconv1 = Dense(4 * dim)
|
||||
self.act = tf.keras.activations.gelu
|
||||
self.pwconv2 = Dense(dim)
|
||||
self.drop_path = DropPath(drop_path)
|
||||
self.dim = dim
|
||||
self.layer_scale_init_value = layer_scale_init_value
|
||||
self.prefix = prefix
|
||||
|
||||
def build(self, input_shape):
|
||||
self.gamma = tf.Variable(
|
||||
initial_value=self.layer_scale_init_value * tf.ones((self.dim)),
|
||||
trainable=True,
|
||||
name=f'{self.prefix}/gamma')
|
||||
self.built = True
|
||||
|
||||
def call(self, x, training):
|
||||
input = x
|
||||
x = self.dwconv(x, training=training)
|
||||
# x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
|
||||
x = self.norm(x, training=training)
|
||||
x = self.pwconv1(x, training=training)
|
||||
x = self.act(x)
|
||||
x = self.pwconv2(x, training=training)
|
||||
if self.gamma is not None:
|
||||
x = self.gamma * x
|
||||
# x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
|
||||
|
||||
x = input + self.drop_path(x, training=training)
|
||||
return x
|
||||
|
||||
|
||||
class ConvNeXt(tf.keras.Model):
|
||||
r""" ConvNeXt
|
||||
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
|
||||
https://arxiv.org/pdf/2201.03545.pdf
|
||||
|
||||
Args:
|
||||
num_classes (int): Number of classes for classification head. Default: 1000
|
||||
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
|
||||
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
|
||||
include_top (bool): whether to add head or just use it as feature extractor. Default: True
|
||||
drop_path_rate (float): Stochastic depth rate. Default: 0.
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
|
||||
"""
|
||||
|
||||
def __init__(self, num_classes=1000, kernel_size=7,
|
||||
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], include_top=True,
|
||||
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
|
||||
):
|
||||
super().__init__()
|
||||
self.include_top = include_top
|
||||
self.downsample_layers = [] # stem and 3 intermediate downsampling conv layers
|
||||
# stem = tf.keras.Sequential([
|
||||
# Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
|
||||
# LayerNormalization(epsilon=1e-6)]
|
||||
# )
|
||||
stem = [Conv2D(dims[0], kernel_size=4, strides=4, padding='same'), LayerNormalization(epsilon=1e-6)]
|
||||
self.downsample_layers.append(stem)
|
||||
for i in range(len(depths) - 1):
|
||||
# downsample_layer = tf.keras.Sequential([
|
||||
# LayerNormalization(epsilon=1e-6),
|
||||
# Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
|
||||
# )
|
||||
downsample_layer = [LayerNormalization(epsilon=1e-6), Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
|
||||
self.downsample_layers.append(downsample_layer)
|
||||
|
||||
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
|
||||
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
|
||||
cur = 0
|
||||
for i in range(len(depths)):
|
||||
stage = tf.keras.Sequential(
|
||||
[Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size,
|
||||
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
|
||||
)
|
||||
self.stages.append(stage)
|
||||
cur += depths[i]
|
||||
|
||||
if self.include_top:
|
||||
self.avg = GlobalAveragePooling2D()
|
||||
self.norm = LayerNormalization(epsilon=1e-6) # final norm layer
|
||||
self.head = Dense(num_classes)
|
||||
else:
|
||||
self.avg = None
|
||||
self.norm = None
|
||||
self.head = None
|
||||
|
||||
def forward_features(self, x, training):
|
||||
for i in range(len(self.downsample_layers)):
|
||||
# x = self.downsample_layers[i](x)
|
||||
x = self.downsample_layers[i][0](x, training=training)
|
||||
x = self.downsample_layers[i][1](x, training=training)
|
||||
x = self.stages[i](x, training=training)
|
||||
return x
|
||||
|
||||
def call(self, x, training):
|
||||
x = self.forward_features(x, training=training)
|
||||
return x
|
||||
|
||||
|
||||
class ConvNeXtReverse(tf.keras.Model):
|
||||
r""" ConvNeXt
|
||||
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
|
||||
https://arxiv.org/pdf/2201.03545.pdf
|
||||
|
||||
Args:
|
||||
num_classes (int): Number of classes for classification head. Default: 1000
|
||||
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
|
||||
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
|
||||
include_top (bool): whether to add head or just use it as feature extractor. Default: True
|
||||
drop_path_rate (float): Stochastic depth rate. Default: 0.
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], kernel_size=7,
|
||||
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
|
||||
):
|
||||
super().__init__()
|
||||
self.upsample_layers = [] # stem and 3 intermediate downsampling conv layers
|
||||
stem = tf.keras.Sequential([
|
||||
Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
|
||||
# UpSampling2D((4, 4), interpolation='nearest'),
|
||||
# Conv2D(dims[0], kernel_size=4, strides=1, padding='same'),
|
||||
# Dense(dims[0]),
|
||||
LayerNormalization(epsilon=1e-6)]
|
||||
)
|
||||
self.upsample_layers.append(stem)
|
||||
for i in range(len(depths) - 1):
|
||||
downsample_layer = tf.keras.Sequential([
|
||||
LayerNormalization(epsilon=1e-6),
|
||||
Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')
|
||||
# UpSampling2D((2,2), interpolation='nearest'),
|
||||
# Conv2D(dims[i+1], kernel_size=2, strides=1, padding='same'),
|
||||
# Dense(dims[i+1]),
|
||||
])
|
||||
self.upsample_layers.append(downsample_layer)
|
||||
|
||||
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
|
||||
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
|
||||
cur = 0
|
||||
for i in range(len(depths)):
|
||||
stage = tf.keras.Sequential(
|
||||
[Block(dim=dims[i], drop_path=dp_rates[cur + j], kernel_size=kernel_size,
|
||||
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
|
||||
)
|
||||
self.stages.append(stage)
|
||||
cur += depths[i]
|
||||
|
||||
self.last = Dense(3, activation='tanh')
|
||||
|
||||
def call(self, x, training):
|
||||
for i in reversed(range(len(self.upsample_layers))):
|
||||
x = self.upsample_layers[i](x, training=training)
|
||||
x = self.stages[i](x, training=training)
|
||||
x = self.last(x, training=training)
|
||||
return x
|
||||
|
||||
|
||||
class ConvNeXtReverseLabels(tf.keras.Model):
|
||||
r""" ConvNeXt
|
||||
A Tensorflow keras impl of : `A ConvNet for the 2020s` -
|
||||
https://arxiv.org/pdf/2201.03545.pdf
|
||||
|
||||
Args:
|
||||
num_classes (int): Number of classes for classification head. Default: 1000
|
||||
depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
|
||||
dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
|
||||
include_top (bool): whether to add head or just use it as feature extractor. Default: True
|
||||
drop_path_rate (float): Stochastic depth rate. Default: 0.
|
||||
layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
|
||||
head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
|
||||
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
|
||||
nr_of_classes=7):
|
||||
super().__init__()
|
||||
self.downsample_layers = [] # stem and 3 intermediate downsampling conv layers
|
||||
stem = tf.keras.Sequential([
|
||||
Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
|
||||
LayerNormalization(epsilon=1e-6)]
|
||||
)
|
||||
self.downsample_layers.append(stem)
|
||||
for i in range(len(depths) - 1):
|
||||
downsample_layer = tf.keras.Sequential([
|
||||
LayerNormalization(epsilon=1e-6),
|
||||
Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')]
|
||||
)
|
||||
self.downsample_layers.append(downsample_layer)
|
||||
|
||||
self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks
|
||||
dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
|
||||
cur = 0
|
||||
for i in range(len(depths)):
|
||||
stage = tf.keras.Sequential(
|
||||
[Block(dim=dims[i], drop_path=dp_rates[cur + j],
|
||||
layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
|
||||
)
|
||||
self.stages.append(stage)
|
||||
cur += depths[i]
|
||||
|
||||
# self.last = Dense(nr_of_classes, activation='softmax')
|
||||
self.last = Dense(nr_of_classes - 1)
|
||||
|
||||
def call(self, x, training):
|
||||
for i in reversed(range(len(self.downsample_layers))):
|
||||
x = self.downsample_layers[i](x, training=training)
|
||||
x = self.stages[i](x, training=training)
|
||||
x = self.last(x, training=training)
|
||||
return x
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
model_urls = {
|
||||
"convnext_tiny_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_tiny_1k_224_ema.h5",
|
||||
"convnext_small_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_small_1k_224_ema.h5",
|
||||
"convnext_base_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_base_22k_1k_224.h5",
|
||||
"convnext_base_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_base_22k_1k_384.h5",
|
||||
"convnext_large_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_large_22k_1k_224.h5",
|
||||
"convnext_large_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_large_22k_1k_384.h5",
|
||||
"convnext_xlarge_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_xlarge_22k_1k_224_ema.h5",
|
||||
"convnext_xlarge_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_xlarge_22k_1k_384_ema.h5",
|
||||
}
|
||||
|
||||
|
||||
model_configs = dict(
|
||||
convnext_tiny=dict(
|
||||
depths=[3, 3, 9, 3],
|
||||
dims=[8, 16, 32, 512]
|
||||
),
|
||||
convnext_small=dict(
|
||||
depths=[3, 3, 27, 3],
|
||||
dims=[96, 192, 384, 768]
|
||||
),
|
||||
convnext_base=dict(
|
||||
depths=[3, 3, 27, 3],
|
||||
dims=[128, 256, 512, 1024]
|
||||
),
|
||||
convnext_large=dict(
|
||||
depths=[3, 3, 27, 3],
|
||||
dims=[192, 384, 768, 1536]
|
||||
),
|
||||
convnext_xlarge=dict(
|
||||
depths=[3, 3, 27, 3],
|
||||
dims=[256, 512, 1024, 2048]
|
||||
),
|
||||
)
|
||||
|
||||
@@ -0,0 +1,140 @@
|
||||
import os
|
||||
import time
|
||||
import random
|
||||
import tensorflow as tf
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
# import plotly.graph_objects as go
|
||||
# import plotly.express as ex
|
||||
|
||||
from custom_utils.custom_utils import load_model
|
||||
from models.base import AEEncoder, VariationalEncoder, Decoder
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
class CVAE(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(CVAE, self).__init__()
|
||||
self.latent_dim = config['latent_dim']
|
||||
self.latent_dim = config['cvae']['latent_dim']
|
||||
self.size = config['resolution']
|
||||
|
||||
self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
if config['static_model'] != '':
|
||||
load_model(self.static_enc, config['static_model'] + '_enc')
|
||||
|
||||
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
|
||||
if config['body_model'] != '':
|
||||
load_model(self.body_enc, config['body_model'] + '_enc')
|
||||
|
||||
if config['freeze_encoders']:
|
||||
self.static_enc.trainable = False
|
||||
self.body_enc.trainable = False
|
||||
|
||||
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
|
||||
tf.keras.layers.Dense(256, activation='relu'),
|
||||
tf.keras.layers.Dense(128, activation='relu'),
|
||||
tf.keras.layers.Dense(self.latent_dim, activation='tanh'),])
|
||||
|
||||
self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
|
||||
self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'], is_label=config['cvae']['is_label'])
|
||||
|
||||
self.add_noise = config['cvae']['add_noise']
|
||||
self.noise_amount = config['cvae']['noise_amount']
|
||||
|
||||
# @tf.function
|
||||
def sample_single(self, uv_template, uv_body_posed, uv_mask, eps=None, nr_of_samples=4):
|
||||
if eps is None:
|
||||
eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
|
||||
|
||||
uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
|
||||
uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
|
||||
uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
|
||||
|
||||
# uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
|
||||
z_static = self.static_enc(uv_template)
|
||||
z_body = self.body_enc(uv_body_posed)
|
||||
|
||||
z_cond = tf.concat([z_static, z_body], axis=-1)
|
||||
z_cond = self.cond_encoder(z_cond)
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z)
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def encode_conds(self, uv_template, uv_body_posed, training):
|
||||
z_static = self.static_enc(uv_template, training=training)
|
||||
z_body = self.body_enc(uv_body_posed, training=training)
|
||||
|
||||
z_cond = tf.concat([z_static, z_body], axis=-1)
|
||||
z_cond = self.cond_encoder(z_cond, training=training)
|
||||
|
||||
return z_cond
|
||||
|
||||
def sample(self, uv_template, uv_body_posed, uv_mask, noise=None):
|
||||
if noise is None:
|
||||
eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
|
||||
else:
|
||||
eps = noise
|
||||
|
||||
z_cond = self.encode_conds(uv_template, uv_body_posed, training=False)
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z, training=False)
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
background = tf.ones((uv_template.shape[0],512,256,1)) * (1 - uv_mask[:,:,:,tf.newaxis])
|
||||
x = tf.concat((background, x), axis=-1)
|
||||
|
||||
x = tf.math.softmax(x)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
# encoders have to be the same for curv_vae and general CVAE
|
||||
def sample_without_encoding(self, z_static, z_body, uv_mask, eps=None):
|
||||
eps = tf.random.normal(shape=(z_static.shape[0], self.latent_dim))
|
||||
|
||||
z_cond = tf.concat([z_static, z_body], axis=-1)
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z)
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def reparameterize(self, mean, logvar):
|
||||
eps = tf.random.normal(shape=mean.shape)
|
||||
return eps * tf.exp(logvar * .5) + mean
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
uv_template, uv_body_posed, uv_curv, uv_mask = inputs
|
||||
|
||||
|
||||
x = tf.concat([uv_template, uv_body_posed, uv_curv], axis=-1)
|
||||
mean, logvar = self.encoder(x)
|
||||
z = self.reparameterize(mean, logvar)
|
||||
|
||||
z_cond = self.encode_conds(uv_template, uv_body_posed, training=training)
|
||||
|
||||
z = tf.concat([z, z_cond], axis=-1)
|
||||
x = self.decoder(z, training=training)
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
background = tf.ones((4,512,256,1)) * (1 - uv_mask[:,:,:,tf.newaxis])
|
||||
x = tf.concat((background, x), axis=-1)
|
||||
|
||||
x = tf.math.softmax(x)
|
||||
|
||||
return x, mean, logvar
|
||||
@@ -0,0 +1,285 @@
|
||||
import os
|
||||
import time
|
||||
import random
|
||||
import tensorflow as tf
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
# import plotly.graph_objects as go
|
||||
# import plotly.express as ex
|
||||
|
||||
from custom_utils.custom_utils import load_model
|
||||
from models.base import AEEncoder, VariationalEncoder, Decoder
|
||||
from models.base import AEEncoderNew, VariationalEncoderNew, DecoderNew
|
||||
|
||||
|
||||
class CVAE(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(CVAE, self).__init__()
|
||||
self.latent_dim = config['cvae']['latent_dim']
|
||||
self.size = config['resolution']
|
||||
|
||||
if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
|
||||
self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
if config['static_model'] != '':
|
||||
load_model(self.static_enc, config['static_model'] + '_enc')
|
||||
|
||||
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
|
||||
if config['body_model'] != '':
|
||||
load_model(self.body_enc, config['body_model'] + '_enc')
|
||||
|
||||
if config['use_curvature']:
|
||||
self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
if config['curvature_model'] != '':
|
||||
load_model(self.curv_enc, config['curvature_model'] + '_enc')
|
||||
|
||||
if config['freeze_encoders']:
|
||||
if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
|
||||
self.static_enc.trainable = False
|
||||
self.body_enc.trainable = False
|
||||
if config['use_curvature']:
|
||||
self.curv_enc.trainable = False
|
||||
|
||||
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
|
||||
tf.keras.layers.Dense(256, activation='relu'),
|
||||
tf.keras.layers.Dense(128, activation='relu'),
|
||||
# tf.keras.layers.Lambda(lambda x: x / 1000),
|
||||
tf.keras.layers.Dense(config['latent_dim'], activation='tanh')])
|
||||
|
||||
self.encoder = VariationalEncoder(config['cvae']['latent_dim'], config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
|
||||
self.decoder = Decoder(config['cvae']['decoder_input'], config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'])
|
||||
|
||||
if config['avg_garment_path'] != '':
|
||||
self.avg_garment = np.load(config['avg_garment_path'])
|
||||
else:
|
||||
self.avg_garment = None
|
||||
|
||||
self.add_noise = config['cvae']['add_noise']
|
||||
self.noise_amount = config['cvae']['noise_amount']
|
||||
self.use_curvature = config['use_curvature']
|
||||
self.only_curvatures = config['only_curvatures'] if 'only_curvatures' in config.keys() and config['only_curvatures'] == True else None
|
||||
self.use_offsets_over_body = not config['use_skinning']
|
||||
|
||||
if self.use_offsets_over_body:
|
||||
assert self.avg_garment is None, 'When we are not using skinning, we can\'t use avg_garment!'
|
||||
if self.avg_garment is not None:
|
||||
assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
|
||||
|
||||
# @tf.function
|
||||
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
|
||||
if noise is None:
|
||||
eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
|
||||
else:
|
||||
eps = noise
|
||||
|
||||
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
|
||||
if add_noise:
|
||||
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
|
||||
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z, training=False)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
if self.use_offsets_over_body:
|
||||
x = x + uv_body_posed
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def reparameterize(self, mean, logvar, training, fix_noise=False):
|
||||
if training == True:
|
||||
if not fix_noise:
|
||||
eps = tf.random.normal(shape=mean.shape)
|
||||
else:
|
||||
eps = tf.random.normal(shape=(1, mean.shape[1]))
|
||||
eps = tf.tile(eps, [mean.shape[0],1])
|
||||
out = eps * tf.exp(logvar * .5) + mean
|
||||
else:
|
||||
out = tf.exp(logvar * .5) + mean
|
||||
return out
|
||||
|
||||
|
||||
def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
|
||||
if self.only_curvatures:
|
||||
z_cond = self.curv_enc(uv_curv, training=training)
|
||||
else:
|
||||
z_body = self.body_enc(uv_body_posed, training=training)
|
||||
z_static = self.static_enc(uv_template, training=training)
|
||||
if self.use_curvature:
|
||||
z_curv = self.curv_enc(uv_curv, training=training)
|
||||
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
|
||||
else:
|
||||
z_cond = tf.concat([z_static, z_body], axis=-1)
|
||||
|
||||
z_cond = self.cond_encoder(z_cond, training=training)
|
||||
|
||||
return z_cond
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
|
||||
if self.only_curvatures:
|
||||
x = tf.concat([uv_unposed, uv_curv], axis=-1)
|
||||
else:
|
||||
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
|
||||
|
||||
mean, logvar = self.encoder(x, training=training)
|
||||
z = self.reparameterize(mean, logvar, training)
|
||||
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=training)
|
||||
|
||||
if self.add_noise and training == True:
|
||||
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
|
||||
|
||||
|
||||
z = tf.concat([z, z_cond], axis=-1)
|
||||
x = self.decoder(z, training=training)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
if self.use_offsets_over_body:
|
||||
x = x + uv_body_posed
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x, mean, logvar
|
||||
|
||||
|
||||
|
||||
|
||||
class CVAENew(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(CVAENew, self).__init__()
|
||||
self.size = config['resolution']
|
||||
|
||||
if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
|
||||
self.static_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_static']['encoder_depths'], config['cond_ae_static']['encoder_dims'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'])
|
||||
if config['static_model'] != '':
|
||||
load_model(self.static_enc, config['static_model'] + '_enc')
|
||||
|
||||
self.body_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['kernel_size'])
|
||||
if config['body_model'] != '':
|
||||
load_model(self.body_enc, config['body_model'] + '_enc')
|
||||
|
||||
if config['use_curvature']:
|
||||
self.curv_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_curv']['encoder_depths'], config['cond_ae_curv']['encoder_dims'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation'])
|
||||
if config['curvature_model'] != '':
|
||||
load_model(self.curv_enc, config['curvature_model'] + '_enc')
|
||||
|
||||
if config['freeze_encoders']:
|
||||
if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
|
||||
self.static_enc.trainable = False
|
||||
self.body_enc.trainable = False
|
||||
if config['use_curvature']:
|
||||
self.curv_enc.trainable = False
|
||||
|
||||
self.encoder = VariationalEncoderNew(config['cvae']['encoder_depths'], config['cvae']['encoder_dims'], kernel_size=config['cvae']['kernel_size'])
|
||||
self.decoder = DecoderNew(config['cvae']['decoder_depths'], config['cvae']['decoder_dims'], config['cvae']['kernel_size'])
|
||||
|
||||
self.reshape = tf.keras.layers.Reshape((8, 4, 64))
|
||||
|
||||
if config['avg_garment_path'] != '':
|
||||
self.avg_garment = np.load(config['avg_garment_path'])
|
||||
else:
|
||||
self.avg_garment = None
|
||||
|
||||
self.add_noise = config['cvae']['add_noise']
|
||||
self.noise_amount = config['cvae']['noise_amount']
|
||||
self.use_curvature = config['use_curvature']
|
||||
self.only_curvatures = config['only_curvatures'] if 'only_curvatures' in config.keys() and config['only_curvatures'] == True else None
|
||||
self.use_offsets_over_body = not config['use_skinning']
|
||||
|
||||
if self.use_offsets_over_body:
|
||||
assert self.avg_garment is None, 'When we are not using skinning, we can\'t use avg_garment!'
|
||||
if self.avg_garment is not None:
|
||||
assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
|
||||
|
||||
# @tf.function
|
||||
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
|
||||
if noise is None:
|
||||
eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
|
||||
else:
|
||||
eps = noise
|
||||
|
||||
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
|
||||
if add_noise:
|
||||
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
|
||||
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z, training=False)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
if self.use_offsets_over_body:
|
||||
x = x + uv_body_posed
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def reparameterize(self, mean, logvar, training, fix_noise=False):
|
||||
if training == True:
|
||||
if not fix_noise:
|
||||
eps = tf.random.normal(shape=mean.shape)
|
||||
else:
|
||||
eps = tf.random.normal(shape=(1, mean.shape[1]))
|
||||
eps = tf.tile(eps, [mean.shape[0],1])
|
||||
out = eps * tf.exp(logvar * .5) + mean
|
||||
else:
|
||||
out = tf.exp(logvar * .5) + mean
|
||||
return out
|
||||
|
||||
|
||||
def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
|
||||
if self.only_curvatures:
|
||||
z_cond = self.curv_enc(uv_curv, training=training)
|
||||
else:
|
||||
z_body = self.body_enc(uv_body_posed, training=training)
|
||||
z_static = self.static_enc(uv_template, training=training)
|
||||
if self.use_curvature:
|
||||
z_curv = self.curv_enc(uv_curv, training=training)
|
||||
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
|
||||
else:
|
||||
z_cond = tf.concat([z_static, z_body], axis=-1)
|
||||
|
||||
return z_cond
|
||||
|
||||
|
||||
def call(self, inputs, training):
|
||||
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
|
||||
if self.only_curvatures:
|
||||
x = tf.concat([uv_unposed, uv_curv], axis=-1)
|
||||
else:
|
||||
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
|
||||
|
||||
mean, logvar = self.encoder(x, training=training)
|
||||
z = self.reparameterize(mean, logvar, training)
|
||||
z = self.reshape(z)
|
||||
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=training)
|
||||
|
||||
if self.add_noise and training == True:
|
||||
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
|
||||
|
||||
|
||||
z = tf.concat([z, z_cond], axis=-1)
|
||||
x = self.decoder(z, training=training)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
if self.use_offsets_over_body:
|
||||
x = x + uv_body_posed
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x, mean, logvar
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
import os
|
||||
import time
|
||||
import random
|
||||
import tensorflow as tf
|
||||
import numpy as np
|
||||
import scipy as sp
|
||||
# import plotly.graph_objects as go
|
||||
# import plotly.express as ex
|
||||
|
||||
from custom_utils.custom_utils import load_model
|
||||
from models.base import AEEncoder, VariationalEncoder, Decoder, VariationalEncoderNew, DecoderNew
|
||||
from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
|
||||
|
||||
|
||||
class PyramidLevel(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(PyramidLevel, self).__init__()
|
||||
self.latent_dim = config['latent_dim']
|
||||
self.size = config['resolution']
|
||||
|
||||
self.encoder = VariationalEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'])
|
||||
self.decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], config['kernel_size'])
|
||||
|
||||
self.upscale_enhancer = None
|
||||
if config['upscale_enhancer'] != '':
|
||||
k_size = config['upscale_enhancer_kernel_size'] if 'upscale_enhancer_kernel_size' in config.keys() else ((3 if config['resolution'] <= 64 else 5))
|
||||
self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=k_size)
|
||||
load_model(self.upscale_enhancer, config['upscale_enhancer'])
|
||||
self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
|
||||
|
||||
if config['avg_garment_path'] != '':
|
||||
self.avg_garment = np.load(config['avg_garment_path'])
|
||||
else:
|
||||
self.avg_garment = None
|
||||
|
||||
|
||||
def enhance(self, low, mask, training=None):
|
||||
x = self.upscale_enhancer([low, mask], training=training)
|
||||
return x
|
||||
|
||||
|
||||
def sample(self, z_cond, uv_mask, noise=None):
|
||||
if noise is None:
|
||||
eps = tf.random.normal(shape=(uv_mask.shape[0], self.latent_dim))
|
||||
else:
|
||||
eps = noise
|
||||
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def reparameterize(self, mean, logvar, training, fix_noise=False):
|
||||
if training == True:
|
||||
if not fix_noise:
|
||||
eps = tf.random.normal(shape=mean.shape)
|
||||
else:
|
||||
eps = tf.random.normal(shape=(1, mean.shape[1]))
|
||||
eps = tf.tile(eps, [mean.shape[0],1])
|
||||
out = eps * tf.exp(logvar * .5) + mean
|
||||
else:
|
||||
out = tf.exp(logvar * .5) + mean
|
||||
return out
|
||||
|
||||
|
||||
def call(self, inputs, z_cond, training=None):
|
||||
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs
|
||||
|
||||
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
|
||||
mean, logvar = self.encoder(x, training=training)
|
||||
z = self.reparameterize(mean, logvar, training)
|
||||
|
||||
z = tf.concat([z, z_cond], axis=-1)
|
||||
x = self.decoder(z, training=training)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
# if self.upscale_enhancer is not None:
|
||||
# pred_low = self.enhance(x, uv_mask, training=training)
|
||||
# x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
|
||||
# x = x * uv_mask_high[:,:,:,tf.newaxis]
|
||||
|
||||
return x, mean, logvar
|
||||
|
||||
|
||||
|
||||
class PyramidLevelNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(PyramidLevelNew, self).__init__()
|
||||
self.size = config['resolution']
|
||||
|
||||
self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'])
|
||||
self.decoder = DecoderNew(config['decoder_depths'], config['decoder_dims'], config['kernel_size'])
|
||||
|
||||
self.upscale_enhancer = None
|
||||
if config['upscale_enhancer'] != '':
|
||||
k_size = config['upscale_enhancer_kernel_size'] if 'upscale_enhancer_kernel_size' in config.keys() else ((3 if config['resolution'] <= 64 else 5))
|
||||
self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=k_size)
|
||||
load_model(self.upscale_enhancer, config['upscale_enhancer'])
|
||||
self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
|
||||
|
||||
if config['avg_garment_path'] != '':
|
||||
self.avg_garment = np.load(config['avg_garment_path'])
|
||||
else:
|
||||
self.avg_garment = None
|
||||
|
||||
self.reshape = tf.keras.layers.Reshape((8, 4, 64))
|
||||
|
||||
|
||||
def enhance(self, low, mask, training=None):
|
||||
x = self.upscale_enhancer([low, mask], training=training)
|
||||
return x
|
||||
|
||||
|
||||
def sample(self, z_cond, uv_mask, noise=None):
|
||||
if noise is None:
|
||||
eps = tf.random.normal(shape=(uv_mask.shape[0], self.latent_dim))
|
||||
else:
|
||||
eps = noise
|
||||
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def reparameterize(self, mean, logvar, training, fix_noise=False):
|
||||
if training == True:
|
||||
if not fix_noise:
|
||||
eps = tf.random.normal(shape=mean.shape)
|
||||
else:
|
||||
eps = tf.random.normal(shape=(1, mean.shape[1]))
|
||||
eps = tf.tile(eps, [mean.shape[0],1])
|
||||
out = eps * tf.exp(logvar * .5) + mean
|
||||
else:
|
||||
out = tf.exp(logvar * .5) + mean
|
||||
return out
|
||||
|
||||
|
||||
def call(self, inputs, z_cond, training=None):
|
||||
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs
|
||||
|
||||
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
|
||||
mean, logvar = self.encoder(x, training=training)
|
||||
z = self.reparameterize(mean, logvar, training)
|
||||
z = self.reshape(z)
|
||||
|
||||
z = tf.concat([z, z_cond], axis=-1)
|
||||
x = self.decoder(z, training=training)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
# if self.upscale_enhancer is not None:
|
||||
# pred_low = self.enhance(x, uv_mask, training=training)
|
||||
# x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
|
||||
# x = x * uv_mask_high[:,:,:,tf.newaxis]
|
||||
|
||||
return x, mean, logvar
|
||||
@@ -0,0 +1,365 @@
|
||||
import os
|
||||
import time
|
||||
import random
|
||||
import tensorflow as tf
|
||||
import numpy as np
|
||||
# import plotly.graph_objects as go
|
||||
# import plotly.express as ex
|
||||
|
||||
from custom_utils.custom_utils import load_model
|
||||
from models.base import AEEncoder, AEEncoderNew
|
||||
from models.pyramid_level import PyramidLevel, PyramidLevelNew
|
||||
from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
|
||||
|
||||
|
||||
class Pyramid(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(Pyramid, self).__init__()
|
||||
|
||||
self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
if config['pyramid']['static_model'] != '':
|
||||
load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
|
||||
|
||||
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
|
||||
if config['pyramid']['body_model'] != '':
|
||||
load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
|
||||
|
||||
self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
if config['pyramid']['curvature_model'] != '':
|
||||
load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
|
||||
|
||||
if config['pyramid']['freeze_encoders']:
|
||||
self.static_enc.trainable = False
|
||||
self.body_enc.trainable = False
|
||||
self.curv_enc.trainable = False
|
||||
|
||||
self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
|
||||
tf.keras.layers.Dense(256, activation='relu'),
|
||||
tf.keras.layers.Dense(128, activation='relu'),
|
||||
tf.keras.layers.Dense(config['pyramid']['latent_dim'], activation='tanh'),])
|
||||
|
||||
self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
|
||||
|
||||
# if config['is_cape']:
|
||||
# if len(config['pyramid']['resolutions']) == 2:
|
||||
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# if len(config['pyramid']['resolutions']) == 3:
|
||||
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# if len(config['pyramid']['resolutions']) == 4:
|
||||
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# else:
|
||||
if len(config['pyramid']['resolutions']) == 2:
|
||||
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
if len(config['pyramid']['resolutions']) == 3:
|
||||
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
|
||||
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
if len(config['pyramid']['resolutions']) == 4:
|
||||
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
|
||||
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
|
||||
self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
if len(config['pyramid']['resolutions']) == 5:
|
||||
self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)])
|
||||
self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
|
||||
self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
|
||||
self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
|
||||
self.add_noise = config['pyramid']['add_noise']
|
||||
self.noise_amount = config['pyramid']['noise_amount']
|
||||
self.use_offsets = config['use_offsets']
|
||||
|
||||
self.upscale_enhancer = None
|
||||
if config['pyramid']['final_upscale_enhancer'] != '':
|
||||
k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5))
|
||||
self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size)
|
||||
load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
|
||||
self.upscale_enhancer.trainable = False
|
||||
|
||||
def enhance(self, low, mask, training):
|
||||
x = self.upscale_enhancer([low, mask], training=training)
|
||||
return x
|
||||
|
||||
def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
|
||||
z_static = self.static_enc(uv_template, training)
|
||||
z_body = self.body_enc(uv_body_posed, training)
|
||||
z_curv = self.curv_enc(uv_curv, training)
|
||||
|
||||
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
|
||||
z_cond = self.cond_encoder(z_cond, training=training)
|
||||
|
||||
return z_cond
|
||||
|
||||
def call(self, inputs, training):
|
||||
levels_input, conds = inputs
|
||||
uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
|
||||
|
||||
z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training)
|
||||
|
||||
if self.add_noise == True and training == True:
|
||||
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
|
||||
|
||||
out_levels = []
|
||||
pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
|
||||
if self.use_offsets:
|
||||
pred = pred + levels_input[0][2] # uv_body_posed_32
|
||||
if 1 != len(self.levels):
|
||||
pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
|
||||
pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
|
||||
out_levels.append([pred, mean, logvar])
|
||||
pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
|
||||
else:
|
||||
pred = self.enhance(pred, levels_input[0][-2], training=training)
|
||||
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
|
||||
out_levels.append([pred, mean, logvar])
|
||||
|
||||
for i in range(1, len(self.levels)):
|
||||
pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
|
||||
pred = pred + pred_level
|
||||
if i != len(self.levels) - 1:
|
||||
pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
|
||||
out_levels.append([pred, mean, logvar])
|
||||
pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
|
||||
pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
|
||||
else:
|
||||
if self.upscale_enhancer is not None:
|
||||
pred = self.enhance(pred, levels_input[-1][-2], training=training)
|
||||
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
|
||||
out_levels.append([pred, mean, logvar])
|
||||
|
||||
return out_levels
|
||||
|
||||
# # 32
|
||||
# pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
|
||||
# pred_0 = self.enhance(pred_0, levels_input[0][-2], training=training)
|
||||
# pred_0 = tf.image.resize(pred_0, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
# pred_0 = pred_0 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
|
||||
# return [ [pred_0, mean_0, logvar_0] ]
|
||||
|
||||
# # 32 64
|
||||
# pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
|
||||
# pred_0 = self.levels[0].enhance(pred_0, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
|
||||
# pred_0 = tf.image.resize(pred_0, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
|
||||
# pred = pred_0 * levels_input[1][-2][:,:,:,tf.newaxis]
|
||||
|
||||
# pred_1, mean_1, logvar_1 = self.levels[1](levels_input[1], z_cond, training=training)
|
||||
# pred_1 = pred + pred_1
|
||||
# pred_1 = self.enhance(pred_1, levels_input[-1][-2], training=training)
|
||||
# pred_1 = tf.image.resize(pred_1, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
# pred_1 = pred_1 * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
|
||||
|
||||
# return [ [pred_0, mean_0, logvar_0], [pred_1, mean_1, logvar_1] ]
|
||||
|
||||
|
||||
# class Pyramid(tf.keras.Model):
|
||||
|
||||
# def __init__(self, config):
|
||||
# super(Pyramid, self).__init__()
|
||||
|
||||
# self.static_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
# if config['pyramid']['static_model'] != '':
|
||||
# load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
|
||||
|
||||
# self.body_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
# if config['pyramid']['body_model'] != '':
|
||||
# load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
|
||||
|
||||
# self.curv_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
# if config['pyramid']['curvature_model'] != '':
|
||||
# load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
|
||||
|
||||
# if config['pyramid']['freeze_encoders']:
|
||||
# self.static_enc.trainable = False
|
||||
# self.body_enc.trainable = False
|
||||
# self.curv_enc.trainable = False
|
||||
|
||||
# self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'),
|
||||
# tf.keras.layers.Dense(256, activation='relu'),
|
||||
# tf.keras.layers.Dense(128, activation='relu'),
|
||||
# tf.keras.layers.Dense(config['pyramid']['latent_dim']),])
|
||||
|
||||
# self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae': config['cond_ae']} }) for i in range(len(config['pyramid']['levels']))]
|
||||
|
||||
|
||||
# self.upscale_enhancer = None
|
||||
# if config['pyramid']['final_upscale_enhancer'] != '':
|
||||
# self.upscale_enhancer = UpscalerBasicNew(low_res=config['pyramid']['resolutions'][-1])
|
||||
# load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
|
||||
# self.upscale_enhancer.trainable = False
|
||||
|
||||
# def enhance(self, low, mask, training):
|
||||
# x = self.upscale_enhancer([low, mask], training=training)
|
||||
# return x
|
||||
|
||||
# def encode_conds(self, uv_template, uv_body_posed, uv_curv):
|
||||
# z_static = self.static_enc(uv_template)
|
||||
# z_body = self.body_enc(uv_body_posed)
|
||||
# z_curv = self.curv_enc(uv_curv)
|
||||
|
||||
# z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
|
||||
# z_cond = self.cond_encoder(z_cond)
|
||||
|
||||
# return z_cond
|
||||
|
||||
# def call(self, inputs, training):
|
||||
# levels_input, conds = inputs
|
||||
# uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
|
||||
# z_cond = self.encode_conds(uv_static, uv_body, uv_curv)
|
||||
|
||||
# out_levels = []
|
||||
# pred_0, mean_0, logvar_0 = self.levels[0](levels_input[0], z_cond, training=training)
|
||||
# pred_0 = self.levels[0].enhance(pred_0, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
|
||||
# # pred_0 = self.enhance(pred_0, levels_input[0][-2], training=training)
|
||||
|
||||
# ##################
|
||||
# # if more than 1 level
|
||||
# # x = tf.image.resize(pred_0, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.LANCZOS5)
|
||||
# # x = x * levels_input[1][-2][:,:,:,tf.newaxis]
|
||||
|
||||
# # if only 1 level
|
||||
# pred_0 = tf.image.resize(pred_0, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
# ##################
|
||||
|
||||
|
||||
# # pred_1, mean_1, logvar_1 = self.levels[1](levels_input[1], z_cond, training=training)
|
||||
# # pred_1 = pred_0 + pred_1
|
||||
# # pred_1 = self.levels[1].enhance(pred, levels_input[1][-2], training=training)
|
||||
|
||||
# # for i in range(1, len(self.levels)):
|
||||
# # pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond, training=training)
|
||||
# # pred = pred + pred_level
|
||||
# # if i != len(self.levels) - 1:
|
||||
# # pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
|
||||
# # out_levels.append([pred, mean, logvar])
|
||||
# # pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.LANCZOS5)
|
||||
# # pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
|
||||
# # else:
|
||||
# # if self.upscale_enhancer is not None:
|
||||
# # pred = self.enhance(pred, levels_input[-1][-2], training=training)
|
||||
# # pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.LANCZOS5)
|
||||
# # out_levels.append([pred, mean, logvar])
|
||||
# # pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
|
||||
|
||||
# out_levels = [
|
||||
# [pred_0, mean_0, logvar_0]
|
||||
# ]
|
||||
|
||||
# return out_levels
|
||||
|
||||
|
||||
class PyramidNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, config):
|
||||
super(PyramidNew, self).__init__()
|
||||
|
||||
self.static_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_static']['encoder_depths'], config['cond_ae_static']['encoder_dims'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'])
|
||||
if config['pyramid']['static_model'] != '':
|
||||
load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
|
||||
|
||||
self.body_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['kernel_size'])
|
||||
if config['pyramid']['body_model'] != '':
|
||||
load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
|
||||
|
||||
self.curv_enc : AEEncoderNew = AEEncoderNew(config['cond_ae_curv']['encoder_depths'], config['cond_ae_curv']['encoder_dims'], config['cond_ae_curv']['kernel_size'], config['cond_ae_curv']['activation'])
|
||||
if config['pyramid']['curvature_model'] != '':
|
||||
load_model(self.curv_enc, config['pyramid']['curvature_model'] + '_enc')
|
||||
|
||||
if config['pyramid']['freeze_encoders']:
|
||||
self.static_enc.trainable = False
|
||||
self.body_enc.trainable = False
|
||||
self.curv_enc.trainable = False
|
||||
|
||||
self.levels = [PyramidLevelNew({ **config['pyramid']['levels'][i] }) for i in range(len(config['pyramid']['levels']))]
|
||||
|
||||
|
||||
# if len(config['pyramid']['resolutions']) == 2:
|
||||
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# if len(config['pyramid']['resolutions']) == 3:
|
||||
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
|
||||
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# if len(config['pyramid']['resolutions']) == 4:
|
||||
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
|
||||
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
|
||||
# self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
# if len(config['pyramid']['resolutions']) == 5:
|
||||
# self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)])
|
||||
# self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)])
|
||||
# self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)])
|
||||
# self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)])
|
||||
|
||||
self.add_noise = config['pyramid']['add_noise']
|
||||
self.noise_amount = config['pyramid']['noise_amount']
|
||||
self.use_offsets = config['use_offsets']
|
||||
|
||||
self.upscale_enhancer = None
|
||||
if config['pyramid']['final_upscale_enhancer'] != '':
|
||||
k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5))
|
||||
self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size)
|
||||
load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
|
||||
self.upscale_enhancer.trainable = False
|
||||
|
||||
def enhance(self, low, mask, training):
|
||||
x = self.upscale_enhancer([low, mask], training=training)
|
||||
return x
|
||||
|
||||
def encode_conds(self, uv_template, uv_body_posed, uv_curv, training):
|
||||
z_static = self.static_enc(uv_template, training)
|
||||
z_body = self.body_enc(uv_body_posed, training)
|
||||
z_curv = self.curv_enc(uv_curv, training)
|
||||
|
||||
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
|
||||
|
||||
return z_cond
|
||||
|
||||
def call(self, inputs, training):
|
||||
levels_input, conds = inputs
|
||||
uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
|
||||
|
||||
z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training)
|
||||
|
||||
if self.add_noise == True and training == True:
|
||||
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
|
||||
|
||||
out_levels = []
|
||||
pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training)
|
||||
if self.use_offsets:
|
||||
pred = pred + levels_input[0][2] # uv_body_posed_32
|
||||
if 1 != len(self.levels):
|
||||
pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
|
||||
pred = tf.image.resize(pred, (self.levels[1].size, self.levels[1].size//2), tf.image.ResizeMethod.BILINEAR)
|
||||
out_levels.append([pred, mean, logvar])
|
||||
pred = pred * levels_input[1][-2][:,:,:,tf.newaxis]
|
||||
else:
|
||||
# pred = self.enhance(pred, levels_input[0][-2], training=training)
|
||||
pred = self.levels[0].enhance(pred, levels_input[0][-2], training=training) # levels_input[i][-2]=uv_mask_res
|
||||
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.BILINEAR)
|
||||
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
|
||||
out_levels.append([pred, mean, logvar])
|
||||
|
||||
for i in range(1, len(self.levels)):
|
||||
input = levels_input[i]
|
||||
input[0] = input[0] - pred
|
||||
input_min, input_max = tf.reduce_min(input[0], axis=[1,2,3]), tf.reduce_max(input[0], axis=[1,2,3])
|
||||
input[0] = (input[0] - input_min[:,None,None,None]) / (input_max - input_min)[:,None,None,None]
|
||||
input[0] = input[0] * input[-2][:,:,:,None] # uv_mask_res
|
||||
pred_level, mean, logvar = self.levels[i](input, z_cond, training=training)
|
||||
pred_level = ( (pred_level * (input_max - input_min)[:,None,None,None]) + input_min[:,None,None,None] ) * input[-2][:,:,:,None] # uv_mask_res
|
||||
pred = pred + pred_level
|
||||
if i != len(self.levels) - 1:
|
||||
pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
|
||||
out_levels.append([pred, mean, logvar])
|
||||
pred = tf.image.resize(pred, (self.levels[i+1].size, self.levels[i+1].size//2), tf.image.ResizeMethod.BILINEAR)
|
||||
pred = pred * levels_input[i+1][-2][:,:,:,tf.newaxis]
|
||||
else:
|
||||
# if self.upscale_enhancer is not None:
|
||||
# pred = self.enhance(pred, levels_input[-1][-2], training=training)
|
||||
pred = self.levels[-1].enhance(pred, levels_input[-1][-2], training=training)
|
||||
pred = tf.image.resize(pred, (512, 256), tf.image.ResizeMethod.BILINEAR)
|
||||
pred = pred * levels_input[-1][-1][:,:,:,tf.newaxis] # full res uv mask
|
||||
out_levels.append([pred, mean, logvar])
|
||||
|
||||
return out_levels
|
||||
@@ -0,0 +1,172 @@
|
||||
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(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'),
|
||||
])
|
||||
|
||||
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user