Compare commits
9
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 luni
|
a82e1e7d9b | ||
|
luni
|
53b630c787 | ||
|
luni
|
cbfa129999 | ||
|
luni
|
61c15a3ec4 | ||
|
luni
|
408b984b4e | ||
|
luni
|
19cd69295d | ||
|
luni
|
cc723222dc | ||
|
luni
|
54a77273c5 | ||
|
luni
|
8c0436cea8 |
@@ -1,8 +1,3 @@
|
||||
nohup python cvae.py --name="lhunor_weights_test_6" > nohup/iccv/weights_test_6.out 2>&1 &
|
||||
nohup python evaluate_ae.py --name iccv_weights_test_6 > nohup/iccv/evaluate_iccv_weights_test_6.out 2>&1 &
|
||||
nohup python calculate_error.py --name iccv_weights_test_6 > nohup/iccv/error_iccv_weights_test_6.out 2>&1 &
|
||||
|
||||
|
||||
one final convolutional layer for smoothing, on the final coordinates
|
||||
|
||||
skirts!!!!
|
||||
|
||||
@@ -1,188 +0,0 @@
|
||||
import os
|
||||
import sys
|
||||
from configs.config_abl_all_post import config as config_a
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = '5'
|
||||
import tensorflow as tf
|
||||
|
||||
gpus = tf.config.experimental.list_physical_devices('GPU')
|
||||
if gpus:
|
||||
try:
|
||||
# Currently, memory growth needs to be the same across GPUs
|
||||
for gpu in gpus:
|
||||
tf.config.experimental.set_memory_growth(gpu, True)
|
||||
logical_gpus = tf.config.experimental.list_logical_devices('GPU')
|
||||
print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
|
||||
except RuntimeError as e:
|
||||
# Memory growth must be set before GPUs have been initialized
|
||||
print(e)
|
||||
|
||||
import glob
|
||||
import time
|
||||
import argparse
|
||||
import wandb
|
||||
import tqdm
|
||||
import random
|
||||
import numpy as np
|
||||
import plotly.graph_objects as go
|
||||
import plotly.express as ex
|
||||
from PIL import Image
|
||||
from postprocessing import fix_collision
|
||||
|
||||
if config_a['is_skirt']:
|
||||
from losses import edge_loss_skirt as edge_loss
|
||||
else:
|
||||
from losses import edge_loss
|
||||
|
||||
|
||||
|
||||
from losses import collision_ratio, kl_loss #, concept_loss
|
||||
# from losses import collision_penalty as collision_penalty
|
||||
from losses import collision_loss as collision_penalty
|
||||
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
|
||||
from dataset import DatasetFactory, StaticDataset
|
||||
from custom_utils.smpl import SMPL
|
||||
from models.cvae_model import CVAE
|
||||
from custom_utils.supersmpl_io import writeOBJ
|
||||
from plotly.subplots import make_subplots
|
||||
|
||||
|
||||
# uv2mesh = UV2Mesh(res=256, is_skirt=config_a['is_skirt'], is_adaptive=False)
|
||||
|
||||
|
||||
|
||||
def plot(mesh_gt, pred_a, pred_b, body, mask):
|
||||
fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
|
||||
|
||||
_, faces_garment = _maskFaces(np.zeros((14475, 3), np.float32), F, mask)
|
||||
faces = np.concatenate((F, faces_garment + 14475), axis=0)
|
||||
mesh_gt = np.concatenate((body, mesh_gt[mask, :]), axis=0)
|
||||
pred_a = np.concatenate((body, pred_a), axis=0)
|
||||
pred_b = np.concatenate((body, pred_b), axis=0)
|
||||
colors = np.concatenate((np.ones_like(body, dtype=np.int32)*100, np.ones_like(mesh_gt, dtype=np.int32)*230), axis=0)
|
||||
# colors = np.ones_like(mesh_gt, dtype=np.int32)*230
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
x=mesh_gt[:,0],
|
||||
y=mesh_gt[:,2],
|
||||
z=mesh_gt[:,1],
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
text="Ground Truth",
|
||||
vertexcolor=colors,
|
||||
),
|
||||
row=1, col=1)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
x=pred_a[:,0],
|
||||
y=pred_a[:,2],
|
||||
z=pred_a[:,1],
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
text="Pred A",
|
||||
vertexcolor=colors,
|
||||
),
|
||||
row=1, col=2)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
x=pred_b[:,0],
|
||||
y=pred_b[:,2],
|
||||
z=pred_b[:,1],
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
text="Pred B",
|
||||
vertexcolor=colors,
|
||||
),
|
||||
row=1, col=3)
|
||||
|
||||
|
||||
|
||||
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
|
||||
)
|
||||
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
|
||||
yaxis = dict(nticks=4, range=[-1,1],),
|
||||
zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
fig.update_layout(scene = helper,
|
||||
scene2 = helper,
|
||||
scene3 = helper,
|
||||
)
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
run = wandb.init(project='plots', name=config_a['name'], notes=config_a['notes'], dir='/wandb')
|
||||
|
||||
parser = argparse.ArgumentParser(description='')
|
||||
parser.add_argument('--name', type=str)
|
||||
args = vars(parser.parse_args())
|
||||
|
||||
skirt_modifier = '_skirt' if config_a['is_skirt'] else ''
|
||||
|
||||
folder = config_a['test_dir']
|
||||
# smpl = SMPL(use_one_hot=False)
|
||||
|
||||
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
|
||||
|
||||
iteration = ''
|
||||
|
||||
# load_model(model, 'checkpoints/' + 'abl_all_120k')
|
||||
|
||||
|
||||
factory = DatasetFactory()
|
||||
dataset_a = factory.get_dataset(folder=config_a['test_dir'],
|
||||
batch_size=config_a['batch_size'],
|
||||
signature=config_a['signature_dynamic'],
|
||||
resolutions=[config_a['resolution']],
|
||||
in_memory=False,
|
||||
do_repeat=False,
|
||||
use_one_hot=config_a['use_one_hot'],
|
||||
use_no_orient=config_a['remove_global_orientation'],
|
||||
shuffle=False,
|
||||
testing=True,
|
||||
# frames_file='frames_abl.txt',
|
||||
)
|
||||
static_dataset_a = StaticDataset(folder=config_a['test_dir'],
|
||||
signature=config_a['signature_static'],
|
||||
resolutions=[config_a['resolution']],
|
||||
in_memory=True)
|
||||
|
||||
iterator_a = iter(dataset_a)
|
||||
|
||||
|
||||
faces_folder = '/home/workspace/data/results/support/faces' + ('_skirt' if config_a['is_skirt'] else '')
|
||||
ids = os.listdir(faces_folder)
|
||||
ids = [id.split('.')[0] for id in ids]
|
||||
faces = {id : np.load(faces_folder + '/' + id + '.npy') for id in ids }
|
||||
|
||||
results_folder = '/home/workspace/data/results/post'
|
||||
|
||||
count = 0
|
||||
while True:
|
||||
batch = iterator_a.get_next()
|
||||
|
||||
batch = static_dataset_a.expand_batch(batch_to_dict(batch, config_a['signature_dynamic']))
|
||||
|
||||
mask = batch['mesh_mask'].numpy().astype(bool)
|
||||
gt = batch['mesh_posed'].numpy()
|
||||
body = batch['mesh_body_posed'].numpy()
|
||||
|
||||
for i in range(mask.shape[0]):
|
||||
id, frame_nr = batch['id'][i].numpy().decode("utf-8"), batch['frame_nr'][i].numpy().decode("utf-8")
|
||||
pred_a = np.load(results_folder + '/al_all_post-/' + id + '_' + frame_nr + '.npy')
|
||||
pred_b = np.load(results_folder + '/iccv_weights_test_6-/' + id + '_' + frame_nr + '.npy')
|
||||
# pred_c = np.load(results_folder + '/iccv_weights_test_6-/' + id + '_' + frame_nr + '.npy')
|
||||
|
||||
fig = plot(gt[i], pred_a, pred_b, body[i], mask[i])
|
||||
|
||||
wandb.log({f'{id}_{frame_nr}': fig}, step = count)
|
||||
count += 1
|
||||
|
||||
print(count)
|
||||
|
||||
@@ -7,7 +7,7 @@ config = dict()
|
||||
# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
|
||||
# DON'T FORGET TO CHECK AVG GARMENT USED
|
||||
|
||||
config['use_skinning'] = False # and deshaping
|
||||
config['use_skinning'] = True # and deshaping
|
||||
config['remove_global_orientation'] = True
|
||||
config['use_curvature'] = True
|
||||
config['use_extra_losses'] = True # edge, collision, 3d mesh
|
||||
@@ -19,9 +19,9 @@ config['only_curvatures'] = False
|
||||
ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
|
||||
assert ablation_count == 1 or ablation_count == 0
|
||||
|
||||
config['gpu'] = '5'
|
||||
config['name'] = 'iccv_weights_test_8'
|
||||
config['notes'] = 'smaller weight loss coeff, 1e-6, convolutional weight learning, changed convnext reverse'
|
||||
config['gpu'] = '1'
|
||||
config['name'] = 'al_all_post'
|
||||
config['notes'] = ''
|
||||
|
||||
config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv
|
||||
if config['use_skinning']:
|
||||
@@ -86,15 +86,14 @@ config['use_one_hot'] = False
|
||||
config['cvae'] = dict()
|
||||
config['cvae']['is_label'] = False
|
||||
config['cvae']['add_noise'] = True
|
||||
config['cvae']['noise_amount'] = 0.00001
|
||||
config['cvae']['noise_amount'] = 0.001
|
||||
|
||||
config['cvae']['kernel_size'] = 7
|
||||
config['cvae']['encoder_depths'] = [3, 3, 3, 9, 3]
|
||||
config['cvae']['encoder_dims'] = [32, 64, 128, 256, 64]
|
||||
config['cvae']['decoder_depths'] = [3, 9, 3, 3, 3, 3]
|
||||
config['cvae']['decoder_dims'] = [3, 32, 128, 128, 256, 512]
|
||||
config['cvae']['decoder_weights_depths'] = [3, 9, 3, 3, 3]
|
||||
config['cvae']['decoder_weights_dims'] = [24, 32, 64, 128, 256]
|
||||
config['cvae']['decoder_depths'] = config['cvae']['encoder_depths']
|
||||
config['cvae']['decoder_dims'] = config['cvae']['encoder_dims']
|
||||
|
||||
|
||||
config['cond_ae_static'] = dict()
|
||||
config['cond_ae_static']['is_label'] = False
|
||||
|
||||
@@ -19,7 +19,7 @@ config['only_curvatures'] = True
|
||||
ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
|
||||
assert ablation_count == 1 or ablation_count == 0
|
||||
|
||||
config['gpu'] = '5'
|
||||
config['gpu'] = '3'
|
||||
config['name'] = 'abl_overfit_jumpsuit_2'
|
||||
config['notes'] = 'remove other conditions from input as well'
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '0'
|
||||
config['name'] = 'pyr_abl_32'
|
||||
config['name'] = 'pyr_abl_32_r1'
|
||||
config['notes'] = ''
|
||||
|
||||
# DON'T FORGET TO RESET SSIM LOSS USAGE
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
config = dict()
|
||||
|
||||
config['gpu'] = '3'
|
||||
config['name'] = 'pyr_abl_32_skirt_2'
|
||||
config['notes'] = 'continued from pyr_abl_32_skirt-180k'
|
||||
config['name'] = 'pyr_abl_32_skirt_r1'
|
||||
config['notes'] = ''
|
||||
|
||||
# DON'T FORGET TO RESET SSIM LOSS USAGE
|
||||
# DON'T FORGET TO CHECK MODEL LOADING
|
||||
@@ -24,7 +24,10 @@ config['pyramid']['resolutions'] = [32]
|
||||
config['pyramid']['pretrained_pyramid'] = ''
|
||||
config['pyramid']['final_upscale_enhancer'] = 'checkpoints/32_512_skirt'
|
||||
|
||||
config['train_dir'] = '/home/workspace/local-data/train_skirt'
|
||||
config['pretrained_model'] = ''
|
||||
config['pretrained_model_iteration'] = ''
|
||||
|
||||
config['train_dir'] = '/home/workspace/data/dataset/local-cache/train_skirt'
|
||||
config['valid_dir'] = '/home/workspace/local-data/valid_skirt'
|
||||
config['test_dir'] = '/home/workspace/data/dataset/test_skirt_ready_2'
|
||||
config['batch_size'] = 4
|
||||
@@ -35,7 +38,7 @@ config['pyramid']['static_model'] = 'checkpoints/static_1_skirt'
|
||||
config['pyramid']['body_model'] = 'checkpoints/body_1_skirt_2_2'
|
||||
config['pyramid']['curvature_model'] = 'checkpoints/curv_1_skirt'
|
||||
|
||||
config['dataset_in_memory'] = False
|
||||
config['dataset_in_memory'] = True
|
||||
config['resolution'] = None
|
||||
config['latent_dim'] = 128
|
||||
|
||||
@@ -65,6 +68,7 @@ config['signature_static'] = { 'uv_static': (512, 256, 3),
|
||||
config['pyramid']['is_label'] = False
|
||||
config['pyramid']['latent_dim'] = config['latent_dim']
|
||||
config['pyramid']['add_noise'] = True
|
||||
config['pyramid']['noise_amount'] = 0.001
|
||||
|
||||
config['pyramid']['loss_weights'] = [1,1,1,1,1]
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@ from unittest import skip
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
|
||||
GlobalAveragePooling2D, Layer, Activation, Lambda,
|
||||
GlobalAveragePooling2D, Layer,
|
||||
LayerNormalization, Flatten, Reshape, UpSampling2D)
|
||||
|
||||
def drop_path(inputs, drop_prob, is_training):
|
||||
@@ -222,31 +222,29 @@ class ConvNeXtReverse(tf.keras.Model):
|
||||
|
||||
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., learn_weights=False,
|
||||
drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
|
||||
):
|
||||
super().__init__()
|
||||
self.upsample_layers = []
|
||||
|
||||
self.upsample_layers = [] # stem and 3 intermediate downsampling conv layers
|
||||
stem = tf.keras.Sequential([
|
||||
LayerNormalization(epsilon=1e-6),
|
||||
Conv2DTranspose(dims[0], kernel_size=2, strides=2, padding='same'),
|
||||
])
|
||||
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) - 2):
|
||||
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'),
|
||||
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)
|
||||
|
||||
downsample_layer = tf.keras.Sequential([
|
||||
Conv2DTranspose(dims[i+2], kernel_size=2, strides=2, padding='same'),
|
||||
LayerNormalization(epsilon=1e-6),
|
||||
])
|
||||
self.upsample_layers.append(downsample_layer)
|
||||
|
||||
self.stages = []
|
||||
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)):
|
||||
@@ -257,19 +255,12 @@ class ConvNeXtReverse(tf.keras.Model):
|
||||
self.stages.append(stage)
|
||||
cur += depths[i]
|
||||
|
||||
if not learn_weights:
|
||||
self.last = Dense(3, activation='tanh')
|
||||
# self.last = Activation('tanh')
|
||||
else:
|
||||
self.last = Activation('softmax')
|
||||
# self.last = Dense(24, activation='softmax')
|
||||
|
||||
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
|
||||
|
||||
|
||||
@@ -0,0 +1,283 @@
|
||||
import os
|
||||
import sys
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = '4'
|
||||
|
||||
|
||||
import tqdm
|
||||
import argparse
|
||||
import random
|
||||
import string
|
||||
import time
|
||||
import numpy as np
|
||||
import trimesh as tm
|
||||
import cmapy as cm
|
||||
import scipy as sp
|
||||
import scipy.io as sio
|
||||
import plotly.graph_objects as go
|
||||
from pyntcloud import PyntCloud
|
||||
from scipy.sparse import csr_matrix, coo_matrix
|
||||
from scipy.sparse.csgraph import floyd_warshall, dijkstra
|
||||
from plotly.subplots import make_subplots
|
||||
from custom_utils.custom_utils import _maskFaces
|
||||
from postprocessing import fix_collision
|
||||
|
||||
from custom_utils.supersmpl_io import readPC2, readPC2Frame, writePC2, readOBJ
|
||||
from custom_utils.custom_utils import _maskFaces
|
||||
from auto_encoder_curv import dice_dense
|
||||
|
||||
def postprocess(mesh, faces, mask=None, body=None, do_smoothing=True):
|
||||
if body is not None:
|
||||
mesh_fixed = fix_collision(mesh, body[np.newaxis,:,:], mask[np.newaxis,:])[0].numpy()[mask, :]
|
||||
else:
|
||||
mesh_fixed = mesh
|
||||
|
||||
if do_smoothing:
|
||||
m = tm.Trimesh(mesh_fixed, faces, process=False, validate=False)
|
||||
outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
|
||||
outliers = np.nonzero(outliers)[0]
|
||||
outliers_n_neighbours = get_closest_n_points(mesh_fixed, outliers, 2)
|
||||
outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
|
||||
faces_fixed = np.concatenate((faces, outliers_faces), axis=0)
|
||||
m = tm.Trimesh(mesh_fixed, faces_fixed, process=False, validate=False)
|
||||
tm.smoothing.filter_humphrey(m, iterations=1, alpha=0.6, beta=0.4)
|
||||
|
||||
mesh = m.vertices
|
||||
else:
|
||||
mesh = mesh_fixed
|
||||
|
||||
return mesh
|
||||
|
||||
def export_to_ply(mesh, file):
|
||||
ply = tm.exchange.ply.export_ply(mesh, encoding='binary', vertex_normal=True, include_attributes=True)
|
||||
output_file = open(file, "wb+")
|
||||
output_file.write(ply)
|
||||
output_file.close()
|
||||
|
||||
|
||||
def get_closest_n_points(verts, idx, n):
|
||||
dist = sp.spatial.distance.cdist(verts[idx, :], verts)
|
||||
dist[dist == 0] = 1e9
|
||||
idx_sort = np.argsort(dist, axis=-1)
|
||||
closest = idx_sort[:, :n]
|
||||
return closest
|
||||
|
||||
|
||||
def calculate_dist(folder, neighbourhood):
|
||||
out = dict()
|
||||
frames = sorted(os.listdir(os.path.join(folder, 'mesh_posed')))[int(args.start):int(args.end)]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
id = frame.split('_')[0]
|
||||
if id in out.keys():
|
||||
continue
|
||||
|
||||
mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
|
||||
k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood)
|
||||
out[id] = k_neighbors
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def histedges_equalN(x, nbin):
|
||||
npt = len(x)
|
||||
return np.interp(np.linspace(0, npt, nbin + 1),
|
||||
np.arange(npt),
|
||||
np.sort(x))
|
||||
|
||||
|
||||
def get_bins(values, nr_of_bins):
|
||||
flattened = []
|
||||
for c in values:
|
||||
flattened.extend(c.flatten())
|
||||
flattened = np.array(flattened)
|
||||
bins = histedges_equalN(flattened, nr_of_bins)
|
||||
return bins
|
||||
|
||||
|
||||
def generate_curvatures(v, faces, mask, neighbourhood=25, nr_of_labels=6):
|
||||
# id = frame.split('_')[0]
|
||||
# frame_nr = int(frame.split('_')[1].split('.')[0])
|
||||
|
||||
# mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
|
||||
# _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
|
||||
# v = np.load(os.path.join(folder, 'mesh_posed', f'{id}_{frame_nr}.npz'))['mesh_posed'][mask,:]
|
||||
|
||||
m = tm.Trimesh(v, faces, process=False, validate=False)
|
||||
|
||||
outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
|
||||
outliers = np.nonzero(outliers)[0]
|
||||
outliers_n_neighbours = get_closest_n_points(v, outliers, 2)
|
||||
outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
|
||||
faces = np.concatenate((faces, outliers_faces), axis=0)
|
||||
|
||||
verts = v
|
||||
|
||||
m = tm.Trimesh(verts, faces, process=False, validate=False)
|
||||
export_to_ply(m, temp_file)
|
||||
|
||||
cloud = PyntCloud.from_file(temp_file)
|
||||
|
||||
k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood)
|
||||
#k_neighbors = k_neighbors_cache[id] # np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood)
|
||||
ev = cloud.add_scalar_field("eigen_values", k_neighbors=k_neighbors)
|
||||
cloud.add_scalar_field("curvature", ev=ev)
|
||||
curvature = np.array(cloud.points["curvature(" + str(neighbourhood+1) + ")"])
|
||||
|
||||
bins = np.load(f'custom_data/train_bins.npy') # get_bins(curvature, nr_of_labels)
|
||||
labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)
|
||||
|
||||
frame_labels = np.zeros((curvature.shape[0]), dtype=np.float32)
|
||||
for j in range(bins.shape[0] - 2):
|
||||
frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
|
||||
frame_labels[curvature >= bins[-2]] = labels[-1]
|
||||
|
||||
return frame_labels
|
||||
|
||||
import tensorflow as tf
|
||||
|
||||
def visualize_curvature_labels(folder, results_folder, frame):
|
||||
id = frame.split('_')[0]
|
||||
frame_nr = frame.split('_')[1].split('.')[0]
|
||||
|
||||
mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
|
||||
_, faces = _maskFaces(np.zeros((14475,3)), F, mask)
|
||||
verts = np.load(os.path.join(folder, 'mesh_posed', f'{id}_{frame_nr}.npz'))['mesh_posed'][mask,:]
|
||||
pred = np.load(os.path.join(results_folder, frame + '.npy'))
|
||||
pred = postprocess(pred, faces)
|
||||
|
||||
curvature_labels_gt = generate_curvatures(verts, faces, mask, 25)
|
||||
curvature_labels_pred = generate_curvatures(pred, faces, mask, 25)
|
||||
# dice = dice_dense(curvature_labels_pred[None,:], curvature_labels_gt[None,:])
|
||||
dice = dice_dense(tf.one_hot(curvature_labels_pred, 7)[None,:], tf.one_hot(curvature_labels_gt, 7)[None,:]).numpy()
|
||||
colors_binned_all = []
|
||||
|
||||
for curvature_labels in [curvature_labels_gt, curvature_labels_pred]:
|
||||
colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, int(curvature_labels.max()+1))])
|
||||
colors_binned = np.zeros((curvature_labels.shape[0], 3))
|
||||
for i in range(int(curvature_labels.max())+1):
|
||||
# print(np.count_nonzero(curvature_labels == i))
|
||||
colors_binned[curvature_labels == i, :] = colors[i]
|
||||
colors_binned_all.append(colors_binned)
|
||||
|
||||
fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=-pred[:,0],
|
||||
y=pred[:,2],
|
||||
z=pred[:,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
vertexcolor=np.ones((verts.shape[0],3))*np.array([255,200,200])[None,:]
|
||||
),
|
||||
row=1, col=2)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=verts[:,0],
|
||||
y=verts[:,1],
|
||||
z=verts[:,2],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
vertexcolor=colors_binned_all[0]
|
||||
),
|
||||
row=1, col=1)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=-pred[:,0],
|
||||
y=pred[:,2],
|
||||
z=pred[:,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
vertexcolor=colors_binned_all[1],
|
||||
text=f'dice: {dice}'
|
||||
),
|
||||
row=1, col=3)
|
||||
|
||||
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
|
||||
scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
|
||||
# scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
|
||||
# scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1),
|
||||
)
|
||||
helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
|
||||
yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
|
||||
zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),)
|
||||
|
||||
fig.update_layout(scene = helper,
|
||||
scene2 = helper,
|
||||
scene3 = helper,
|
||||
# scene4 = helper,
|
||||
# scene5 = helper,
|
||||
)
|
||||
|
||||
|
||||
# meshes = [go.Mesh3d(
|
||||
# x=x[:,0],
|
||||
# y=x[:,1],
|
||||
# z=x[:,2],
|
||||
# i = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2],
|
||||
# vertexcolor=colors_binned
|
||||
# )]
|
||||
|
||||
# fig = go.Figure(data=meshes)
|
||||
# fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1))
|
||||
# fig.update_layout(
|
||||
# scene = dict(
|
||||
# xaxis = dict(nticks=4, range=[-1,1],),
|
||||
# yaxis = dict(nticks=4, range=[-1,1],),
|
||||
# zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
# )
|
||||
|
||||
return fig, dice
|
||||
|
||||
import wandb
|
||||
|
||||
def init_wandb():
|
||||
run = wandb.init(project='plots', name='review_4', dir='/home/workspace/local-data/temp/')
|
||||
artifact = wandb.Artifact('code', type='code')
|
||||
artifact.add_file('curvature_analysis.py')
|
||||
run.log_artifact(artifact)
|
||||
|
||||
if __name__ == '__main__':
|
||||
# init_wandb()
|
||||
|
||||
is_skirt = False
|
||||
skirt_modifier = '_skirt' if is_skirt else ''
|
||||
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
|
||||
# k_neighbours(k=25)
|
||||
|
||||
folder = '/home/workspace/local-data/train'
|
||||
results_folder = '/home/workspace/data/results/post/abl_overfit_jumpsuit_2-'
|
||||
|
||||
frames_file = 'frames_jumpsuit_30k.txt'
|
||||
frames = []
|
||||
with open(os.path.join(folder, frames_file), 'r') as f:
|
||||
lines = f.readlines()
|
||||
lines = [line.strip() for line in lines]
|
||||
frames = lines
|
||||
|
||||
temp_file = '/home/workspace/local-data/temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply'
|
||||
if not is_skirt:
|
||||
dist = np.load('custom_data/dijkstra.npy')
|
||||
else:
|
||||
dist = np.load('custom_data/dijkstra_skirt.npy')
|
||||
|
||||
i = 0
|
||||
while True:
|
||||
frame = np.random.choice(frames)
|
||||
# frame = '03265_193'
|
||||
frame = '03102_49'
|
||||
fig, dice = visualize_curvature_labels(folder, results_folder, frame)
|
||||
wandb.log({f'{frame}-dice:{dice}': fig}, step=i)
|
||||
i += 1
|
||||
print(frame)
|
||||
Binary file not shown.
@@ -2,7 +2,6 @@ import json
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
import scipy as sp
|
||||
import cmapy as cm
|
||||
import plotly.graph_objects as go
|
||||
from plotly.subplots import make_subplots
|
||||
|
||||
@@ -217,35 +216,6 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, colli
|
||||
return fig
|
||||
|
||||
|
||||
def plot_weights(weights, body, faces):
|
||||
fig = go.Figure()
|
||||
colors = np.array([cm.color('jet', 1 - w) for w in weights.numpy()])
|
||||
|
||||
# if flip_xy:
|
||||
# mesh_gt = flip_axis(mesh_gt)
|
||||
# mesh_pred_unposed = flip_axis(mesh_pred_unposed)
|
||||
# collision = flip_axis(collision, ind=14475)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
x=body[:,0],
|
||||
y=body[:,2],
|
||||
z=body[:,1],
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
vertexcolor=colors
|
||||
))
|
||||
|
||||
|
||||
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),)
|
||||
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
|
||||
yaxis = dict(nticks=4, range=[-1,1],),
|
||||
zaxis = dict(nticks=4, range=[-1,1],),)
|
||||
fig.update_layout(scene = helper,)
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
def plot_3d_collision(gt, collision_1, collision_2, collision_faces, collision_color, flip_xy=False):
|
||||
fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
|
||||
|
||||
|
||||
@@ -406,8 +406,7 @@ class SMPL:
|
||||
return T_out, J_out
|
||||
|
||||
|
||||
def run_garment_with_weights(self, pose, betas, v_cloth, weights, J=None):
|
||||
R = self.run_rodrigues(pose)
|
||||
def run_garment_with_weights(self, R, betas, v_cloth, weights, J=None):
|
||||
|
||||
if J is None:
|
||||
J = self.get_joints_template(betas)
|
||||
@@ -431,12 +430,13 @@ class SMPL:
|
||||
J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
|
||||
|
||||
T = tf.tensordot(weights, results, axes=((1), (0)))
|
||||
rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.shape[0], 1), dtype=tf.float32)), axis=1)
|
||||
rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
|
||||
v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
|
||||
v = tf.reshape(v, (-1, 4))[:, :3]
|
||||
|
||||
return v, tf.squeeze(J)
|
||||
|
||||
# gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
|
||||
T_out, J_out = tf.map_fn(single_run, (R, J, v_cloth, weights), (tf.float32, tf.float32))
|
||||
|
||||
J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import os
|
||||
import sys
|
||||
from configs.config_abl_all_post import config
|
||||
from configs.config_abl_overfit_jumpsuit_2 import config
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
|
||||
import tensorflow as tf
|
||||
@@ -38,10 +38,10 @@ else:
|
||||
from losses import collision_ratio, kl_loss, normal_loss #, concept_loss
|
||||
# from losses import collision_penalty as collision_penalty
|
||||
from losses import collision_loss as collision_penalty
|
||||
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, UV2MeshLabel, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision, plot_weights
|
||||
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
|
||||
from dataset import DatasetFactory, StaticDataset
|
||||
from custom_utils.smpl import SMPL
|
||||
from models.cvae_model import CVAENew as CVAE
|
||||
from models.cvae_model import CVAE
|
||||
|
||||
|
||||
def init_wandb():
|
||||
@@ -62,7 +62,6 @@ def init_wandb():
|
||||
run.log_artifact(artifact)
|
||||
|
||||
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False, is_smpl_orig=config['is_cape'])
|
||||
uv2mesh_label = UV2MeshLabel(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
|
||||
|
||||
|
||||
def kl_loss(mean, logvar):
|
||||
@@ -72,15 +71,9 @@ def kl_loss(mean, logvar):
|
||||
return kl_loss
|
||||
|
||||
|
||||
def blend_weights_reg(pred_w, gender=None):
|
||||
orig = smpl.weights
|
||||
l1 = tf.reduce_mean(tf.abs(pred_w - orig[None, 0, :, :]))
|
||||
return l1
|
||||
|
||||
|
||||
# 3d loss on posed !!!!!!!!!!!!
|
||||
|
||||
@tf.function
|
||||
# @tf.function
|
||||
def train_step(batch):
|
||||
|
||||
with tf.GradientTape() as tape:
|
||||
@@ -89,10 +82,9 @@ def train_step(batch):
|
||||
else:
|
||||
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
|
||||
|
||||
pred, mean, logvar, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
|
||||
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
mesh_weigths = uv2mesh_label(weights)
|
||||
|
||||
if config['use_skinning']:
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
@@ -102,31 +94,25 @@ def train_step(batch):
|
||||
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
|
||||
else:
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
||||
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
|
||||
|
||||
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_pred_posed = smpl.run_garment_with_weights(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped, mesh_weigths)
|
||||
mesh_pred_posed = mesh_pred
|
||||
|
||||
# is the output already masked?
|
||||
|
||||
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
|
||||
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
|
||||
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
|
||||
# l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
|
||||
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
|
||||
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
|
||||
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
|
||||
# conc_loss = concept_loss(pred, batch['uv_unposed'])
|
||||
kl = kl_loss(mean, logvar)
|
||||
weights_l = blend_weights_reg(mesh_weigths)
|
||||
|
||||
reconstr_loss = l2_mesh_loss + 1e-6 * weights_l
|
||||
reconstr_loss = l2_loss
|
||||
if config['use_extra_losses']:
|
||||
if 'block_collision_loss' in config.keys() and config['block_collision_loss'] == True:
|
||||
reconstr_loss = reconstr_loss + (e_loss) * config['reg_loss_coeff']
|
||||
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss) * config['reg_loss_coeff']
|
||||
else:
|
||||
reconstr_loss = reconstr_loss + (e_loss + c_loss) * config['reg_loss_coeff']
|
||||
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
|
||||
if config['use_norm_loss'] and step > 2000:
|
||||
reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
|
||||
|
||||
@@ -136,14 +122,14 @@ def train_step(batch):
|
||||
logvar = tf.reduce_mean(logvar)
|
||||
|
||||
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
|
||||
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_posed) ** 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'])
|
||||
gradients = tape.gradient(loss, 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/euclid_error': metric,
|
||||
'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio,
|
||||
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
|
||||
'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss, 'train/weights_loss': weights_l,
|
||||
# 'train/ms_ssim': ms_ssim, 'train/l2_loss': l2_loss,
|
||||
'train/l2_loss': l2_loss, 'train/l2_mesh_loss': l2_mesh_loss, 'train/normal_loss': n_loss,
|
||||
# 'train/ms_ssim': ms_ssim,
|
||||
}
|
||||
|
||||
|
||||
@@ -154,10 +140,9 @@ def valid_step(batch):
|
||||
else:
|
||||
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
|
||||
|
||||
pred, mean, logvar, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
|
||||
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
mesh_weigths = uv2mesh_label(weights)
|
||||
|
||||
if config['use_skinning']:
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
@@ -167,28 +152,22 @@ def valid_step(batch):
|
||||
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
|
||||
else:
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
||||
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
|
||||
|
||||
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_pred_posed = smpl.run_garment_with_weights(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped, mesh_weigths)
|
||||
mesh_pred_posed = mesh_pred
|
||||
|
||||
# is the output already masked?
|
||||
|
||||
e_loss = tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
|
||||
c_loss, c_ratio = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
|
||||
c_loss, c_ratio = tf.reduce_mean(c_loss), tf.reduce_mean(c_ratio)
|
||||
# l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
|
||||
l2_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(((uv_orig - pred) ** 2), axis=-1), axis=[1, 2]))
|
||||
l2_mesh_loss = tf.reduce_mean(1e3 * tf.reduce_mean(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
|
||||
n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
|
||||
# conc_loss = concept_loss(pred, batch['uv_unposed'])
|
||||
kl = kl_loss(mean, logvar)
|
||||
weights_l = blend_weights_reg(mesh_weigths)
|
||||
|
||||
reconstr_loss = l2_mesh_loss + 1e-6 * weights_l
|
||||
reconstr_loss = l2_loss
|
||||
if config['use_extra_losses']:
|
||||
reconstr_loss = reconstr_loss + (e_loss + c_loss) * config['reg_loss_coeff']
|
||||
reconstr_loss = reconstr_loss + l2_mesh_loss + (e_loss + c_loss) * config['reg_loss_coeff']
|
||||
if config['use_norm_loss']:
|
||||
reconstr_loss = reconstr_loss + n_loss * config['reg_loss_coeff']
|
||||
|
||||
@@ -198,12 +177,12 @@ def valid_step(batch):
|
||||
logvar = tf.reduce_mean(logvar)
|
||||
|
||||
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
|
||||
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_posed) ** 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'])
|
||||
# c_ratio = tf.reduce_mean(collision_ratio(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']))
|
||||
|
||||
return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/collision_loss': c_loss, 'valid/euclid_error': metric,
|
||||
'valid/reconstr_loss': reconstr_loss, 'valid/l2_mesh_loss': l2_mesh_loss, 'valid/collision_ratio': c_ratio,
|
||||
'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar, 'valid/weights_loss': weights_l,
|
||||
'valid/mean': mean, 'valid/logvar': logvar, 'valid/kl_loss': kl, 'valid/var': np.e ** logvar,
|
||||
'valid/normal_loss': n_loss,# 'valid/ms_ssim': ms_ssim,
|
||||
}
|
||||
|
||||
@@ -215,7 +194,7 @@ def test_step(batch):
|
||||
else:
|
||||
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
|
||||
|
||||
pred, mean, logvar, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
|
||||
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
|
||||
@@ -229,7 +208,7 @@ def test_step(batch):
|
||||
mesh_orig_shaped = mesh_orig + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_orig_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_orig_shaped)
|
||||
else:
|
||||
mesh_pred_posed = smpl.run_garment_with_weights(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred, weights)
|
||||
mesh_pred_posed = mesh_pred
|
||||
mesh_orig_posed = mesh_orig
|
||||
|
||||
mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
|
||||
@@ -271,10 +250,9 @@ def plot_everything(batch: dict, type: str):
|
||||
else:
|
||||
mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
|
||||
|
||||
pred, mean, logvar, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
|
||||
pred, mean, logvar = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
|
||||
|
||||
mesh_pred = uv2mesh(pred).numpy()
|
||||
mesh_weigths = uv2mesh_label(weights)
|
||||
|
||||
orig_img = uv_orig
|
||||
uv_static_img = batch['uv_static']
|
||||
@@ -299,12 +277,7 @@ def plot_everything(batch: dict, type: str):
|
||||
mesh_pred[i, mask[i], :] += shape_offsets.numpy()[i, mask[i], :]
|
||||
mesh_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred).numpy()
|
||||
else:
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
||||
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
|
||||
|
||||
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_posed = smpl.run_garment_with_weights(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped, mesh_weigths).numpy()
|
||||
mesh_posed = mesh_pred
|
||||
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
|
||||
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
|
||||
|
||||
@@ -323,16 +296,13 @@ def plot_everything(batch: dict, type: str):
|
||||
collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(garment_posed))]
|
||||
collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(garment_posed))]
|
||||
|
||||
# figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], garment_posed[i], faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
|
||||
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i],) for i in range(4) }
|
||||
figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i],) for i in range(4) }
|
||||
figs_weights = { type + '/weights_' + str(i): plot_weights(mesh_weigths[0,:,i], body[0], F) for i in range(24) }
|
||||
# figs_weights = { type + '/weights_' + str(i): plot_weights(smpl.weights[0,:,i], body[0], F) for i in range(24) }
|
||||
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
|
||||
figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i], flip_xy=(not config['use_skinning'] and config['remove_global_orientation'])) for i in range(4) }
|
||||
|
||||
mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0))
|
||||
logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0))
|
||||
|
||||
return { **figs, **figs_coll, **figs_weights, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
|
||||
return { **figs, **figs_coll, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
|
||||
|
||||
|
||||
def sample(folder):
|
||||
@@ -506,7 +476,7 @@ def test():
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
init_wandb()
|
||||
# init_wandb()
|
||||
parser = argparse.ArgumentParser(description='')
|
||||
parser.add_argument('--name', type=str)
|
||||
args = vars(parser.parse_args())
|
||||
|
||||
+18
-16
@@ -1,6 +1,6 @@
|
||||
import os
|
||||
import sys
|
||||
from configs.post.config_abl_pyr_64_128_5 import config
|
||||
from configs.config_abl_pyr_32_skirt import config
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
|
||||
import tensorflow as tf
|
||||
@@ -41,7 +41,7 @@ from losses import collision_loss as collision_penalty
|
||||
from custom_utils.custom_utils import batch_to_dict, load_model
|
||||
from dataset import DatasetFactory, StaticDataset
|
||||
from custom_utils.smpl import SMPL
|
||||
from models.pyramid_model_optimized import PyramidNew as Pyramid
|
||||
from models.pyramid_model_optimized import Pyramid
|
||||
from postprocessing import fix_collision
|
||||
|
||||
|
||||
@@ -59,9 +59,11 @@ def init_wandb():
|
||||
artifact.add_file('models/cond_encoder.py')
|
||||
artifact.add_file('dataset.py')
|
||||
artifact.add_file('convnext/models/convnext_tf.py')
|
||||
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.post.')]
|
||||
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs')]
|
||||
for config_file_name in config_file_names:
|
||||
artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
|
||||
file_name = '/'.join(config_file_name.split('.')) + '.py'
|
||||
if os.path.exists(file_name):
|
||||
artifact.add_file(file_name)
|
||||
run.log_artifact(artifact)
|
||||
|
||||
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
|
||||
@@ -208,7 +210,7 @@ def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, col
|
||||
return fig
|
||||
|
||||
|
||||
@tf.function
|
||||
# @tf.function
|
||||
def train_step(batch):
|
||||
levels_input, conds = group_inputs_for_levels(batch, config)
|
||||
|
||||
@@ -261,8 +263,8 @@ def train_step(batch):
|
||||
else:
|
||||
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
|
||||
|
||||
if config['use_norm_loss'] and step > 2000:
|
||||
reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
|
||||
# if config['use_norm_loss'] and step > 2000:
|
||||
# reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
|
||||
|
||||
kl_total += kl
|
||||
|
||||
@@ -334,8 +336,8 @@ def valid_step(batch):
|
||||
else:
|
||||
reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss)
|
||||
|
||||
if config['use_norm_loss'] and step > 2000:
|
||||
reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
|
||||
# if config['use_norm_loss'] and step > 2000:
|
||||
# reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff']
|
||||
|
||||
kl_total += kl
|
||||
|
||||
@@ -731,7 +733,7 @@ def speed_step(levels_input, conds):
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
init_wandb()
|
||||
# init_wandb()
|
||||
parser = argparse.ArgumentParser(description='')
|
||||
parser.add_argument('--name', type=str, default='lhunor')
|
||||
args = vars(parser.parse_args())
|
||||
@@ -790,7 +792,7 @@ if __name__ == '__main__':
|
||||
# generate_results()
|
||||
|
||||
ckpt = tf.train.Checkpoint(model=model)
|
||||
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/post/' + config['name'], max_to_keep=20)
|
||||
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=20)
|
||||
|
||||
# load_model(model, 'checkpoints/pyr_abl_32_skirt', iteration='180000')
|
||||
|
||||
@@ -807,18 +809,18 @@ if __name__ == '__main__':
|
||||
model.static_enc.trainable = False
|
||||
model.body_enc.trainable = False
|
||||
model.curv_enc.trainable = False
|
||||
|
||||
model.cond_encoder.trainable = False
|
||||
|
||||
# load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer'])
|
||||
# model.levels[0].upscale_enhancer.trainable = False
|
||||
|
||||
|
||||
for i in range(len(model.levels)):
|
||||
for i in range(len(model.levels)-1):
|
||||
load_model(model.levels[i].upscale_enhancer, config['pyramid']['levels'][i]['upscale_enhancer'])
|
||||
model.levels[i].upscale_enhancer.trainable = False
|
||||
# if config['pyramid']['resolutions'][-1] != 512:
|
||||
# load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
|
||||
# model.upscale_enhancer.trainable = False
|
||||
if config['pyramid']['resolutions'][-1] != 512:
|
||||
load_model(model.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
|
||||
model.upscale_enhancer.trainable = False
|
||||
|
||||
for i in range(len(model.levels)-1):
|
||||
model.levels[i].trainable = False
|
||||
|
||||
+2
-34
@@ -2,41 +2,9 @@ FROM tensorflow/tensorflow:2.9.0-gpu
|
||||
|
||||
RUN apt-get update
|
||||
RUN apt-get upgrade -y
|
||||
|
||||
# Add the current user to the image
|
||||
ARG USER_NAME=lhunor
|
||||
ARG USER_ID=11017
|
||||
ARG GROUP_NAME=guest
|
||||
ARG GROUP_ID=11000
|
||||
|
||||
# Update the system and install required packages
|
||||
RUN apt-get update && apt-get upgrade -y && apt-get install -y sudo git
|
||||
|
||||
# Create a new group with the same GID as the host
|
||||
RUN groupadd --gid $GROUP_ID $GROUP_NAME
|
||||
|
||||
# Create a new user with the same UID and GID as the host
|
||||
RUN useradd --uid $USER_ID --gid $GROUP_ID --shell /bin/bash --create-home $USER_NAME
|
||||
|
||||
# don't require password with sudo, for convenience
|
||||
# not the safest thing to do, but hopefully okay inside the container
|
||||
RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
|
||||
|
||||
# Add the new user to the sudo group
|
||||
RUN usermod -aG sudo $USER_NAME
|
||||
|
||||
# Set the default user to the new user
|
||||
USER $USER_NAME
|
||||
|
||||
WORKDIR //home-net
|
||||
|
||||
RUN apt-get install git ffmpeg libsm6 libxext6 libopenexr-dev zip unzip htop nano -y
|
||||
RUN pip install numpy matplotlib opencv-python wandb OpenEXR tensorflow-graphics scipy plotly pandas tensorboardX opencv-python cmapy
|
||||
|
||||
# defining the volumes that will be mounted
|
||||
VOLUME [ "/data-net" ]
|
||||
VOLUME [ "/data-local" ]
|
||||
VOLUME [ "/home-net" ]
|
||||
|
||||
# Start the container with a bash shell
|
||||
CMD ["/bin/bash"]
|
||||
WORKDIR /home/workspace/data
|
||||
# RUN wandb login
|
||||
+6
-14
@@ -1,6 +1,6 @@
|
||||
import os
|
||||
import sys
|
||||
from configs.config_abl_all_post import config
|
||||
from configs.config_abl_overfit_jumpsuit_2 import config
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
|
||||
import tensorflow as tf
|
||||
@@ -39,16 +39,15 @@ else:
|
||||
from losses import collision_ratio, kl_loss #, concept_loss
|
||||
# from losses import collision_penalty as collision_penalty
|
||||
from losses import collision_loss as collision_penalty
|
||||
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, UV2MeshLabel, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
|
||||
from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
|
||||
from dataset import DatasetFactory, StaticDataset
|
||||
from custom_utils.smpl import SMPL
|
||||
from models.cvae_model import CVAENew as CVAE
|
||||
from models.cvae_model import CVAE
|
||||
from custom_utils.supersmpl_io import writeOBJ
|
||||
# from unif_metric import chamfer_and_score
|
||||
|
||||
|
||||
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_adaptive=False, is_smpl_orig=config['is_cape'])
|
||||
uv2mesh_label = UV2MeshLabel(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
|
||||
|
||||
@tf.function
|
||||
def step(batch):
|
||||
@@ -57,10 +56,9 @@ def step(batch):
|
||||
else:
|
||||
uv_orig = batch['uv_posed']
|
||||
|
||||
pred, _, _, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
|
||||
pred, _, _ = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
mesh_weigths = uv2mesh_label(weights)
|
||||
|
||||
if config['use_skinning']:
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
@@ -70,13 +68,7 @@ def step(batch):
|
||||
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_pred_posed = smpl.run_garment(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)
|
||||
else:
|
||||
shape = batch['pose_shape'][:, 72:]
|
||||
gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
|
||||
shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
|
||||
|
||||
mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,tf.newaxis]
|
||||
mesh_pred_posed = smpl.run_garment_with_weights(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped, mesh_weigths)
|
||||
|
||||
mesh_pred_posed = mesh_pred
|
||||
|
||||
# mesh_pred_posed_fixed = fix_collision(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
|
||||
|
||||
@@ -112,7 +104,7 @@ if __name__ == '__main__':
|
||||
|
||||
iteration = ''
|
||||
|
||||
load_model(model, 'checkpoints/post/' + checkpoint_name, iteration=iteration)
|
||||
load_model(model, 'checkpoints/' + checkpoint_name, iteration=iteration)
|
||||
# load_model(model, 'checkpoints/' + 'abl_all_120k')
|
||||
|
||||
output_folder = '/home/workspace/data/results/post/' + checkpoint_name + '-' + iteration
|
||||
|
||||
+2
-2
@@ -171,13 +171,13 @@ class AEEncoderNew(tf.keras.Model):
|
||||
|
||||
class DecoderNew(tf.keras.Model):
|
||||
|
||||
def __init__(self, depths, dims, kernel_size=7, is_label=False, learn_weights=False):
|
||||
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, learn_weights=learn_weights)
|
||||
self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size)
|
||||
else:
|
||||
self.decoder = ConvNeXtReverseLabels(depths=depths, dims=dims, nr_of_classes=7)
|
||||
|
||||
|
||||
@@ -43,7 +43,7 @@ class CVAE(tf.keras.Model):
|
||||
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.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'])
|
||||
@@ -181,7 +181,6 @@ class CVAENew(tf.keras.Model):
|
||||
|
||||
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.decoder_weights = DecoderNew(config['cvae']['decoder_weights_depths'], config['cvae']['decoder_weights_dims'], config['cvae']['kernel_size'], learn_weights=True)
|
||||
|
||||
self.reshape = tf.keras.layers.Reshape((8, 4, 64))
|
||||
|
||||
@@ -194,12 +193,12 @@ class CVAENew(tf.keras.Model):
|
||||
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 = False
|
||||
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!'
|
||||
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):
|
||||
@@ -215,7 +214,6 @@ class CVAENew(tf.keras.Model):
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z, training=False)
|
||||
weights = self.decoder_weights(z, training=False)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
@@ -225,7 +223,7 @@ class CVAENew(tf.keras.Model):
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x, weights
|
||||
return x
|
||||
|
||||
|
||||
def reparameterize(self, mean, logvar, training, fix_noise=False):
|
||||
@@ -274,7 +272,6 @@ class CVAENew(tf.keras.Model):
|
||||
|
||||
z = tf.concat([z, z_cond], axis=-1)
|
||||
x = self.decoder(z, training=training)
|
||||
weights = self.decoder_weights(z, training=training)
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
@@ -284,5 +281,5 @@ class CVAENew(tf.keras.Model):
|
||||
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x, mean, logvar, weights
|
||||
return x, mean, logvar
|
||||
|
||||
|
||||
Reference in New Issue
Block a user