luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

separate sub models from their training (wip)

Browse Source
This commit is contained in:
luni
2022-09-14 10:03:14 +00:00
parent 4caaa1af27
commit ac851d2de7
6 changed files with 449 additions and 371 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_body.py
+29 -25
View File
@@ -1,5 +1,5 @@
import os import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0" os.environ["CUDA_VISIBLE_DEVICES"] = "3"
import time import time
import argparse import argparse
@@ -17,10 +17,10 @@ from losses import edge_loss, uv_mask_loss
from dataset import DatasetFactory, StaticDataset from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL from custom_utils.smpl import SMPL
from models.ae_body_model import AutoEncoder from models.ae_body_model import AutoEncoder, AutoEncoder128, AutoEncoder32
name='ae_body_8' name='ae_body_8_d128'
def init_wandb(): def init_wandb():
run = wandb.init(project='vae', name=name, run = wandb.init(project='vae', name=name,
notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \ notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \
@@ -31,7 +31,8 @@ def init_wandb():
flattened bottleneck with 128 latent_dim\ remove global orientation\ flattened bottleneck with 128 latent_dim\ remove global orientation\
fixed gender issue with deoriented bodies, again, aaand again\ fixed gender issue with deoriented bodies, again, aaand again\
trained on 100 garments\ trained on 100 garments\
remove multiplications with garment uv mask') remove multiplications with garment uv mask\
')
artifact = wandb.Artifact('code', type='code') artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_body.py') artifact.add_file('auto_encoder_body.py')
artifact.add_file('losses.py') artifact.add_file('losses.py')
@@ -98,7 +99,7 @@ def train_step(batch):
with tf.GradientTape() as tape: with tf.GradientTape() as tape:
pred = model(batch['uv_body_posed']) pred = model(batch['uv_body_posed' + resolution_modifier])
# pred = (pred + zero_body) * batch['mesh_mask'][:,:,tf.newaxis] # pred = (pred + zero_body) * batch['mesh_mask'][:,:,tf.newaxis]
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
@@ -109,13 +110,14 @@ def train_step(batch):
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones_like(batch['mesh_mask']))) e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones_like(batch['mesh_mask'])))
c_loss = 0 # 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask'])) c_loss = 0 # 1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
# l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['mesh_unposed'] - pred)), axis=-1) * batch['mesh_mask']) # l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['mesh_unposed'] - pred)), axis=-1) * batch['mesh_mask'])
l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed'] - pred) * uv_mask_body[tf.newaxis,tf.newaxis,:,:])), axis=1)) l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed' + resolution_modifier] - pred) * uv_mask_body[tf.newaxis,tf.newaxis,:,:])), axis=1))
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed'] * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, max_val=2)) # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed' + resolution_modifier] * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, max_val=2))
ms_ssim = tf.reduce_mean(tf.image.ssim(tf.transpose(pred * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, tf.transpose(batch['uv_body_posed' + resolution_modifier] * uv_mask_body[tf.newaxis,tf.newaxis,:,:], [0,2,3,1]) + 1, max_val=2))
fft_loss = 0 # tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_body_posed'], tf.zeros_like(batch['uv_body_posed'])))), axis=[1,2,3])) fft_loss = 0 # tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_body_posed'], tf.zeros_like(batch['uv_body_posed'])))), axis=[1,2,3]))
# loss = fft_loss_pc # + e_loss + c_loss # loss = fft_loss_pc # + e_loss + c_loss
loss = 1e3 * (1 - ms_ssim) loss = 1e3 * (1 - ms_ssim)
mesh_orig_2 = uv2mesh(batch['uv_body_posed']) mesh_orig_2 = uv2mesh(batch['uv_body_posed' + resolution_modifier])
# mesh_skinned = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred) # mesh_skinned = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred)
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask']) metric_2 = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig_2 - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
@@ -153,12 +155,12 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str): def plot_everything(batch: dict, type: str):
pred = model(batch['uv_body_posed']) pred = model(batch['uv_body_posed' + resolution_modifier])
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_body_posed'] mesh_orig = batch['mesh_body_posed']
orig_img = tf.transpose(batch['uv_body_posed'], [0,2,3,1]) orig_img = tf.transpose(batch['uv_body_posed' + resolution_modifier], [0,2,3,1])
pred_img = tf.transpose(pred, [0,2,3,1]) pred_img = tf.transpose(pred, [0,2,3,1])
img = tf.concat((orig_img, pred_img), axis=2) img = tf.concat((orig_img, pred_img), axis=2)
imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)] imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
@@ -170,7 +172,8 @@ def plot_everything(batch: dict, type: str):
# gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])] # gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
gt = [mesh_orig[i].numpy()[:, :] for i in range(pred.shape[0])] gt = [mesh_orig[i].numpy()[:, :] for i in range(pred.shape[0])]
body = batch['mesh_body_posed'].numpy() body = uv2mesh(batch['uv_body_posed' + resolution_modifier]).numpy()
# body = batch['mesh_body_posed'].numpy()
# mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy() # mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
# garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])] # garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
garment_posed = garment_unposed garment_posed = garment_unposed
@@ -222,37 +225,38 @@ if __name__ == '__main__':
args = vars(parser.parse_args()) args = vars(parser.parse_args())
smpl = SMPL() smpl = SMPL()
uv_mask_body = np.load('custom_utils/mask_256x512.npy') # uv_mask_body = np.load('custom_utils/mask_256x512.npy')
uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2) # uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
uv_mask_body = tf.ones((128, 64))
resolution_modifier = '_128'
resolutions = [128]
factory = DatasetFactory() factory = DatasetFactory()
signature_dynamic = { 'id': (), signature_dynamic = { 'id': (),
'pose_shape': (72 + 11,), 'pose_shape': (72 + 11,),
# 'uv_posed': (3, 512, 256,), 'uv_body_posed_128': (3, 512, 256,),
'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,), 'mesh_body_posed': (14475, 3,),
# 'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,),
# 'uv_unposed': (3, 512, 256,)
} }
signature_static = { #'uv_static': (3, 512, 256,), signature_static = {
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,), 'mesh_mask': (14475,),
# 'uv_mask': (512, 256,), # 'uv_mask': (512, 256,),
# 'uv_shaped_body': (3, 512, 256,),
} }
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic, use_filled=True) dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic, use_filled=True, resolutions=resolutions)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=True) static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=True, resolutions=resolutions)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=32, signature=signature_dynamic) dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=32, signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static) static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
model = AutoEncoder() # model = AutoEncoder()
# lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0005, decay_steps=30000, decay_rate=0.1, staircase=False) # model = AutoEncoder32()
model = AutoEncoder128()
optimizer = tf.optimizers.Adam(learning_rate=0.0001) optimizer = tf.optimizers.Adam(learning_rate=0.0001)
folder = 'models/ae_logs/' folder = 'models/ae_logs/'
ckpt = tf.train.Checkpoint(model=model) ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3) manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
ckpt.restore(manager.latest_checkpoint)
train() train()
auto_encoder_curv.py
+4 -117
View File
@@ -5,7 +5,7 @@ import wandb
import random import random
import numpy as np import numpy as np
import tensorflow as tf import tensorflow as tf
os.environ["CUDA_VISIBLE_DEVICES"] = "0" os.environ["CUDA_VISIBLE_DEVICES"] = "1"
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
@@ -13,13 +13,13 @@ import plotly.graph_objects as go
import plotly.express as ex import plotly.express as ex
from tf_utils import check_tf_version, allow_memory_growth from tf_utils import check_tf_version, allow_memory_growth
from losses import collision_loss as collision_penalty
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
from dataset import DatasetFactory, StaticDataset from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL from custom_utils.smpl import SMPL
from models.ae_curv_model import AutoEncoder
name='ae_curvature_7_2'
name='ae_curvature_x_2'
testing = False testing = False
def init_wandb(): def init_wandb():
if not testing: if not testing:
@@ -35,43 +35,6 @@ def init_wandb():
uv2mesh = UV2Mesh(res=256) uv2mesh = UV2Mesh(res=256)
class DatasetSimple(tf.data.Dataset):
def _generator(folder):
# Opening the file
folder = folder.decode("utf-8")
uv_folder = os.path.join(folder, 'uv_posed')
frames = os.listdir(uv_folder)
random.shuffle(frames)
for frame in frames:
id, frame_nr = frame.split('.')[0].split('_')
mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
yield uv_unposed, mesh_unposed, uv_mask, mesh_mask
def __new__(cls, folder, res=256):
return tf.data.Dataset.from_generator(
cls._generator,
output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32),
tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32),
tf.TensorSpec(shape = (14475,), dtype = tf.float32), ),
args=(folder, )
)
def filter_resolutions_featuremaps(resolutions, featuremaps, res):
index = resolutions.index(res)
filtered_resolutions = resolutions[:index + 1]
filtered_featuremaps = featuremaps[:index + 1]
return filtered_resolutions, filtered_featuremaps
def plot_3d(mesh_gt, mesh_pred, faces): def plot_3d(mesh_gt, mesh_pred, faces):
mesh_pred[:,1] += 2 mesh_pred[:,1] += 2
fig = go.Figure(data=[ fig = go.Figure(data=[
@@ -109,82 +72,6 @@ def plot_3d(mesh_gt, mesh_pred, faces):
return fig return fig
class Bottleneck(tf.keras.Model):
def __init__(self, dim, latent_dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
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=None):
x = inputs
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
return x
class Decoder(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(16 * 8 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((16, 8, bottleneck_size))
self.decoder = ConvNeXtReverseLabels(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])), nr_of_classes=7)
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder(tf.keras.Model):
def __init__(self):
super(AutoEncoder, self).__init__()
bottleneck_size = 32
latent_dim = 128
self.decoder = Decoder(bottleneck_size, latent_dim)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, bottleneck_size], include_top=False)
def call(self, inputs, training=None):
x, uv_mask = inputs
# x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
# x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
# x = x * uv_mask[:,tf.newaxis,:,:]
x = x * uv_mask[:,:,:,tf.newaxis]
return x
g_mean_iou = tf.keras.metrics.MeanIoU(num_classes=7) g_mean_iou = tf.keras.metrics.MeanIoU(num_classes=7)
auto_encoder_static.py
+20 -229
View File
@@ -1,3 +1,4 @@
import imp
import os import os
import time import time
import argparse import argparse
@@ -6,25 +7,18 @@ import random
import numpy as np import numpy as np
import tensorflow as tf import tensorflow as tf
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
from tensorboardX import SummaryWriter
import plotly.graph_objects as go import plotly.graph_objects as go
import plotly.express as ex import plotly.express as ex
from utils import str_to_bool from tf_utils import check_tf_version, allow_memory_growth
from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
from load_models import load_generator, load_discriminator
from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss
# from losses import collision_penalty as collision_penalty
from losses import collision_loss as collision_penalty
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
from stylegan2_ae.generator import Generator
from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder
from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor
from dataset import DatasetFactory, StaticDataset from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL from custom_utils.smpl import SMPL
from models.ae_static_model import AutoEncoder
name='ae_static_4'
name='ae_static_5'
def init_wandb(): def init_wandb():
os.environ["CUDA_VISIBLE_DEVICES"] = "0" os.environ["CUDA_VISIBLE_DEVICES"] = "0"
run = wandb.init(project='vae', name=name, run = wandb.init(project='vae', name=name,
@@ -32,7 +26,8 @@ def init_wandb():
fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \ fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
. lowered learning rate because it goes to nan. shallow third stage. slower lr decay. with validation. higher lr rate.\ . lowered learning rate because it goes to nan. shallow third stage. slower lr decay. with validation. higher lr rate.\
convnext first stage down/upscale by factor fo 2 instead of 4, and 5 stages instead of 4 so bottlneck same size\ convnext first stage down/upscale by factor fo 2 instead of 4, and 5 stages instead of 4 so bottlneck same size\
flattened latent dim to 128') flattened latent dim to 128\
trai on 100 garments')
artifact = wandb.Artifact('code', type='code') artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_static.py') artifact.add_file('auto_encoder_static.py')
artifact.add_file('losses.py') artifact.add_file('losses.py')
@@ -43,79 +38,6 @@ def init_wandb():
uv2mesh = UV2Mesh(res=256) uv2mesh = UV2Mesh(res=256)
class DatasetSimple(tf.data.Dataset):
def _generator(folder):
# Opening the file
folder = folder.decode("utf-8")
uv_folder = os.path.join(folder, 'uv_posed')
frames = os.listdir(uv_folder)
random.shuffle(frames)
for frame in frames:
id, frame_nr = frame.split('.')[0].split('_')
mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32)
uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32)
mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed']
uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1])
yield uv_unposed, mesh_unposed, uv_mask, mesh_mask
def __new__(cls, folder, res=256):
return tf.data.Dataset.from_generator(
cls._generator,
output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32),
tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32),
tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32),
tf.TensorSpec(shape = (14475,), dtype = tf.float32), ),
args=(folder, )
)
def filter_resolutions_featuremaps(resolutions, featuremaps, res):
index = resolutions.index(res)
filtered_resolutions = resolutions[:index + 1]
filtered_featuremaps = featuremaps[:index + 1]
return filtered_resolutions, filtered_featuremaps
# def plot_3d(mesh_gt, mesh_pred, faces):
# mesh_pred[:,1] += 2
# fig = go.Figure(data=[
# go.Mesh3d(
# # 8 vertices of a cube
# x=mesh_gt[:,0],
# y=mesh_gt[:,1],
# z=mesh_gt[:,2],
# # i, j and k give the vertices of triangles
# i = faces[:,0],
# j = faces[:,1],
# k = faces[:,2]
# ),
# go.Mesh3d(
# # 8 vertices of a cube
# x=mesh_pred[:,0],
# y=mesh_pred[:,1],
# z=mesh_pred[:,2],
# # i, j and k give the vertices of triangles
# i = faces[:,0],
# j = faces[:,1],
# k = faces[:,2],
# )
# ])
# fig.update_layout(scene_aspectmode='manual',
# scene_aspectratio=dict(x=1, y=2, z=1))
# fig.update_layout(
# scene = dict(
# xaxis = dict(nticks=4, range=[-1,1],),
# yaxis = dict(nticks=4, range=[-1,3],),
# zaxis = dict(nticks=4, range=[-1,1],),)
# )
# return fig
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color): def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
mesh_pred_unposed[:,1] += 1 mesh_pred_unposed[:,1] += 1
@@ -166,146 +88,14 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
return fig return fig
class Bottleneck(tf.keras.Model):
def __init__(self, dim, latent_dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
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=None):
x = inputs
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
return x
class Decoder(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(8 * 4 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((8, 4, bottleneck_size))
self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3, 3])), dims=list(reversed([bottleneck_size, 128, 128, 64, 16])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder(tf.keras.Model):
def __init__(self):
super(AutoEncoder, self).__init__()
bottleneck_size = 32
latent_dim = 128
self.encoder = ConvNeXt(depths=[3, 3, 9, 3, 3], dims=[32, 64, 128, 128, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.decoder = Decoder(bottleneck_size, latent_dim)
def call(self, inputs, training=None):
x, uv_mask = inputs
x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
# x = self.pooling(x)
# x = self.reshape(x)
# x = self.flatten(x)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
x = x * uv_mask[:,tf.newaxis,:,:]
return x
def draw_cicle(shape,diamiter):
'''
Input:
shape : tuple (height, width)
diameter : scalar
Output:
np.array of shape that says True within a circle with diamiter = around center
'''
assert len(shape) == 2
TF = np.zeros(shape,dtype=bool)
center = np.array(TF.shape)/2.0
for iy in range(shape[0]):
for ix in range(shape[1]):
TF[iy,ix] = (iy- center[0])**2 + (ix - center[1])**2 < diamiter **2
return (TF)
TFcircleIN = draw_cicle(shape=(256,256), diamiter=15)
TFcircleIN = np.ones((512,256,3), dtype=bool) * np.concatenate((TFcircleIN, TFcircleIN), axis=0)[:,:,np.newaxis]
TFcircleOUT = ~TFcircleIN
def fft_loss_per_channel(pred, gt):
pred = tf.complex(pred, tf.zeros_like(pred))
gt = tf.complex(gt, tf.zeros_like(gt))
fft_pred_top = tf.stack([tf.signal.fftshift(tf.signal.fft2d(pred[:,i,:256,:])) for i in range(3)], axis=-1)
fft_gt_top = tf.stack([tf.signal.fftshift(tf.signal.fft2d(gt[:,i,:256,:])) for i in range(3)], axis=-1)
fft_pred_bottom = tf.stack([tf.signal.fftshift(tf.signal.fft2d(pred[:,i,256:,:])) for i in range(3)], axis=-1)
fft_gt_bottom = tf.stack([tf.signal.fftshift(tf.signal.fft2d(gt[:,i,256:,:])) for i in range(3)], axis=-1)
fft_pred = tf.concat((fft_pred_top, fft_pred_bottom), axis=1)
fft_gt = tf.concat((fft_gt_top, fft_gt_bottom), axis=1)
fft_pred_low = tf.where(TFcircleIN, fft_pred, tf.zeros_like(fft_pred))
fft_gt_low = tf.where(TFcircleIN, fft_gt, tf.zeros_like(fft_gt))
fft_pred_high = tf.where(TFcircleOUT, fft_pred, tf.zeros_like(fft_pred))
fft_gt_high = tf.where(TFcircleOUT, fft_gt, tf.zeros_like(fft_gt))
low_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred_low - fft_gt_low), axis=[1,2,3]))
high_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred_high - fft_gt_high), axis=[1,2,3]))
total = tf.reduce_mean(tf.reduce_sum(tf.math.abs(fft_pred - fft_gt), axis=[1,2,3]))
return total, low_loss, high_loss
# ex.imshow(tf.math.real(gt[0,0,:256,:])).show()
# img = tf.signal.fftshift(fft_gt[0][0])
# ex.imshow(tf.math.abs(tf.math.log(img))).show()
# low = tf.where(TFcircleIN, img, tf.zeros_like(img))
# ex.imshow(tf.math.real(tf.signal.ifft2d(tf.signal.ifftshift(low)))).show()
# high = tf.where(TFcircleOUT, img, tf.zeros_like(img))
# # ex.imshow(tf.math.abs(tf.math.log(high))).show()
# ex.imshow(tf.math.real(tf.signal.ifft2d(tf.signal.ifftshift(high)))).show()
print()
@tf.function @tf.function
def train_step(batch): def train_step(batch):
with tf.GradientTape() as tape: with tf.GradientTape() as tape:
pred = model([batch['uv_static'], batch['uv_mask']]) pred = model(batch['uv_static'])
uv_mask = batch['uv_mask']
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_static'] mesh_orig = batch['mesh_static']
@@ -313,10 +103,9 @@ def train_step(batch):
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 0 #1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask'])) c_loss = 0 #1e-2 * tf.reduce_mean(collision_penalty(mesh_pred, batch['mesh_body_posed'], batch['mesh_mask']))
# l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['mesh_unposed'] - pred)), axis=-1) * batch['mesh_mask']) # l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['mesh_unposed'] - pred)), axis=-1) * batch['mesh_mask'])
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1) * batch['uv_mask']) l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1) * uv_mask)
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_static'], [0,2,3,1]) + 1, max_val=2)) ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) * uv_mask[:,:,:,tf.newaxis] + 1, tf.transpose(batch['uv_static'], [0,2,3,1]) * uv_mask[:,:,:,tf.newaxis] + 1, max_val=2))
fft_loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_static'], tf.zeros_like(batch['uv_static'])))), axis=[1,2,3])) fft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_static'], tf.zeros_like(batch['uv_static'])))), axis=[1,2,3]))
fft_loss_pc, fft_loss_pc_low, fft_loss_pc_high = 0,0,0 #fft_loss_per_channel(pred, batch['uv_unposed'])
loss = 1e3 * (1 - ms_ssim) loss = 1e3 * (1 - ms_ssim)
mesh_orig_2 = uv2mesh(batch['uv_static']) mesh_orig_2 = uv2mesh(batch['uv_static'])
@@ -329,20 +118,20 @@ def train_step(batch):
# gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients] # gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
optimizer.apply_gradients(zip(gradients, model.trainable_variables)) optimizer.apply_gradients(zip(gradients, model.trainable_variables))
return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/fft_oc_low': fft_loss_pc_low, 'train/fft_pc_high': fft_loss_pc_high, 'train/fft_loss': fft_loss, 'train/euclid_error_reconst': metric_2, 'train/euclid_error_posed': euclid_posed }
'train/fft_loss_per_channel': fft_loss_pc, 'train/euclid_error_posed': euclid_posed }
@tf.function @tf.function
def valid_step(batch): def valid_step(batch):
pred = model([batch['uv_static'], batch['uv_mask']]) pred = model(batch['uv_static'])
uv_mask = batch['uv_mask']
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_static'] mesh_orig = batch['mesh_static']
e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
c_loss = 0 c_loss = 0
l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1) * batch['uv_mask']) l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1) * uv_mask)
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_static'], [0,2,3,1]) + 1, max_val=2)) ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(tf.transpose(pred, [0,2,3,1]) + 1, tf.transpose(batch['uv_static'], [0,2,3,1]) + 1, max_val=2))
fft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_static'], tf.zeros_like(batch['uv_static'])))), axis=[1,2,3])) fft_loss = 0 #tf.reduce_mean(tf.reduce_sum(tf.math.abs(tf.signal.fft3d(tf.complex(pred, tf.zeros_like(pred))) - tf.signal.fft3d(tf.complex(batch['uv_static'], tf.zeros_like(batch['uv_static'])))), axis=[1,2,3]))
loss = 1e3 * (1 - ms_ssim) loss = 1e3 * (1 - ms_ssim)
@@ -357,7 +146,7 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str): def plot_everything(batch: dict, type: str):
pred = model([batch['uv_static'], batch['uv_mask']]) pred = model(batch['uv_static'])
mesh_pred = uv2mesh(pred) mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_static'] mesh_orig = batch['mesh_static']
@@ -414,6 +203,8 @@ if __name__ == '__main__':
smpl = SMPL() smpl = SMPL()
resolution_modifier = '_32'
resolutions = [32]
factory = DatasetFactory() factory = DatasetFactory()
signature_dynamic = { 'id': (), signature_dynamic = { 'id': (),
models/ae_body_model.py
+143
View File
@@ -78,3 +78,146 @@ class AutoEncoder(tf.keras.Model):
# x = x * uv_mask_body[tf.newaxis,tf.newaxis,:,:] # x = x * uv_mask_body[tf.newaxis,tf.newaxis,:,:]
return x return x
class Bottleneck32(tf.keras.Model):
def __init__(self, dim, latent_dim):
super(Bottleneck32, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
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=None):
x = inputs
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
return x
class Decoder32(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder32, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder32(tf.keras.Model):
def __init__(self):
super(AutoEncoder32, self).__init__()
bottleneck_size = 512
latent_dim = 128
self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck32(bottleneck_size, latent_dim)
self.decoder = Decoder32(bottleneck_size, latent_dim)
def call(self, inputs, training=None):
x = inputs
x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
# x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
x = tf.transpose(x, [0,3,1,2])
return x
class Bottleneck128(tf.keras.Model):
def __init__(self, dim, latent_dim):
super(Bottleneck128, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
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=None):
x = inputs
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
return x
class Decoder128(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder128, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
self.decoder = ConvNeXtReverse(depths=list(reversed([16, 16, 9, 9])), dims=list(reversed([bottleneck_size, 256, 128, 64])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder128(tf.keras.Model):
def __init__(self):
super(AutoEncoder128, self).__init__()
bottleneck_size = 512
latent_dim = 128
self.encoder = ConvNeXt(depths=[16, 16, 9, 9], dims=[64, 128, 256, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck128(bottleneck_size, latent_dim)
self.decoder = Decoder128(bottleneck_size, latent_dim)
def call(self, inputs, training=None):
x = inputs
x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
# x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5)
x = tf.transpose(x, [0,3,1,2])
return x
models/ae_curv_model.py
+130
View File
@@ -0,0 +1,130 @@
import tensorflow as tf
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverseLabels, Block
class Bottleneck(tf.keras.Model):
def __init__(self, dim, latent_dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
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=None):
x = inputs
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
return x
class Decoder(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(16 * 8 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((16, 8, bottleneck_size))
self.decoder = ConvNeXtReverseLabels(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([bottleneck_size, 128, 64, 16])), nr_of_classes=7)
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder(tf.keras.Model):
def __init__(self):
super(AutoEncoder, self).__init__()
bottleneck_size = 32
latent_dim = 128
self.decoder = Decoder(bottleneck_size, latent_dim)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, bottleneck_size], include_top=False)
def call(self, inputs, training=None):
x, uv_mask = inputs
# x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
# x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
# x = x * uv_mask[:,tf.newaxis,:,:]
x = x * uv_mask[:,:,:,tf.newaxis]
return x
class Decoder32(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder32, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
self.decoder = ConvNeXtReverseLabels(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])), nr_of_classes=7)
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder32(tf.keras.Model):
def __init__(self):
super(AutoEncoder32, self).__init__()
bottleneck_size = 512
latent_dim = 128
self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.decoder = Decoder32(bottleneck_size, latent_dim)
def call(self, inputs, training=None):
x, uv_mask = inputs
# x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
# x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
# x = x * uv_mask[:,tf.newaxis,:,:]
# x = x * uv_mask[:,:,:,tf.newaxis]
return x
models/ae_static_model.py
+123
View File
@@ -0,0 +1,123 @@
import tensorflow as tf
from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block
class Bottleneck(tf.keras.Model):
def __init__(self, dim, latent_dim):
super(Bottleneck, self).__init__()
self.bottleneck = tf.keras.Sequential(
[Block(dim=dim, drop_path=0,
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=None):
x = inputs
x = self.bottleneck(x)
x = self.flatten(x)
x = self.dense(x)
return x
class Decoder(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder, self).__init__()
self.dense = tf.keras.layers.Dense(8 * 4 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((8, 4, bottleneck_size))
self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3, 3])), dims=list(reversed([bottleneck_size, 128, 128, 64, 16])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder(tf.keras.Model):
def __init__(self):
super(AutoEncoder, self).__init__()
bottleneck_size = 32
latent_dim = 128
self.encoder = ConvNeXt(depths=[3, 3, 9, 3, 3], dims=[32, 64, 128, 128, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.decoder = Decoder(bottleneck_size, latent_dim)
def call(self, inputs, training=None):
x = inputs
x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
# x = x * uv_mask[:,tf.newaxis,:,:]
return x
class Decoder32(tf.keras.Model):
def __init__(self, bottleneck_size, latent_dim):
super(Decoder32, self).__init__()
self.dense = tf.keras.layers.Dense(4 * 2 * bottleneck_size)
self.reshape = tf.keras.layers.Reshape((4, 2, bottleneck_size))
self.decoder = ConvNeXtReverse(depths=list(reversed([27, 27])), dims=list(reversed([bottleneck_size, 256])))
def call(self, inputs, training=None):
x = inputs
x = self.dense(x)
x = self.reshape(x)
x = self.decoder(x)
return x
class AutoEncoder32(tf.keras.Model):
def __init__(self):
super(AutoEncoder32, self).__init__()
bottleneck_size = 512
latent_dim = 128
self.encoder = ConvNeXt(depths=[27, 27], dims=[256, bottleneck_size], include_top=False)
self.bottleneck = Bottleneck(bottleneck_size, latent_dim)
self.decoder = Decoder32(bottleneck_size, latent_dim)
def call(self, inputs, training=None):
x = inputs
x = tf.transpose(x, [0,2,3,1])
x = self.encoder(x, training=training)
x = self.bottleneck(x, training=training)
x = self.decoder(x, training=training)
x = tf.transpose(x, [0,3,1,2])
# x = x * uv_mask[:,:, tf.newaxis]
# x = x * uv_mask[:,tf.newaxis,:,:]
return x