making ready for skirts
This commit is contained in:
@@ -142,14 +142,15 @@ if __name__ == '__main__':
|
||||
parser.add_argument('--name', type=str)
|
||||
args = vars(parser.parse_args())
|
||||
|
||||
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)
|
||||
|
||||
# TODO: different resolutions not working yet, model should be ready though
|
||||
resolution_modifier = ''
|
||||
if config['resolution'] != 512:
|
||||
resolution_modifier = '_' + str(config['resolution'])
|
||||
|
||||
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)
|
||||
|
||||
factory = DatasetFactory()
|
||||
|
||||
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'])
|
||||
|
||||
@@ -22,7 +22,7 @@ from models.cond_encoder import AutoEncoder
|
||||
testing = False
|
||||
def init_wandb():
|
||||
if not testing:
|
||||
run = wandb.init(project='vae', name=config['name'], notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. same as 7')
|
||||
run = wandb.init(project='vae', name=config['name'], notes=config['notes'])
|
||||
artifact = wandb.Artifact('code', type='code')
|
||||
artifact.add_file('auto_encoder_curv.py')
|
||||
artifact.add_file('losses.py')
|
||||
@@ -79,7 +79,7 @@ def train_step(batch):
|
||||
with tf.GradientTape() as tape:
|
||||
pred = model([batch['uv_curvature'], batch['uv_mask']], training=True)
|
||||
|
||||
ce_uv_loss = cce(batch['uv_curvature'], pred)
|
||||
ce_uv_loss = scce(tf.math.argmax(batch['uv_curvature'], axis=-1), pred)
|
||||
ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))
|
||||
loss = ce_uv_loss
|
||||
|
||||
@@ -166,7 +166,7 @@ if __name__ == '__main__':
|
||||
parser.add_argument('--name', type=str)
|
||||
args = vars(parser.parse_args())
|
||||
|
||||
cce = tf.keras.losses.CategoricalCrossentropy()
|
||||
scce = tf.keras.losses.SparseCategoricalCrossentropy()
|
||||
|
||||
resolution_modifier = ''
|
||||
if config['resolution'] != 512:
|
||||
|
||||
+10
-4
@@ -12,13 +12,18 @@ from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch
|
||||
import plotly.graph_objects as go
|
||||
import plotly.express as ex
|
||||
|
||||
if config['is_skirt']:
|
||||
from losses import edge_loss_skirt as edge_loss
|
||||
else:
|
||||
from losses import edge_loss
|
||||
|
||||
from tf_utils import check_tf_version, allow_memory_growth
|
||||
from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss
|
||||
from dataset import DatasetFactory, StaticDataset
|
||||
from custom_utils.custom_utils import batch_to_dict, plot_3d
|
||||
from models.cond_encoder import AutoEncoder
|
||||
|
||||
|
||||
|
||||
def init_wandb():
|
||||
run = wandb.init(project='vae', name=config['name'],
|
||||
notes='convnext + convnext, add dense with tanh at the end of reverse convnext,\
|
||||
@@ -35,7 +40,7 @@ def init_wandb():
|
||||
artifact.add_file('convnext/models/convnext_tf.py')
|
||||
run.log_artifact(artifact)
|
||||
|
||||
uv2mesh = UV2Mesh(res=256)
|
||||
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'])
|
||||
|
||||
|
||||
|
||||
@@ -93,7 +98,7 @@ def plot_everything(batch: dict, type: str):
|
||||
img = tf.concat((orig_img, pred_img), axis=2)
|
||||
imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
|
||||
|
||||
F = np.load("custom_utils/smpl_faces.npy")
|
||||
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
|
||||
mask = batch['mesh_mask'].numpy().astype(bool)
|
||||
garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
|
||||
gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
|
||||
@@ -101,7 +106,7 @@ def plot_everything(batch: dict, type: str):
|
||||
body = gt
|
||||
garment_posed = mesh_pred
|
||||
collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
|
||||
collision_faces = [np.concatenate((faces[i], faces[i] + np.sum(mask[i])), axis=0) for i in range(len(garment_posed))]
|
||||
collision_faces = [np.concatenate((F, faces[i] + np.sum(mask[i])), axis=0) for i in range(len(garment_posed))]
|
||||
collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
|
||||
|
||||
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) }
|
||||
@@ -162,6 +167,7 @@ if __name__ == '__main__':
|
||||
resolution_modifier = ''
|
||||
if config['resolution'] != 512:
|
||||
resolution_modifier = '_' + str(config['resolution'])
|
||||
skirt_modifier = '_skirt' if config['is_skirt'] else ''
|
||||
|
||||
factory = DatasetFactory()
|
||||
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
from configs.config_base import config
|
||||
|
||||
config['gpu'] = '1'
|
||||
config['gpu'] = '3'
|
||||
config['lr'] = 0.0001
|
||||
config['name'] = 'ae_body_10_3'
|
||||
config['name'] = 'ae_body_10_3_skirt'
|
||||
|
||||
config['notes'] = 'convnext + convnext, add dense with tanh at the end of reverse convnext,\
|
||||
flattened bottleneck with 128 latent_dim\ remove global orientation\
|
||||
@@ -13,7 +13,7 @@ config['notes'] = 'convnext + convnext, add dense with tanh at the end of revers
|
||||
larger batch size, 32 instead of 4\
|
||||
'
|
||||
|
||||
config['dataset_in_memory'] = False
|
||||
config['dataset_in_memory'] = True
|
||||
config['batch_size'] = 32
|
||||
|
||||
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
from configs.config_base import config
|
||||
|
||||
config['gpu'] = '2'
|
||||
config['gpu'] = '0'
|
||||
config['lr'] = 0.001
|
||||
config['name'] = 'x' #'ae_curv_9'
|
||||
config['name'] = 'ae_curv_9_skirt' #'ae_curv_9'
|
||||
config['notes'] = 'convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. switched back to sparse CCE, normal one was giving NaN'
|
||||
|
||||
config['cond_ae']['is_label'] = True
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
from configs.config_base import config
|
||||
|
||||
config['gpu'] = '1'
|
||||
config['gpu'] = '2'
|
||||
config['lr'] = 0.0005
|
||||
config['name'] = 'ae_static_7_2'
|
||||
config['name'] = 'ae_static_7_2_skirt'
|
||||
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
|
||||
|
||||
+15
-2
@@ -2,15 +2,17 @@ config = dict()
|
||||
|
||||
|
||||
config['batch_size'] = 4
|
||||
config['train_dir'] = '/home/workspace/local-data/train_100'
|
||||
config['valid_dir'] = '/home/workspace/local-data/train_100'
|
||||
config['train_dir'] = '/home/workspace/local-data/train_skirt_25_2'
|
||||
config['valid_dir'] = '/home/workspace/local-data/train_skirt_25_2'
|
||||
config['use_filled'] = False
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = True
|
||||
|
||||
config['latent_dim'] = 128
|
||||
|
||||
config['cond_ae'] = dict()
|
||||
config['cond_ae']['is_label'] = False
|
||||
config['cond_ae']['upscale_output'] = False
|
||||
config['cond_ae']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae']['kernel_size'] = 7
|
||||
config['cond_ae']['bottleneck_size'] = 64
|
||||
@@ -19,3 +21,14 @@ config['cond_ae']['encoder_dims'] = [32, 64, 128, 256, config['cond_ae']['bottle
|
||||
config['cond_ae']['decoder_input'] = (8, 4, config['cond_ae']['bottleneck_size'])
|
||||
config['cond_ae']['decoder_depths'] = [3, 3, 9, 3, 3]
|
||||
config['cond_ae']['decoder_dims'] = [config['cond_ae']['bottleneck_size'], 256, 128, 64, 16]
|
||||
|
||||
config['cond_ae_body'] = dict()
|
||||
config['cond_ae_body']['latent_dim'] = config['latent_dim']
|
||||
config['cond_ae_body']['kernel_size'] = 7
|
||||
config['cond_ae_body']['upscale_output'] = True
|
||||
config['cond_ae_body']['bottleneck_size'] = 512
|
||||
config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3, 3]
|
||||
config['cond_ae_body']['encoder_dims'] = [32, 64, 128, 256, 256, config['cond_ae_body']['bottleneck_size']]
|
||||
config['cond_ae_body']['decoder_input'] = (4, 2, config['cond_ae_body']['bottleneck_size'])
|
||||
config['cond_ae_body']['decoder_depths'] = [3, 3, 9, 3]
|
||||
config['cond_ae_body']['decoder_dims'] = [config['cond_ae_body']['bottleneck_size'], 128, 64, 16]
|
||||
+11
-10
@@ -1,7 +1,7 @@
|
||||
from configs.config_base import config
|
||||
|
||||
config['gpu'] = '4'
|
||||
config['name'] = 'cvae_low_4_3_d32'
|
||||
config['gpu'] = '2'
|
||||
config['name'] = 'cvae_19_skirt'
|
||||
config['notes'] = '\
|
||||
offsets over average unposed garment. remove global orientation (in posed body input and in pose_shape parameters. \
|
||||
retrained body encoder with deoriented bodies \
|
||||
@@ -11,13 +11,13 @@ remove shape from unposed garment\
|
||||
training on 100 garments with offsets over the average of these\
|
||||
fixed autoencoders!!!!!!!!!\
|
||||
cvae only loads encoders instead AEs\
|
||||
full res loss\
|
||||
'
|
||||
config['dataset_in_memory'] = False
|
||||
config['dataset_in_memory'] = True
|
||||
|
||||
|
||||
config['resolution'] = 32
|
||||
|
||||
config['resolution'] = 512
|
||||
config['is_skirt'] = True
|
||||
config['reg_loss_coeff'] = 0 # set to 0 if don't want to use edge+collision
|
||||
|
||||
if config['resolution'] != 512:
|
||||
config['notes'] += '\
|
||||
@@ -50,13 +50,14 @@ if config['resolution'] != 512:
|
||||
config['signature_static']['uv_mask' + res_modifier] = (config['resolution'], config['resolution']//2)
|
||||
|
||||
if config['resolution'] == 512:
|
||||
config['avg_garment_path'] = 'custom_data/train_100_no_shape_avg_uv.npy'
|
||||
config['avg_garment_path'] = 'custom_data/train_skirt_25_2_no_shape_avg_uv.npy'
|
||||
else:
|
||||
config['avg_garment_path'] = ''
|
||||
|
||||
config['static_model'] = 'checkpoints/ae_static_7_2_x'
|
||||
config['body_model'] = 'checkpoints/ae_body_9_x'
|
||||
config['curvature_model'] = 'checkpoints/ae_curv_9_x'
|
||||
config['static_model'] = 'checkpoints/ae_static_7_2_skirt_x'
|
||||
config['body_model'] = 'checkpoints/ae_body_10_3_skirt_x'
|
||||
config['curvature_model'] = 'checkpoints/ae_curv_9_skirt_x'
|
||||
config['upscale_enhancer'] = 'checkpoints/up_2_3_skirt_x_d32'
|
||||
config['freeze_encoders'] = True
|
||||
|
||||
config['cvae'] = dict()
|
||||
|
||||
+18
-18
@@ -452,7 +452,7 @@ if __name__ == '__main__':
|
||||
parser.add_argument("-s", "--start", help = "start of range", default=0)
|
||||
parser.add_argument("-e", "--end", help = "end of range", default=10)
|
||||
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
|
||||
parser.add_argument("-sk", "--skirt", help = "end of range", default=1)
|
||||
parser.add_argument("-sk", "--skirt", help = "end of range", default=0)
|
||||
args = parser.parse_args()
|
||||
|
||||
# smpl = SMPL('FeatureLines/superSMPL')
|
||||
@@ -528,9 +528,9 @@ if __name__ == '__main__':
|
||||
mask = np.squeeze(sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].astype(bool))
|
||||
# body = readPC2(os.path.join(folder, id, 'Body_posed.pc16'), True)['V']
|
||||
gt = readPC2(os.path.join(folder, id, garment + '_registered.pc16'), True)['V']
|
||||
curvature = np.load(os.path.join(folder, id, 'curvature_label_' + curvature_label_type + '.npy'))
|
||||
mesh_unposed = [np.load(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])]
|
||||
mesh_unposed = np.stack([mesh_unposed[i]['mesh_unposed'] for i in range(gt.shape[0])])
|
||||
# curvature = np.load(os.path.join(folder, id, 'curvature_label_' + curvature_label_type + '.npy'))
|
||||
# mesh_unposed = [np.load(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])]
|
||||
# mesh_unposed = np.stack([mesh_unposed[i]['mesh_unposed'] for i in range(gt.shape[0])])
|
||||
mesh_unposed_no_shape = [np.load(os.path.join(folder, id, 'mesh_unposed_no_shape', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])]
|
||||
mesh_unposed_no_shape = np.stack([mesh_unposed_no_shape[i]['mesh_unposed_no_shape'] for i in range(gt.shape[0])])
|
||||
# body_no_orient = [np.load(os.path.join(folder, id, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz')) for i in range(gt.shape[0])]
|
||||
@@ -605,17 +605,17 @@ if __name__ == '__main__':
|
||||
# uv_body_static = convert_to_uv(uv_res, m2uv, uv_map, body_static, neighbours=3, mask=None, method='rbf', garment_uv_mask=None).astype(np.float32)
|
||||
# uv_body = convert_to_uv_batch(uv_res, m2uv, uv_map, body[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
# uv_posed_body_no_orient = convert_to_uv_batch(uv_res, m2uv, uv_map, body_no_orient[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
uv_gt = convert_to_uv_batch(uv_res, m2uv, uv_map, gt[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
# uv_gt = convert_to_uv_batch(uv_res, m2uv, uv_map, gt[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
# uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
# uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
uv_unposed_no_shape = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed_no_shape[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=garment_uv_mask, garment_uv_mask_border=garment_uv_mask_border).astype(np.float32)
|
||||
# uv_shaped_body = convert_to_uv(uv_res, m2uv, uv_map, shaped_body, neighbours=3, mask=None, method='rbf', garment_uv_mask=None).astype(np.float32)
|
||||
print(time.perf_counter() - start)
|
||||
|
||||
start = time.perf_counter()
|
||||
np.save(os.path.join(folder, id, garment + '_offsets_static_uv.npy'), uv_verts_static)
|
||||
# np.save(os.path.join(folder, id, garment + '_offsets_static_uv.npy'), uv_verts_static)
|
||||
# np.save(os.path.join(folder, id, 'body_static_uv.npy'), uv_body_static)
|
||||
np.save(os.path.join(folder, id, 'mask_uv.npy'), garment_uv_mask)
|
||||
# np.save(os.path.join(folder, id, 'mask_uv.npy'), garment_uv_mask)
|
||||
|
||||
# os.makedirs(os.path.join(folder, id, 'body_posed_uv'), exist_ok=True)
|
||||
# for i in range(uv_body.shape[0]):
|
||||
@@ -625,17 +625,17 @@ if __name__ == '__main__':
|
||||
# for i in range(uv_posed_body_no_orient.shape[0]):
|
||||
# np.savez_compressed(os.path.join(folder, id, 'uv_posed_body_no_orient', garment + '_' + str(i) + '.npz'), uv_posed_body_no_orient=uv_posed_body_no_orient[i, :, :, :])
|
||||
|
||||
os.makedirs(os.path.join(folder, id, 'gt_uv'), exist_ok=True)
|
||||
for i in range(uv_gt.shape[0]):
|
||||
np.savez_compressed(os.path.join(folder, id, 'gt_uv', garment + '_' + str(i) + '.npz'), uv_gt=uv_gt[i, :, :, :])
|
||||
# os.makedirs(os.path.join(folder, id, 'gt_uv'), exist_ok=True)
|
||||
# for i in range(uv_gt.shape[0]):
|
||||
# np.savez_compressed(os.path.join(folder, id, 'gt_uv', garment + '_' + str(i) + '.npz'), uv_gt=uv_gt[i, :, :, :])
|
||||
|
||||
os.makedirs(os.path.join(folder, id, 'curvature_uv'), exist_ok=True)
|
||||
for i in range(uv_curvature.shape[0]):
|
||||
np.savez_compressed(os.path.join(folder, id, 'curvature_uv', garment + '_' + str(i) + '.npz'), curvature_uv=uv_curvature[i, :, :])
|
||||
# os.makedirs(os.path.join(folder, id, 'curvature_uv'), exist_ok=True)
|
||||
# for i in range(uv_curvature.shape[0]):
|
||||
# np.savez_compressed(os.path.join(folder, id, 'curvature_uv', garment + '_' + str(i) + '.npz'), curvature_uv=uv_curvature[i, :, :])
|
||||
|
||||
os.makedirs(os.path.join(folder, id, 'uv_unposed'), exist_ok=True)
|
||||
for i in range(uv_unposed.shape[0]):
|
||||
np.savez_compressed(os.path.join(folder, id, 'uv_unposed', garment + '_' + str(i) + '.npz'), uv_unposed=uv_unposed[i, :, :, :])
|
||||
# os.makedirs(os.path.join(folder, id, 'uv_unposed'), exist_ok=True)
|
||||
# for i in range(uv_unposed.shape[0]):
|
||||
# np.savez_compressed(os.path.join(folder, id, 'uv_unposed', garment + '_' + str(i) + '.npz'), uv_unposed=uv_unposed[i, :, :, :])
|
||||
|
||||
os.makedirs(os.path.join(folder, id, 'uv_unposed_no_shape'), exist_ok=True)
|
||||
for i in range(uv_unposed_no_shape.shape[0]):
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
from faulthandler import disable
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
import plotly.graph_objects as go
|
||||
from plotly.subplots import make_subplots
|
||||
|
||||
|
||||
def m2uv_array(m2uv):
|
||||
@@ -19,12 +19,12 @@ def m2uv_array(m2uv):
|
||||
|
||||
|
||||
class UV2Mesh():
|
||||
def __init__(self, res=256):
|
||||
def __init__(self, res=256, is_skirt=False):
|
||||
super(UV2Mesh, self).__init__()
|
||||
self.res = res
|
||||
uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
|
||||
uv_map = np.load('custom_utils/stretched_uv_map_' + ('skirt' if is_skirt else '3') + '.npy')
|
||||
self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32)
|
||||
m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
|
||||
m2uv = np.load('custom_utils/m2uv' + ('_skirt' if is_skirt else '') + '.npy', allow_pickle=True).item()
|
||||
m2uv = m2uv_array(m2uv)
|
||||
self.uv_pixels = self.uv_map_discreet[m2uv]
|
||||
# self.batch_size = batch_size
|
||||
@@ -136,49 +136,56 @@ def maskFacesBatch(V, F, mask):
|
||||
|
||||
|
||||
def plot_3d(mesh_gt, mesh_pred_unposed, faces, collision, collision_faces, collision_color):
|
||||
mesh_pred_unposed[:,2] += 1
|
||||
collision[:,2] += 2
|
||||
fig = go.Figure(data=[
|
||||
go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=mesh_gt[:,0],
|
||||
y=mesh_gt[:,2],
|
||||
z=mesh_gt[:,1],
|
||||
# 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_unposed[:,0],
|
||||
y=mesh_pred_unposed[:,2],
|
||||
z=mesh_pred_unposed[:,1],
|
||||
# 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=collision[:,0],
|
||||
y=collision[:,2],
|
||||
z=collision[:,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = collision_faces[:,0],
|
||||
j = collision_faces[:,1],
|
||||
k = collision_faces[:,2],
|
||||
vertexcolor=collision_color
|
||||
)
|
||||
])
|
||||
fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ])
|
||||
|
||||
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],),)
|
||||
fig.add_trace(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=mesh_gt[:,0],
|
||||
y=mesh_gt[:,2],
|
||||
z=mesh_gt[:,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2]
|
||||
),
|
||||
row=1, col=1)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=mesh_pred_unposed[:,0],
|
||||
y=mesh_pred_unposed[:,2],
|
||||
z=mesh_pred_unposed[:,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2]
|
||||
),
|
||||
row=1, col=2)
|
||||
|
||||
fig.add_trace(go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=collision[:,0],
|
||||
y=collision[:,2],
|
||||
z=collision[:,1],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = collision_faces[:,0],
|
||||
j = collision_faces[:,1],
|
||||
k = collision_faces[:,2],
|
||||
vertexcolor=collision_color
|
||||
),
|
||||
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
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -26,10 +26,15 @@ import plotly.graph_objects as go
|
||||
import plotly.express as ex
|
||||
from PIL import Image
|
||||
|
||||
if config['is_skirt']:
|
||||
from losses import edge_loss_skirt as edge_loss
|
||||
else:
|
||||
from losses import edge_loss
|
||||
|
||||
from utils import str_to_bool
|
||||
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 edge_loss, uv_mask_loss, bend_loss, kl_loss
|
||||
from losses import uv_mask_loss, bend_loss, kl_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
|
||||
@@ -52,7 +57,7 @@ def init_wandb():
|
||||
artifact.add_file('convnext/models/convnext_tf.py')
|
||||
run.log_artifact(artifact)
|
||||
|
||||
uv2mesh = UV2Mesh(res=256)
|
||||
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'])
|
||||
|
||||
|
||||
def log_normal_pdf(sample, mean, logvar, raxis=1):
|
||||
@@ -72,10 +77,7 @@ def kl_loss(mean, logvar):
|
||||
def train_step(batch):
|
||||
|
||||
with tf.GradientTape() as tape:
|
||||
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
|
||||
|
||||
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']],
|
||||
training=True, use_mask=True)
|
||||
pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
|
||||
|
||||
mesh_pred = uv2mesh(pred)
|
||||
mesh_orig = batch['mesh_unposed']
|
||||
@@ -94,7 +96,7 @@ def train_step(batch):
|
||||
ms_ssim = tf.image.ssim_multiscale(pred + 1, batch['uv_unposed'] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))
|
||||
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale((tf.transpose(pred, [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], (tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1) * batch['uv_mask'][:,:,:,tf.newaxis], max_val=2))
|
||||
|
||||
reconstr_loss = l2_loss + e_loss + c_loss
|
||||
reconstr_loss = l2_loss + (e_loss + c_loss) * config['reg_loss_coeff']
|
||||
reconstr_loss = tf.reduce_mean(reconstr_loss)
|
||||
kl = kl_loss(mean, logvar)
|
||||
loss = reconstr_loss + balancing_coeff * kl
|
||||
@@ -152,8 +154,7 @@ def valid_step(batch):
|
||||
|
||||
|
||||
def plot_everything(batch: dict, type: str):
|
||||
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
|
||||
pred, mean, logvar, z = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], training=False, use_mask=True)
|
||||
pred, mean, logvar = model([batch['uv_unposed'], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
|
||||
|
||||
mesh_pred = uv2mesh(pred).numpy()
|
||||
mesh_orig = batch['mesh_unposed']
|
||||
@@ -169,7 +170,7 @@ def plot_everything(batch: dict, type: str):
|
||||
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])])
|
||||
|
||||
F = np.load("custom_utils/smpl_faces.npy")
|
||||
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
|
||||
mask = batch['mesh_mask'].numpy().astype(bool)
|
||||
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
|
||||
garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
|
||||
@@ -198,7 +199,7 @@ def sample(folder):
|
||||
load_model(model,folder)
|
||||
count = 0
|
||||
for batch in dataset_train:
|
||||
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
|
||||
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
|
||||
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
|
||||
nr_of_samples=2
|
||||
|
||||
@@ -322,6 +323,7 @@ if __name__ == '__main__':
|
||||
args = vars(parser.parse_args())
|
||||
|
||||
smpl = SMPL()
|
||||
skirt_modifier = '_skirt' if config['is_skirt'] else ''
|
||||
|
||||
factory = DatasetFactory()
|
||||
|
||||
|
||||
+27
-28
@@ -27,7 +27,12 @@ import plotly.express as ex
|
||||
from PIL import Image
|
||||
|
||||
from tf_utils import check_tf_version, allow_memory_growth, split_gpu_for_testing
|
||||
from losses import d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss, bend_loss
|
||||
|
||||
if config['is_skirt']:
|
||||
from losses import edge_loss_skirt as edge_loss
|
||||
else:
|
||||
from losses import edge_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
|
||||
@@ -51,7 +56,7 @@ def init_wandb():
|
||||
artifact.add_file('convnext/models/convnext_tf.py')
|
||||
run.log_artifact(artifact)
|
||||
|
||||
uv2mesh = UV2Mesh(res=256)
|
||||
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'])
|
||||
|
||||
|
||||
def log_normal_pdf(sample, mean, logvar, raxis=1):
|
||||
@@ -71,10 +76,8 @@ def kl_loss(mean, logvar):
|
||||
def train_step(batch):
|
||||
|
||||
with tf.GradientTape() as tape:
|
||||
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
|
||||
uv_curvature_low = tf.one_hot(batch['uv_curvature' + res_modifier], 7)
|
||||
pred, mean, logvar, z = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], uv_curvature_low,
|
||||
batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], use_mask=True)
|
||||
pred, mean, logvar = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier],
|
||||
batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
|
||||
|
||||
pred_enhanced = enhance(upscale_enhancer, pred, batch['uv_mask' + res_modifier], training=False)
|
||||
pred_full_res = tf.image.resize(pred_enhanced, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) * batch['uv_mask'][:,:,:,tf.newaxis]
|
||||
@@ -91,22 +94,18 @@ def train_step(batch):
|
||||
c_loss = collision_penalty(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
|
||||
l1_loss = tf.reduce_sum(tf.reduce_sum((tf.math.abs(batch['uv_unposed' + res_modifier] - pred)), axis=-1), axis=[1, 2])
|
||||
l2_loss = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed' + res_modifier] - pred) ** 2), axis=-1), axis=[1, 2])
|
||||
l2_loss_full = 1e3 * tf.reduce_mean(tf.reduce_sum(((batch['uv_unposed'] - pred_full_res) ** 2), axis=-1), axis=[1, 2])
|
||||
if resolution == 32:
|
||||
ms_ssim = tf.image.ssim(pred + 1, batch['uv_unposed' + res_modifier] + 1, max_val=3)
|
||||
ms_ssim_full = tf.image.ssim(pred_full_res + 1, batch['uv_unposed'] + 1, max_val=3)
|
||||
else:
|
||||
ms_ssim = tf.image.ssim_multiscale(pred + 1, batch['uv_unposed' + res_modifier] + 1, max_val=3, power_factors=(0.5, 0.3, 0.1, 0.1))
|
||||
|
||||
reconstr_loss = l2_loss_full #+ e_loss + c_loss
|
||||
reconstr_loss = l2_loss #+ e_loss + c_loss
|
||||
reconstr_loss = tf.reduce_mean(reconstr_loss)
|
||||
kl = kl_loss(mean, logvar)
|
||||
loss = reconstr_loss + balancing_coeff * kl
|
||||
ms_ssim = tf.reduce_mean(ms_ssim)
|
||||
ms_ssim_full = tf.reduce_mean(ms_ssim_full)
|
||||
l1_loss = tf.reduce_mean(l1_loss)
|
||||
l2_loss = tf.reduce_mean(l2_loss)
|
||||
l2_loss_full = tf.reduce_mean(l2_loss_full)
|
||||
e_loss = tf.reduce_mean(e_loss)
|
||||
c_loss = tf.reduce_mean(c_loss)
|
||||
mean = tf.reduce_mean(mean)
|
||||
@@ -116,7 +115,6 @@ def train_step(batch):
|
||||
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/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim,
|
||||
'train/l2_loss_full_res': l2_loss_full, 'train/ms_ssim_full_res': ms_ssim_full,
|
||||
'train/reconstr_loss': reconstr_loss,
|
||||
'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
|
||||
'train/reg_coeff': reg_coeff, 'train/l2_loss': l2_loss }
|
||||
@@ -163,29 +161,28 @@ def valid_step(batch):
|
||||
|
||||
|
||||
def plot_everything(batch: dict, type: str):
|
||||
uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
|
||||
uv_curvature_low = tf.one_hot(batch['uv_curvature' + res_modifier], 7)
|
||||
pred, mean, logvar, z = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], uv_curvature_low,
|
||||
batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], uv_curvature, batch['uv_mask']], use_mask=True)
|
||||
pred, mean, logvar = model([batch['uv_unposed' + res_modifier], batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier],
|
||||
batch['uv_mask' + res_modifier], batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']])
|
||||
|
||||
pred = enhance(upscale_enhancer, pred, batch['uv_mask' + res_modifier], training=False)
|
||||
mesh_pred = uv2mesh(pred).numpy()
|
||||
pred_full_res = tf.image.resize(pred, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) * batch['uv_mask'][:,:,:,tf.newaxis]
|
||||
mesh_pred = uv2mesh(pred_full_res).numpy()
|
||||
mesh_orig = batch['mesh_unposed']
|
||||
|
||||
orig_img = batch['uv_unposed' + res_modifier]
|
||||
uv_static_img = batch['uv_static' + res_modifier]
|
||||
uv_body_posed_img = batch['uv_body_posed' + res_modifier]
|
||||
pred_img = pred.numpy()
|
||||
pred_img[pred_img > 1] = 0
|
||||
pred_img[pred_img < -1] = 0
|
||||
img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
|
||||
imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
|
||||
orig_img = (batch['uv_unposed'].numpy() + 1) / 2
|
||||
uv_static_img = (batch['uv_static'].numpy() + 1) / 2
|
||||
uv_body_posed_img = (batch['uv_body_posed'].numpy() + 1) / 2
|
||||
uv_curvature = (np.tile(batch['uv_curvature'].numpy().argmax(-1)[:,:,:,np.newaxis], [1,1,1,3]) / 6)
|
||||
pred_img = pred_full_res.numpy()
|
||||
pred_img = (pred_img - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis]) / (pred_img.max(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis] - pred_img.min(axis=(1,2,3))[:,np.newaxis,np.newaxis,np.newaxis])
|
||||
img = tf.concat((uv_static_img, uv_body_posed_img, uv_curvature, orig_img, pred_img), axis=2)
|
||||
imgs = [wandb.Image(img[i]) for i in range(4)]
|
||||
|
||||
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])])
|
||||
|
||||
F = np.load("custom_utils/smpl_faces.npy")
|
||||
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy")
|
||||
mask = batch['mesh_mask'].numpy().astype(bool)
|
||||
garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
|
||||
garment_unposed = [garment_unposed[i] + shape_offsets.numpy()[i, mask[i], :] for i in range(mask.shape[0])]
|
||||
@@ -366,11 +363,13 @@ if __name__ == '__main__':
|
||||
|
||||
smpl = SMPL()
|
||||
|
||||
resolution = 32
|
||||
resolution = config['resolution']
|
||||
res_modifier = ''
|
||||
if resolution != 512:
|
||||
res_modifier = '_' + str(resolution)
|
||||
|
||||
skirt_modifier = '_skirt' if config['is_skirt'] else ''
|
||||
|
||||
factory = DatasetFactory()
|
||||
|
||||
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[512, config['resolution']], in_memory=config['dataset_in_memory'])
|
||||
@@ -380,7 +379,7 @@ if __name__ == '__main__':
|
||||
|
||||
model = CVAE(config)
|
||||
|
||||
upscale_enhancer = get_upscale_enhancer('checkpoints/up_2_3', res=resolution)
|
||||
upscale_enhancer = get_upscale_enhancer(config['upscale_enhancer'], res=resolution)
|
||||
|
||||
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
|
||||
balancing_coeff = config['starting_balancing_coeff']
|
||||
|
||||
@@ -142,6 +142,10 @@ def collision_loss(V, B, mask, thr=.004):
|
||||
laplacian_loss = sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo()
|
||||
edges = np.stack(np.nonzero(laplacian_loss.todense())).transpose()
|
||||
|
||||
laplacian_loss_skirt = sio.loadmat("custom_utils/ssmpl_skirt0.mat")['L'].tocoo()
|
||||
edges_skirt = np.stack(np.nonzero(laplacian_loss.todense())).transpose()
|
||||
|
||||
|
||||
def get_edge_mask(masks):
|
||||
return tf.reduce_min(tf.gather(masks, edges, axis=1), axis=-1)
|
||||
|
||||
@@ -158,6 +162,22 @@ def edge_loss(x, y, mask):
|
||||
loss = tf.reduce_sum(d_e ** 2, axis=-1)# / tf.reduce_sum(e_mask, axis=-1)
|
||||
return loss
|
||||
|
||||
def get_edge_mask_skirt(masks):
|
||||
return tf.reduce_min(tf.gather(masks, edges_skirt, axis=1), axis=-1)
|
||||
|
||||
def edge_loss_skirt(x, y, mask):
|
||||
e_mask = get_edge_mask_skirt(mask)
|
||||
x_e = tf.gather(x, edges_skirt[:,0], axis=1) - tf.gather(x, edges_skirt[:,1], axis=1)
|
||||
x_e = tf.reduce_sum(x_e ** 2, -1)
|
||||
x_e = tf.sqrt(x_e + 1e-9)
|
||||
y_e = tf.gather(y, edges_skirt[:,0], axis=1) - tf.gather(y, edges_skirt[:,1], axis=1)
|
||||
y_e = tf.reduce_sum(y_e ** 2, -1)
|
||||
y_e = tf.sqrt(y_e + 1e-9)
|
||||
d_e = y_e - x_e
|
||||
d_e = d_e * e_mask
|
||||
loss = tf.reduce_sum(d_e ** 2, axis=-1)# / tf.reduce_sum(e_mask, axis=-1)
|
||||
return loss
|
||||
|
||||
|
||||
class PairwiseDistance(tf.keras.layers.Layer):
|
||||
def __init__(self, **kwargs):
|
||||
|
||||
+11
-64
@@ -22,7 +22,7 @@ class CVAE(tf.keras.Model):
|
||||
if config['static_model'] != '':
|
||||
load_model(self.static_enc, config['static_model'] + '_enc')
|
||||
|
||||
self.body_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
|
||||
if config['body_model'] != '':
|
||||
load_model(self.body_enc, config['body_model'] + '_enc')
|
||||
|
||||
@@ -49,38 +49,6 @@ class CVAE(tf.keras.Model):
|
||||
self.avg_garment = None
|
||||
|
||||
|
||||
|
||||
# @tf.function
|
||||
def sample_single(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False):
|
||||
if not fix_noise:
|
||||
eps = tf.random.normal(shape=(nr_of_samples, self.latent_dim))
|
||||
else:
|
||||
eps = tf.random.normal(shape=(1, self.latent_dim))
|
||||
eps = tf.tile(eps, [nr_of_samples,1])
|
||||
|
||||
uv_template = tf.tile(tf.transpose(uv_template, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
|
||||
uv_body_posed = tf.tile(tf.transpose(uv_body_posed, [1,2,0])[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
|
||||
if len(uv_curv.shape) == 3:
|
||||
uv_curv = tf.tile(uv_curv[tf.newaxis,:,:,:], [nr_of_samples,1,1,1])
|
||||
uv_mask = tf.tile(uv_mask[tf.newaxis,:,:], [nr_of_samples,1,1])
|
||||
|
||||
# uv_template, uv_body_posed = tf.transpose(uv_template, [0,2,3,1]), tf.transpose(uv_body_posed, [0,2,3,1])
|
||||
z_static = self.static_enc(uv_template)
|
||||
z_body = self.body_enc(uv_body_posed)
|
||||
z_curv = self.curv_enc(uv_curv)
|
||||
|
||||
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
|
||||
z_cond = self.cond_encoder(z_cond)
|
||||
|
||||
z = tf.concat([eps, z_cond], axis=-1)
|
||||
|
||||
x = self.decoder(z)
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask=None, noise=None):
|
||||
if noise is None:
|
||||
eps = tf.random.normal(shape=(uv_template.shape[0], self.latent_dim))
|
||||
@@ -110,24 +78,9 @@ class CVAE(tf.keras.Model):
|
||||
return eps * tf.exp(logvar * .5) + mean
|
||||
|
||||
|
||||
def decode(self, z, apply_sigmoid=False):
|
||||
logits = self.decoder(z)
|
||||
if apply_sigmoid:
|
||||
probs = tf.sigmoid(logits)
|
||||
return probs
|
||||
return logits
|
||||
|
||||
|
||||
def encode_conds(self, uv_template, uv_body_posed, uv_curv):
|
||||
if uv_template.shape[1] == 3:
|
||||
uv_template = tf.transpose(uv_template, [0,2,3,1])
|
||||
if uv_body_posed.shape[1] == 3:
|
||||
uv_body_posed = tf.transpose(uv_body_posed, [0,2,3,1])
|
||||
if uv_curv.shape[1] == 3:
|
||||
uv_curv = tf.transpose(uv_curv, [0,2,3,1])
|
||||
|
||||
z_static = self.static_enc(uv_template)
|
||||
z_body = self.body_enc(uv_body_posed)
|
||||
z_static = self.static_enc(uv_template)
|
||||
z_curv = self.curv_enc(uv_curv)
|
||||
|
||||
z_cond = tf.concat([z_static, z_body, z_curv], axis=-1)
|
||||
@@ -136,20 +89,14 @@ class CVAE(tf.keras.Model):
|
||||
return z_cond
|
||||
|
||||
|
||||
def call(self, inputs, training=None, use_mask=True):
|
||||
def call(self, inputs, training=None):
|
||||
|
||||
if use_mask:
|
||||
if self.size == 512:
|
||||
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
|
||||
uv_template_high, uv_body_posed_high, uv_curv_high = uv_template, uv_body_posed, uv_curv
|
||||
else:
|
||||
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_template_high, uv_body_posed_high, uv_curv_high, uv_mask_high = inputs
|
||||
if self.size == 512:
|
||||
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask = inputs
|
||||
uv_template_high, uv_body_posed_high, uv_curv_high = uv_template, uv_body_posed, uv_curv
|
||||
else:
|
||||
if self.size == 512:
|
||||
uv_unposed, uv_template, uv_body_posed, uv_curv = inputs
|
||||
uv_template_high, uv_body_posed_high, uv_curv_high = uv_template, uv_body_posed, uv_curv
|
||||
else:
|
||||
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_template_high, uv_body_posed_high, uv_curv_high = inputs
|
||||
uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_template_high, uv_body_posed_high, uv_curv_high, uv_mask_high = inputs
|
||||
|
||||
|
||||
x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1)
|
||||
mean, logvar = self.encoder(x, training=training)
|
||||
@@ -161,8 +108,8 @@ class CVAE(tf.keras.Model):
|
||||
|
||||
if self.avg_garment is not None:
|
||||
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
||||
if use_mask:
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x, mean, logvar, z
|
||||
x = x * uv_mask[:,:,:,tf.newaxis]
|
||||
|
||||
return x, mean, logvar
|
||||
|
||||
|
||||
@@ -75,9 +75,6 @@ class UpscalerBasic(tf.keras.Model):
|
||||
def call(self, inputs, training=None):
|
||||
img, mask = inputs
|
||||
|
||||
if img.shape[1] == 3:
|
||||
img = tf.transpose(img, [0,2,3,1])
|
||||
|
||||
x = self.flatten(img)
|
||||
x = self.seq(x, training=training)
|
||||
x = self.reshape(x)
|
||||
|
||||
+391
@@ -0,0 +1,391 @@
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import argparse
|
||||
|
||||
import numpy as np
|
||||
import scipy.io as sio
|
||||
import scipy as sp
|
||||
import plotly.graph_objects as go
|
||||
from scipy.spatial.transform import Rotation
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
|
||||
import tensorflow as tf
|
||||
gpus = tf.config.experimental.list_physical_devices('GPU')
|
||||
if gpus:
|
||||
try:
|
||||
for gpu in gpus:
|
||||
tf.config.experimental.set_memory_growth(gpu, True)
|
||||
|
||||
except RuntimeError as e:
|
||||
print(e)
|
||||
|
||||
|
||||
from custom_utils.smpl import SMPL
|
||||
from custom_utils.smpl_2 import SMPL as SMPL_2
|
||||
from custom_utils.supersmpl_io import readPC2
|
||||
|
||||
|
||||
|
||||
# run = wandb.init(config=tf.compat.v1.flags.FLAGS, sync_tensorboard=True, project='unpose', name='new_smpl_2', notes='')
|
||||
# artifact = wandb.Artifact('code', type='code')
|
||||
# artifact.add_file('Data/unpose.py')
|
||||
# artifact.add_file('Data/verify_unposed.py')
|
||||
# artifact.add_file('FeatureLines/superSMPL/smpl.py')
|
||||
# artifact.add_file('FeatureLines/superSMPL/smpl_2.py')
|
||||
# run.log_artifact(artifact)
|
||||
|
||||
|
||||
def read_garment(folder, id):
|
||||
garment_dir = os.path.join(folder, id)
|
||||
|
||||
garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')]
|
||||
assert len (garments) == 1
|
||||
|
||||
garment = garments[0]
|
||||
|
||||
posed = readPC2(os.path.join(folder, id, garment + '_registered.pc16'), True)['V']
|
||||
info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
|
||||
pose, shape, gender = info['poses'].astype(np.float32), info['shape'], info['gender']
|
||||
shape = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32)
|
||||
mask = np.load(os.path.join(folder, id, garment + '_mask.npy'))
|
||||
offsets = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3]
|
||||
template = sio.loadmat(os.path.join(folder, id, 'info_body.mat'))['B']
|
||||
mesh_static = (template + offsets) * mask[:, np.newaxis]
|
||||
|
||||
posed, mesh_static = posed[:, mask, :], mesh_static[mask, :]
|
||||
mask_idx = np.array(range(14475))[mask]
|
||||
|
||||
return posed, mask, pose, shape, mesh_static, mask_idx
|
||||
|
||||
|
||||
def plot_3d(mesh_orig, mesh_pred, unposed, faces):
|
||||
_mesh_pred = mesh_pred.copy()
|
||||
_unposed = unposed.copy()
|
||||
_mesh_pred[:,1] += 1
|
||||
_unposed[:,1] += 3
|
||||
fig = go.Figure(data=[
|
||||
go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=mesh_orig[:,0],
|
||||
y=mesh_orig[:,1],
|
||||
z=mesh_orig[:,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],
|
||||
),
|
||||
go.Mesh3d(
|
||||
# 8 vertices of a cube
|
||||
x=_unposed[:,0],
|
||||
y=_unposed[:,1],
|
||||
z=_unposed[:,2],
|
||||
# i, j and k give the vertices of triangles
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2],
|
||||
# vertexcolor=collision_color
|
||||
)
|
||||
])
|
||||
|
||||
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
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/train_skirt')
|
||||
parser.add_argument("-s", "--start", help = "start of range", default=0)
|
||||
parser.add_argument("-e", "--end", help = "end of range", default=1)
|
||||
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
|
||||
args = parser.parse_args()
|
||||
|
||||
F = np.load('custom_utils/smpl_faces_skirt.npy')
|
||||
folder = args.folder
|
||||
# smpl = SMPL()
|
||||
smpl_2 = SMPL_2('custom_utils')
|
||||
iterations = 10000
|
||||
|
||||
ids = os.listdir(folder)[int(args.start):int(args.end)]
|
||||
|
||||
for id in tqdm(ids):
|
||||
if not os.path.isdir(os.path.join(folder, id)):
|
||||
continue
|
||||
|
||||
# if len(id) == 5:
|
||||
# continue
|
||||
|
||||
os.makedirs(os.path.join(folder, id, 'mesh_unposed'), exist_ok=True)
|
||||
|
||||
posed, mask, pose, shape, mesh_static, mask_idx = read_garment(folder, id)
|
||||
gender = int(shape[-1] > 0)
|
||||
init = Rotation.from_euler('zyz', [90,90,90], degrees=True).apply(mesh_static.copy()).astype(np.float32)
|
||||
|
||||
existing_frames = os.listdir(os.path.join(folder, id, 'mesh_unposed'))
|
||||
if len(existing_frames) == posed.shape[0]:
|
||||
continue
|
||||
|
||||
|
||||
R = smpl_2.run_rodrigues(np.transpose(pose, [1,0]))
|
||||
|
||||
frames = posed.shape[0]
|
||||
|
||||
# rot = Rotation.from_euler('z', [90], degrees=True)
|
||||
# rest_pose = np.zeros((1,72), dtype=np.float32)
|
||||
# R_rest = smpl.run_rodrigues(rest_pose)
|
||||
# body = smpl.run(R_rest, shape[tf.newaxis,:])[0].numpy()[:, mask,:]
|
||||
|
||||
# # verts = smpl.run(R, tf.tile(shape[tf.newaxis,:], [frames,1]))[0].numpy()[:,mask,:]
|
||||
# # verts = tf.stack([rot.apply(verts[i]).astype(np.float32) for i in range(frames)])
|
||||
# verts = tf.tile(body, [frames,1,1])
|
||||
|
||||
# start = time.perf_counter()
|
||||
# result = smpl.run_garment(R, tf.tile(shape[tf.newaxis,:], [frames,1]), verts, filter_idx=mask_idx)
|
||||
# orig_time = time.perf_counter() - start
|
||||
# print(orig_time)
|
||||
|
||||
# start = time.perf_counter()
|
||||
# result_2 = smpl_2.run_garment(R, shape[:10], gender, verts, filter_idx=mask_idx)
|
||||
# new_time = time.perf_counter() - start
|
||||
# print(new_time)
|
||||
|
||||
# # test = np.abs(result.numpy() - result_2.numpy()) > 1e-4
|
||||
# # x = np.sum(test, axis=-1).sum(-1).sum(-1)
|
||||
# euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((result - result_2) ** 2, axis=-1))).numpy()
|
||||
|
||||
# print('speedup: ' + str(orig_time/new_time) + ' euclid_error: ' + str(euclid_error))
|
||||
|
||||
# # continue
|
||||
|
||||
# batch = np.random.choice(np.array(range(frames)))
|
||||
# _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
|
||||
# fig = go.Figure(data=[
|
||||
# go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=result[batch,:,0],
|
||||
# y=result[batch,:,1],
|
||||
# z=result[batch,:,2],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2],
|
||||
# text='speedup: ' + str(orig_time/new_time) + ' euclid_error: ' + str(euclid_error),
|
||||
# ),
|
||||
# go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=result_2[batch,:,0],
|
||||
# y=result_2[batch,:,1],
|
||||
# z=result_2[batch,:,2],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2],
|
||||
# text='speedup: ' + str(orig_time/new_time) + ' euclid_error: ' + str(euclid_error),
|
||||
# )],
|
||||
# )
|
||||
|
||||
# wandb.log({id: fig})
|
||||
|
||||
# continue
|
||||
|
||||
|
||||
|
||||
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(1000.0, decay_steps=40000, decay_rate=0.1, staircase=False)
|
||||
optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule)
|
||||
unposed = tf.Variable(tf.tile(init[tf.newaxis,:,:], [frames,1,1]))
|
||||
|
||||
|
||||
@tf.function
|
||||
def skinning(verts):
|
||||
result = smpl_2.run_garment(R, shape[:10], gender, verts, filter_idx=mask_idx)
|
||||
return result
|
||||
|
||||
|
||||
@tf.function
|
||||
def error():
|
||||
pred = skinning(unposed)
|
||||
loss = tf.reduce_mean(tf.math.abs(pred - posed))
|
||||
|
||||
return loss
|
||||
|
||||
start = time.perf_counter()
|
||||
euclid_error = 10000
|
||||
|
||||
|
||||
while True:
|
||||
optimizer.minimize(error, [unposed])
|
||||
|
||||
i = optimizer.iterations.numpy()
|
||||
if i % 1000 == 0:
|
||||
pred = skinning(unposed)
|
||||
# loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(pred - posed) ** 2, axis=-1))
|
||||
loss = tf.reduce_mean(tf.math.abs(pred - posed))
|
||||
euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((posed - pred) ** 2, axis=-1)))
|
||||
print('id: ' + id + ' iteration: ' + str(i) + ' loss: ' + str(loss.numpy()) + ' error: ' + str(euclid_error.numpy()) + ' time: ' + str(time.perf_counter() - start))
|
||||
start = time.perf_counter()
|
||||
|
||||
if euclid_error < 0.0001:
|
||||
break
|
||||
|
||||
if euclid_error < 0.0005 and i > 15000:
|
||||
break
|
||||
|
||||
if euclid_error < 0.001 and i > 30000:
|
||||
break
|
||||
|
||||
if i > 40000:
|
||||
break
|
||||
|
||||
# b = np.random.choice(np.array(range(frames)))
|
||||
# _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
|
||||
# plot_3d(posed[b], pred.numpy()[b], unposed.numpy()[b], faces).show()
|
||||
|
||||
result = np.zeros((frames, 14475, 3), dtype = np.float32)
|
||||
result[:, mask,:] = unposed.numpy()
|
||||
for frame in range(frames):
|
||||
np.savez_compressed(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(frame) + '.npz'), mesh_unposed=result[frame])
|
||||
|
||||
# batch = 0
|
||||
# go.Figure(data=[
|
||||
# go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=posed[batch,:,0],
|
||||
# y=posed[batch,:,1],
|
||||
# z=posed[batch,:,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=pred[batch,:,0],
|
||||
# y=pred[batch,:,1],
|
||||
# z=pred[batch,:,2],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2],
|
||||
# )],
|
||||
# ).show()
|
||||
|
||||
|
||||
continue
|
||||
|
||||
|
||||
|
||||
offset = 1
|
||||
|
||||
|
||||
for frame in range(posed.shape[0]):
|
||||
|
||||
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(10.0, decay_steps=10000, decay_rate=0.1, staircase=False)
|
||||
optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule)
|
||||
unposed = tf.Variable(np.random.rand(offset, posed.shape[1], 3).astype(np.float32) * 2 - 1)
|
||||
# unposed = tf.Variable(init)
|
||||
|
||||
@tf.function
|
||||
def skinning(verts):
|
||||
# verts = smpl.run(R[frame:frame+offset], tf.tile(shape[tf.newaxis,:], [offset,1]))[0].numpy()[:,mask,:]
|
||||
|
||||
|
||||
# start = time.perf_counter()
|
||||
result = smpl.run_garment(R[frame:frame+offset], tf.tile(shape[tf.newaxis,:], [offset,1]), verts, filter_idx=mask_idx)
|
||||
# print(time.perf_counter() - start)
|
||||
|
||||
# start = time.perf_counter()
|
||||
# result_2 = smpl_2.run_garment(R[frame:frame+offset], shape[:10], gender, verts, filter_idx=mask_idx)
|
||||
# print(time.perf_counter() - start)
|
||||
|
||||
# test = np.abs(result.numpy() - result_2.numpy()) > 1e-4
|
||||
# x = np.sum(test, axis=-1).sum(-1).sum(-1)
|
||||
# euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((result - result_2) ** 2, axis=-1)))
|
||||
|
||||
# batch = 25
|
||||
# _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
|
||||
# go.Figure(data=[
|
||||
# go.Mesh3d(
|
||||
# # 8 vertices of a cube
|
||||
# x=result[batch,:,0],
|
||||
# y=result[batch,:,1],
|
||||
# z=result[batch,:,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=result_2[batch,:,0],
|
||||
# y=result_2[batch,:,1],
|
||||
# z=result_2[batch,:,2],
|
||||
# # i, j and k give the vertices of triangles
|
||||
# i = faces[:,0],
|
||||
# j = faces[:,1],
|
||||
# k = faces[:,2],
|
||||
# )]).show()
|
||||
|
||||
return result
|
||||
|
||||
@tf.function
|
||||
def error():
|
||||
pred = skinning(unposed)
|
||||
loss = tf.reduce_mean(tf.math.abs(pred - posed[frame]))
|
||||
|
||||
return loss
|
||||
|
||||
start = time.perf_counter()
|
||||
euclid_error = 10000
|
||||
|
||||
for i in range(iterations):
|
||||
optimizer.minimize(error, [unposed])
|
||||
if i % 50 == 0:
|
||||
pred = skinning(unposed)
|
||||
# loss = tf.reduce_mean(tf.reduce_sum(tf.math.abs(pred - posed) ** 2, axis=-1))
|
||||
loss = tf.reduce_mean(tf.math.abs(pred - posed[frame]))
|
||||
euclid_error = tf.reduce_mean(tf.math.sqrt(tf.reduce_sum((posed[frame] - pred) ** 2, axis=-1)))
|
||||
print('id: ' + id + ' frame: ' + str(frame) + ' iteration: ' + str(i) + ' loss: ' + str(loss.numpy()) + ' error: ' + str(euclid_error.numpy()) + ' time: ' + str(time.perf_counter() - start))
|
||||
start = time.perf_counter()
|
||||
|
||||
if euclid_error < 0.001:
|
||||
break
|
||||
|
||||
result = np.zeros((14475, 3), dtype = np.float32)
|
||||
result[mask,:] = unposed.numpy()
|
||||
np.savez_compressed(os.path.join(folder, id, 'mesh_unposed', id + '_' + str(frame) + '.npz'), mesh_unposed=result)
|
||||
|
||||
init = unposed
|
||||
|
||||
# pred, unposed = pred.numpy(), unposed.numpy()
|
||||
# _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
|
||||
# k = 0
|
||||
# plot_3d(posed[k], pred[k], unposed[k], faces).show()
|
||||
|
||||
# print()
|
||||
|
||||
|
||||
|
||||
+31
-18
@@ -1,5 +1,18 @@
|
||||
import os
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "4"
|
||||
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 time
|
||||
import argparse
|
||||
@@ -7,8 +20,7 @@ import wandb
|
||||
import random
|
||||
import cv2
|
||||
import numpy as np
|
||||
import tensorflow as tf
|
||||
from custom_utils.custom_utils import UV2Mesh, maskFacesBatch
|
||||
from custom_utils.custom_utils import UV2Mesh, load_model, maskFacesBatch
|
||||
import plotly.graph_objects as go
|
||||
import plotly.express as ex
|
||||
|
||||
@@ -20,15 +32,16 @@ from custom_utils.smpl import SMPL
|
||||
# from models.upscaler import Upscaler
|
||||
from models.upscaler import UpscalerBasic as Upscaler
|
||||
|
||||
name='up_2_3_d128'
|
||||
name='up_2_3_skirt_d32'
|
||||
def init_wandb():
|
||||
run = wandb.init(project='vae', name=name,
|
||||
notes='convnext upscaling\
|
||||
notes='\
|
||||
changed ms-ssim scale weights\
|
||||
inpainting style\
|
||||
MLP based\
|
||||
smaller batch\
|
||||
use mask for loss\
|
||||
started from up_2_3\
|
||||
')
|
||||
artifact = wandb.Artifact('code', type='code')
|
||||
artifact.add_file('upscale_test.py')
|
||||
@@ -38,7 +51,7 @@ use mask for loss\
|
||||
artifact.add_file('convnext/models/convnext_tf.py')
|
||||
run.log_artifact(artifact)
|
||||
|
||||
uv2mesh = UV2Mesh(res=256)
|
||||
uv2mesh = UV2Mesh(res=256, is_skirt=True)
|
||||
|
||||
|
||||
|
||||
@@ -87,9 +100,9 @@ def train_step(batch):
|
||||
mesh_pred = uv2mesh(pred)
|
||||
mesh_orig = batch['mesh_unposed']
|
||||
|
||||
l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(tf.transpose(batch['uv_unposed'], [0,2,3,1]) - pred * batch['uv_mask'][:,:,:,tf.newaxis]) )), axis=1))
|
||||
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred * batch['uv_mask'][:,:,:,tf.newaxis] + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
|
||||
# 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))
|
||||
l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_unposed'] - pred * batch['uv_mask'][:,:,:,tf.newaxis]) )), axis=1))
|
||||
# ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred * batch['uv_mask'][:,:,:,tf.newaxis] + 1, tf.transpose(batch['uv_unposed'], [0,2,3,1]) + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
|
||||
ms_ssim = tf.reduce_mean(tf.image.ssim(pred * batch['uv_mask'][:,:,:,tf.newaxis] + 1, batch['uv_unposed'] * batch['uv_mask'][:,:,:,tf.newaxis] + 1, max_val=2))
|
||||
loss = 1e3 * (1 - ms_ssim)
|
||||
|
||||
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'])
|
||||
@@ -124,7 +137,7 @@ def plot_everything(batch: dict, type: str):
|
||||
mesh_pred = uv2mesh(pred)
|
||||
mesh_orig = batch['mesh_unposed']
|
||||
|
||||
F = np.load("custom_utils/smpl_faces.npy")
|
||||
F = np.load("custom_utils/smpl_faces_skirt.npy")
|
||||
mask = batch['mesh_mask'].numpy().astype(bool)
|
||||
garment_unposed, faces = maskFacesBatch(mesh_pred.numpy(), F, mask)
|
||||
# gt = [mesh_orig[i].numpy()[mask[i], :] for i in range(pred.shape[0])]
|
||||
@@ -133,7 +146,7 @@ def plot_everything(batch: dict, type: str):
|
||||
figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i]) 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) }
|
||||
|
||||
orig_img = tf.transpose(batch['uv_unposed'], [0,2,3,1])
|
||||
orig_img = batch['uv_unposed']
|
||||
img = tf.concat((orig_img, pred), axis=2)
|
||||
imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
|
||||
|
||||
@@ -142,7 +155,7 @@ def plot_everything(batch: dict, type: str):
|
||||
|
||||
|
||||
def train():
|
||||
valid_iterator = iter(dataset_valid)
|
||||
# valid_iterator = iter(dataset_valid)
|
||||
|
||||
step = 0
|
||||
for batch in dataset_train:
|
||||
@@ -180,7 +193,7 @@ if __name__ == '__main__':
|
||||
# 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)
|
||||
|
||||
low_res = 128
|
||||
low_res = 32
|
||||
uv_mask_body = tf.ones((low_res, low_res//2))
|
||||
resolution_modifier = '_' + str(low_res)
|
||||
|
||||
@@ -197,10 +210,10 @@ if __name__ == '__main__':
|
||||
'uv_mask' + resolution_modifier: (low_res, low_res//2,),
|
||||
}
|
||||
|
||||
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[512, low_res])
|
||||
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False, resolutions=[512, low_res])
|
||||
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)
|
||||
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_skirt_25_2', batch_size=8, signature=signature_dynamic, use_filled=False, resolutions=[512, low_res])
|
||||
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_skirt_25_2', signature=signature_static, use_filled=False, resolutions=[512, low_res])
|
||||
dataset_valid = None #factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=32, signature=signature_dynamic)
|
||||
static_dataset_valid = None #StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
|
||||
|
||||
model = Upscaler(bottleneck_size=512, size=low_res)
|
||||
|
||||
@@ -209,6 +222,6 @@ if __name__ == '__main__':
|
||||
|
||||
ckpt = tf.train.Checkpoint(model=model)
|
||||
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
|
||||
ckpt.restore(manager.latest_checkpoint)
|
||||
load_model(model, 'checkpoints/up_2_3')
|
||||
|
||||
train()
|
||||
@@ -0,0 +1,83 @@
|
||||
# check if upscaler is better then filling
|
||||
import os
|
||||
|
||||
from models.upscaler import UpscalerBasic as Upscaler
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
|
||||
|
||||
import tqdm
|
||||
import numpy as np
|
||||
import plotly.graph_objects as go
|
||||
|
||||
from custom_utils.custom_utils import UV2Mesh, _maskFaces, load_model
|
||||
# from custom_utils.custom_utils import UV2Mesh, maskFacesBatch, _maskFaces
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# folder = '/home/workspace/data/dataset/skirt_test'
|
||||
folder = '/home/workspace/data/dataset/train_skirt_25_2'
|
||||
mask = np.load('custom_utils/mask_256x512_skirt.npy')
|
||||
F = np.load("custom_utils/smpl_faces_skirt.npy")
|
||||
uv2mesh = UV2Mesh(res=256, is_skirt=True)
|
||||
res = 32
|
||||
res_modifier = '_' + str(res)
|
||||
|
||||
|
||||
ids = os.listdir(os.path.join(folder))
|
||||
visited = set()
|
||||
meshes = []
|
||||
count = 0
|
||||
|
||||
for id in tqdm.tqdm(ids):
|
||||
garment_dir = os.path.join(folder, id)
|
||||
|
||||
if not os.path.isdir(garment_dir):
|
||||
continue
|
||||
|
||||
# if os.path.exists(os.path.join(folder, id, 'gt_uv')):
|
||||
# continue
|
||||
|
||||
garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and f != 'info.mat' and not f.endswith('.obj')]
|
||||
# assert len (garments) == 1
|
||||
if len (garments) != 1:
|
||||
continue
|
||||
garment = garments[0]
|
||||
|
||||
# id, frame_nr = frame.split('_')
|
||||
# frame_nr = frame_nr.split('.')[0]
|
||||
frame_nr = '0'
|
||||
print(id)
|
||||
uv_unposed_no_shape = np.load(os.path.join(folder, id, 'uv_unposed_no_shape', garment + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
|
||||
# uv_posed = np.load(os.path.join(folder, id, 'gt_uv', garment + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
|
||||
mesh_mask = np.load(os.path.join(folder, id, garment + '_mask.npy'))
|
||||
uv_mask = np.load(os.path.join(folder, id, 'mask_uv.npy'))
|
||||
|
||||
# body_posed_uv = np.load(os.path.join(folder, id, 'body_posed_uv', garment + '_' + frame_nr + '.npz'))
|
||||
# body_posed_uv = body_posed_uv['uv_body']
|
||||
|
||||
mesh = uv2mesh(uv_unposed_no_shape[np.newaxis,:,:,:]).numpy()[0]
|
||||
# mesh = uv2mesh(body_posed_uv[np.newaxis,:,:,:]).numpy()[0]
|
||||
|
||||
mesh, faces = _maskFaces(mesh, F, mesh_mask)
|
||||
# faces = F
|
||||
x = mesh
|
||||
meshes = [go.Mesh3d(
|
||||
x=x[:,0],
|
||||
y=x[:,2],
|
||||
z=x[:,1],
|
||||
i = faces[:,0],
|
||||
j = faces[:,1],
|
||||
k = faces[:,2]
|
||||
),]
|
||||
|
||||
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],),)
|
||||
)
|
||||
|
||||
fig.show()
|
||||
|
||||
print()
|
||||
Reference in New Issue
Block a user