diff --git a/auto_encoder_body.py b/auto_encoder_body.py index b6f6db0..2861e03 100644 --- a/auto_encoder_body.py +++ b/auto_encoder_body.py @@ -25,9 +25,9 @@ from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict from custom_utils.smpl import SMPL -name='ae_body_7' +name='ae_body_7_2' def init_wandb(): - os.environ["CUDA_VISIBLE_DEVICES"] = "1" + os.environ["CUDA_VISIBLE_DEVICES"] = "4" run = wandb.init(project='vae', name=name, notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \ fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \ @@ -35,7 +35,7 @@ def init_wandb(): two more downscaling than up, so decoder output is resized wto final resolution\ larger batch_size.\ flattened bottleneck with 128 latent_dim\ remove global orientation\ - fixed gender issue with deoriented bodies, again\ + fixed gender issue with deoriented bodies, again, aaand again\ trained on 100 garments') artifact = wandb.Artifact('code', type='code') artifact.add_file('auto_encoder_body.py') @@ -377,14 +377,14 @@ def train(): # data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic)) valid_loss = {}#valid_step(data_valid) - if step % 1000 == 0: + if step % 2500 == 0: train_figs = plot_everything(batch, 'train') valid_figs = {}#plot_everything(data_valid, 'valid') wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step) else: wandb.log({ **train_loss, **valid_loss }, step=step) - if step % 2500 == 0: + if step % 10000 == 0: manager.save(checkpoint_number=step) @@ -410,8 +410,8 @@ if __name__ == '__main__': # 'uv_posed': (3, 512, 256,), 'uv_body_posed': (3, 512, 256,), 'mesh_body_posed': (14475, 3,), - 'mesh_posed': (14475, 3,), - 'mesh_unposed': (14475, 3,), + # 'mesh_posed': (14475, 3,), + # 'mesh_unposed': (14475, 3,), # 'uv_unposed': (3, 512, 256,) } signature_static = { #'uv_static': (3, 512, 256,), @@ -421,8 +421,8 @@ if __name__ == '__main__': # 'uv_shaped_body': (3, 512, 256,), } - dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic) - static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static) + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic, use_filled=False) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False) 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) @@ -432,6 +432,6 @@ if __name__ == '__main__': folder = 'models/ae_logs/' ckpt = tf.train.Checkpoint(model=model) - manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=5) + manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3) train() \ No newline at end of file diff --git a/auto_encoder_curv.py b/auto_encoder_curv.py index f4f3dc7..c9528bf 100644 --- a/auto_encoder_curv.py +++ b/auto_encoder_curv.py @@ -5,33 +5,25 @@ import wandb import random import numpy as np import tensorflow as tf -os.environ["CUDA_VISIBLE_DEVICES"] = "2" +os.environ["CUDA_VISIBLE_DEVICES"] = "0" from custom_utils.custom_utils import UV2Mesh, maskFacesBatch -from tensorboardX import SummaryWriter import plotly.graph_objects as go import plotly.express as ex -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 d_logistic, d_logistic_r1_reg, g_logistic_non_saturating, g_logistic_ns_pathreg, edge_loss, uv_mask_loss -# from losses import collision_penalty as collision_penalty +from tf_utils import check_tf_version, allow_memory_growth from losses import collision_loss as collision_penalty from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels -from stylegan2_ae.generator import Generator -from custom_utils.graph_utils import Reconstructor_V4_Decoder as GNNDecoder -from custom_utils.graph_utils import Reconstructor_v2 as SimpleReconstructor from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict from custom_utils.smpl import SMPL -name='ae_curvature_7' +name='ae_curvature_7_2' testing = False def init_wandb(): if not testing: - run = wandb.init(project='vae', name=name, notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments') + run = wandb.init(project='vae', name=name, notes='convnext + convnext, 16*8*32, ce loss. flattened latent dim of 128. trained on 100 garments. same as 7') artifact = wandb.Artifact('code', type='code') artifact.add_file('auto_encoder_curv.py') artifact.add_file('dataset.py') @@ -359,9 +351,11 @@ if __name__ == '__main__': dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=8, - signature=signature_dynamic) + signature=signature_dynamic, + use_filled=False) static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', - signature=signature_static) + signature=signature_static, + use_filled=False) dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=8, signature=signature_dynamic) diff --git a/convert_to_uv.py b/convert_to_uv.py new file mode 100644 index 0000000..f141bb7 --- /dev/null +++ b/convert_to_uv.py @@ -0,0 +1,654 @@ +import os +import sys +import tqdm +import math +import time +import random +import shutil +import collections +import argparse +import threading +import numpy as np +import scipy as sp +import trimesh as tm +import cmapy as cm +import scipy.io as sio +import scipy.interpolate as spi +import matplotlib.pyplot as plt +import scipy.stats as stat +import cv2 +import plotly.express as ex +import plotly.graph_objects as go + + +from custom_utils.supersmpl_io import readPC2 + + + +os.environ["CUDA_VISIBLE_DEVICES"] = "2" + + +def add_triangle_centers(): + pass + +def dilate(img): + uv_mask_2 = img != 0 + m = img.min() + # norm = img.max() - img.min() + img = (img - m) * uv_mask_2 #* (np.invert(uv_mask_2) * -1 + uv_mask_2) + size = 1 + shape = cv2.MORPH_CROSS # cv.MORPH_CROSS cv.MORPH_ELLIPSE MORPH_RECT + element = cv2.getStructuringElement(shape, (2 * size + 1, 2 * size + 1), (size, size)) + result = cv2.dilate(img, element, iterations=1) + result = (result + m) * uv_mask + return result + +def erode(img): + # m = img.min() + # norm = img.max() - img.min() + # img = (img - m) + size = 1 + shape = cv2.MORPH_CROSS # cv.MORPH_CROSS cv.MORPH_ELLIPSE MORPH_RECT + element = cv2.getStructuringElement(shape, (2 * size + 1, 2 * size + 1), (size, size)) + result = cv2.erode(img, element, iterations=2) + # result = result + m + return result + +def chunks(lst, n): + """Yield successive n-sized chunks from lst.""" + for i in range(0, len(lst), n): + yield lst[i:i + n] + + +class Result(): + def __init__(self) -> None: + self.result = dict() + +result_g = Result() + +class RBFThread(threading.Thread): + def __init__(self, uv_values, id, result): + threading.Thread.__init__(self) + self.uv_values = uv_values + self.id = id + self.result = result + + def run(self): + print('processing: ' + str(self.id)) + rbf = spi.RBFInterpolator(y=uv_map_discreet_unique, d=self.uv_values, neighbors=10) + res = rbf(rbf_dest) + self.result[self.id] = res + pass + +def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rbf', garment_uv_mask=None): + verts_uv = np.zeros((res * 2, res, 3)) + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = verts + uv_values[m2uv[:, 1]] += verts + uv_values = uv_values[uv_map_mask, :] # add garment mask + uv_values /= uv_counter[:, np.newaxis] + return_garment_uv_mask = False + + if mask is not None: + garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) + garment_mask[m2uv[:, 0]] = mask + garment_mask[m2uv[:, 1]] += mask + garment_mask = garment_mask[uv_map_mask] # add garment mask + else: + garment_mask = np.sum(uv_values, axis=-1) != 0 + + # test + verts_uv[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values + # verts_uv_save = verts_uv.copy() + + # garment_mask = np.sum(uv_values, axis=-1) != 0 + unique_uv_values = uv_values[garment_mask, :] + uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :] + rbf_dest_current = rbf_dest + + if garment_uv_mask is None and mask is not None: + garment_uv_mask = np.zeros((res * 2, res)) + garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255 + garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (4 * 2 + 1, 4 * 2 + 1), (4, 4)), iterations=4) + garment_uv_mask_border = garment_uv_mask.astype(bool) + garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool) + garment_uv_mask_border = ~garment_uv_mask_border & garment_uv_mask + garment_uv_mask = garment_uv_mask * uv_mask[:,:,0] + rbf_dest_current = rbf_dest * garment_uv_mask + return_garment_uv_mask = True + + + rbf_dest_current = np.nonzero(rbf_dest_current) + rbf_dest_current = np.array(rbf_dest_current).T + + # uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T + if method == 'rbf': + interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values, neighbors=neighbours, kernel='linear') + elif method == 'nearest': + interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values) + else: + raise NotImplementedError + + final = interpolator(rbf_dest_current) + + result = np.zeros((res * 2, res, 3)) + + # not sure, shouldn't be needed + # rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]] + + result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final + result[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values + if garment_uv_mask is not None: + result = result * garment_uv_mask[:,:,np.newaxis] + + # x = verts_uv #+ eroded + # uv_mask_2 = verts_uv != 0 + # y = result + # fig = plt.figure(figsize=(15, 15)) + # ax = fig.add_subplot(121) + # # im = ((x - x.min()) / (x.max() - x.min())) * uv_mask_2 + # im = ((x + 2) * 64 / 255) * uv_mask_2 + # ax.imshow(im) + # ax = fig.add_subplot(122) + # # im = ((y - y.min()) / (y.max() - y.min())) * uv_mask + # im = ((y + 2) * 64 / 255) * uv_mask # * (rbf_dest + garment_uv_mask)[:,:,np.newaxis] + # ax.imshow(im) + # plt.show() + + if return_garment_uv_mask: + return result, garment_uv_mask, garment_uv_mask_border + else: + return result + + + +def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None): + batch = verts.shape[0] + verts_uv = np.zeros((batch, res * 2, res, 3)) + uv_values = np.zeros((batch, uv_map.shape[0], 3)) + uv_values[:, m2uv[:, 0]] = verts + uv_values[:, m2uv[:, 1]] += verts + uv_values = uv_values[:, uv_map_mask, :] + uv_values /= uv_counter[np.newaxis, :, np.newaxis] + return_garment_uv_mask = False + # test + verts_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values + # verts_uv_save = verts_uv.copy() + + # unique_uv_values_mask = np.sum(uv_values[0, :, :], axis=-1) != 0 + # unique_uv_values = uv_values[:, unique_uv_values_mask, :] + # uv_map_discreet_unique_unique = uv_map_discreet_unique[unique_uv_values_mask, :] + + if mask is not None: + garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) + garment_mask[m2uv[:, 0]] = mask + garment_mask[m2uv[:, 1]] += mask + garment_mask = garment_mask[uv_map_mask] # add garment mask + else: + garment_mask = np.sum(uv_values[0,:,:], axis=-1) != 0 + + unique_uv_values = uv_values[:, garment_mask, :] + uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :] + rbf_dest_current = rbf_dest + + # add back mesh vertices that are not part of the garment + uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[np.invert(garment_mask), :] + rbf_dest_fill = np.zeros((res * 2, res), dtype=bool) + rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True + + if garment_uv_mask is None and mask is not None: + garment_uv_mask = np.zeros((res * 2, res)) + garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255 + garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4) + garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool) + return_garment_uv_mask = True + + if garment_uv_mask is not None: + rbf_dest_current = (rbf_dest + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border + + rbf_dest_current = np.nonzero(rbf_dest_current) + rbf_dest_current = np.array(rbf_dest_current).T + + # uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T + final = [] + for i in range(batch): + print('processing frame: ' + str(i)) + if method == 'rbf': + interpolator = spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[i,:,:], neighbors=neighbours, kernel='linear') + elif method == 'nearest': + interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:,:]) + else: + raise NotImplementedError + final.append(interpolator(rbf_dest_current)) + final = np.asarray(final) + + # def f(uv, d, id): + # print("processing: " + str(id)) + # res = spi.RBFInterpolator(y=uv_map_discreet_unique, d=uv, neighbors=10)(rbf_dest) + # d[id] = res + + + # with Manager() as manager: + # d = manager.dict() + # f_arg = [(uv_values[i, :, :], d, i) for i in range(8)] + # f_arg = zip(uv_values[:8, :, :], d, range(8)) + # with Pool(processes=4) as pool: + # pool.starmap(f, f_arg) + + # a = d.get() + + # processes = [] + # batch_ids = np.linspace(0, batch-1, batch).astype(np.int32).tolist() + # batch_ids_split = list(chunks(batch_ids, 2)) + # for ids in batch_ids_split: + # for i in range(len(ids)): + # p = RBFThread(uv_values[i, :, :], ids[i], result_g.result) + # p.start() + + # for p in processes: + # p.join() + + # final = np.zeros((batch, res, res, 3)) + # for k in sorted(batch_interpolated.keys()): + # final[k, :, :, :] = batch_interpolated[k] + + result = np.zeros((batch, res * 2, res, 3)) + result[:, rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final + # result[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values + result[:, uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1], :] = unique_uv_values + if garment_uv_mask is not None: + result = result * garment_uv_mask[np.newaxis,:,:,np.newaxis] + + if return_garment_uv_mask: + return result, garment_uv_mask + else: + return result + +def convert_to_uv_batch_label(res, m2uv, uv_map, values, neighbours=15, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=None): + batch = values.shape[0] + values = (values + 1) * mask + values_uv = np.zeros((batch, res * 2, res)) + missing = np.zeros((res * 2, res)) + uv_values = np.zeros((batch, uv_map.shape[0])) + uv_values[:, m2uv[:, 0]] = values + uv_values[:, m2uv[:, 1]] = values + uv_values = uv_values[:, uv_map_mask] + return_garment_uv_mask = False + # test + # verts_uv_save = verts_uv.copy() + + # unique_uv_values_mask = np.sum(uv_values[0, :, :], axis=-1) != 0 + # unique_uv_values = uv_values[:, unique_uv_values_mask, :] + # uv_map_discreet_unique_unique = uv_map_discreet_unique[unique_uv_values_mask, :] + + if mask is not None: + garment_mask = np.zeros((uv_map.shape[0]), dtype=bool) + garment_mask[m2uv[:, 0]] = mask + garment_mask[m2uv[:, 1]] += mask + garment_mask = garment_mask[uv_map_mask] # add garment mask + else: + raise Exception + + # uv_values = uv_values * garment_mask + + # test + values_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values + missing[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = (uv_values[0,:] == 0) + + unique_uv_values = uv_values[:, garment_mask] + uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :] + # rbf_dest_current = rbf_dest.astype(bool) + garment_uv_mask_border * uv_mask[:,:,0] + + if garment_uv_mask is None and mask is not None: + garment_uv_mask = np.zeros((res * 2, res)) + garment_uv_mask[uv_map_discreet_unique_unique[:, 0], uv_map_discreet_unique_unique[:, 1]] = 255 + garment_uv_mask = cv2.morphologyEx(garment_uv_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=4) + garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool) + return_garment_uv_mask = True + + rbf_dest_current = rbf_dest * garment_uv_mask + missing * garment_uv_mask + rbf_dest_current = np.nonzero(rbf_dest_current) + rbf_dest_current = np.array(rbf_dest_current).T + + # uv_verts_uniq = np.array(np.nonzero(np.sum(verts_uv, axis=-1))).T + final = [] + for i in range(batch): + print('processing frame: ' + str(i)) + interpolator = spi.NearestNDInterpolator(x=uv_map_discreet_unique_unique, y=unique_uv_values[i,:]) + final.append(interpolator(rbf_dest_current)) + final = np.asarray(final) + + result = np.zeros((batch, res * 2, res)) + result[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = uv_values + result[:, rbf_dest_current[:, 0], rbf_dest_current[:, 1]] = final + if garment_uv_mask is not None: + result = result * garment_uv_mask[np.newaxis,:,:] + + if return_garment_uv_mask: + return result, garment_uv_mask + else: + return result + + +def m2uv_array(m2uv): + result = np.zeros((14475, 2), dtype=np.int32) + for k, v in m2uv.items(): + if len(v) == 2: + result[k] = list(v) + elif len(v) == 1: + x = v.pop() + result[k] = [x, x] + else: + assert False + + return result + + +def get_uv2m(m2uv): + + pass + +def uv_to_mesh(uv, m2uv): + res = uv.shape[0] + mesh = np.zeros((14475, 3)) + uv_pixels = uv_map_discreet[m2uv] + # counter = np.ones((14475)) + # counter[m2uv[:, 0] == m2uv[:, 1]] = 2 + mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis] + return mesh + pass + + +def verify_reconstruction(orig, reconstructed): + error = np.sqrt(np.sum((orig - reconstructed) ** 2, axis=-1)) + colors = np.array([cm.color('jet', 1 - i) for i in error]) + error_sum = np.sum(error) + + + verts_uv = np.zeros((uv_res * 2, uv_res, 3)) + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = colors + uv_values[m2uv[:, 1]] += colors + uv_counter = np.zeros((uv_map.shape[0])) + uv_counter[m2uv[:, 0]] = 1 + uv_counter[m2uv[:, 1]] += 1 + uv_values /= uv_counter[:, np.newaxis] + verts_uv[uv_map_discreet[:, 0], uv_map_discreet[:, 1], :] = uv_values + fig = plt.figure(figsize=(15, 15)) + ax = fig.add_subplot(111) + ax.scatter(uv_map[:, 0], uv_map[:, 1], s=5, c=uv_values/255) + plt.show() + + + + fig = plt.figure(figsize=(15, 15)) + ax = fig.add_subplot(121, projection='3d') + x = orig + ax.set_xlim3d(-1, 1) + ax.set_ylim3d(-1, 1) + ax.set_zlim3d(-1, 1) + ax.set_xlabel("X") + ax.set_ylabel("Y") + ax.set_zlabel("Z") + ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5) + + ax = fig.add_subplot(122, projection='3d') + x = reconstructed#[error > 0.1, :] + ax.set_xlim3d(-1, 1) + ax.set_ylim3d(-1, 1) + ax.set_zlim3d(-1, 1) + ax.set_xlabel("X") + ax.set_ylabel("Y") + ax.set_zlabel("Z") + ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5, c=colors/255) #[error > 0.1, :] + + plt.show() + + pass + +# https://stackoverflow.com/questions/38187286/find-unique-pairs-in-list-of-pairs +def unique_pairs(a): + a.sort(axis=1) + b = np.ascontiguousarray(a).view( + np.dtype((np.void, a.dtype.itemsize * a.shape[1])) + ) + _, inv, ct = np.unique(b, return_inverse=True, return_counts=True) + return ct[inv] == 1 + + +def unique_uv_map_mask(uv_map_discreet): + mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool) + temp = np.zeros((uv_res * 2, uv_res)) + for p in uv_map_discreet: + temp[p[0], p[1]] += 1 + # temp = temp == 1 + for i, p in enumerate(uv_map_discreet): + if temp[p[0], p[1]] == 1: + mask[i] = True + elif temp[p[0], p[1]] == 2: + mask[i] = True + temp[p[0], p[1]] = 0 + return mask + + +def test(verts, faces, values): + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = values + uv_values[m2uv[:, 1]] += values + uv_counter = np.zeros((uv_map.shape[0])) + uv_counter[m2uv[:, 0]] = 1 + uv_counter[m2uv[:, 1]] += 1 + uv_values /= uv_counter[:, np.newaxis] + x = uv_values + verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1) + im = ((x - x.min()) / (x.max() - x.min())) * 255 + tm.Trimesh(verts, faces, vertex_colors=im.astype(np.int32)).show() + + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("-f", "--folder", help = "end of range") + parser.add_argument("-s", "--start", help = "start of range") + parser.add_argument("-e", "--end", help = "end of range") + parser.add_argument("-n", "--name", help = "end of range") + args = parser.parse_args() + + # smpl = SMPL('FeatureLines/superSMPL') + + # F = np.load("D:\\Work\\Code\\featurelines\\FeatureLines\\superSMPL\\smpl_faces.npy") + F = np.load("custom_utils/smpl_faces.npy") + # data = np.load("Blender/uv_data_1.npz", allow_pickle=True) + # uv_faces = data.f.Ft + # folder = 'Z:\\Work\\Huni\\Barcelona\\Dataset\\train' + folder = args.folder + uv_map = np.load('custom_utils/stretched_uv_map_3.npy') + m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item() + m2uv = m2uv_array(m2uv) + uv_mask = np.load('custom_utils/mask_256x512.npy') + curvature_label_type = '6_25' + + uv_res = 256 + uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis] + # uv_mask_dilate = cv2.erode(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_CROSS, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)[:, :, np.newaxis] + # uv_mask_border = uv_mask * np.invert(uv_mask_dilate) + uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis] + + uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32) + uv_map_mask = unique_uv_map_mask(uv_map_discreet) + + uv_counter = np.zeros((uv_map.shape[0])) + uv_counter[m2uv[:, 0]] = 1 + uv_counter[m2uv[:, 1]] += 1 + uv_counter = uv_counter[uv_map_mask] + + uv_map_discreet_unique = uv_map_discreet[uv_map_mask, :] + + rbf_dest = np.ones((uv_res * 2, uv_res)) + rbf_dest[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = 0 + rbf_dest[np.invert(uv_mask[:,:,0])] = 0 + # rbf_dest = np.nonzero(rbf_dest) + # rbf_dest = np.array(rbf_dest).T + # rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]] + # rbf_dest = rbf_dest[unique_pairs(rbf_dest), :] # not good, unique_pairs doesn't consider coord order + + ids = os.listdir(folder)[int(args.start):int(args.end)] + + 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] + + # if os.path.exists(os.path.join(folder, id, 'uv_unposed', garment + '_0.npz')): + # continue + + # uv_verts_static = np.load(os.path.join(folder, id, garment + '_offsets_static_uv.npy')) + # uv_verts_static = ((uv_verts_static - uv_verts_static.min()) / (uv_verts_static.max() - uv_verts_static.min())) + # plt.imshow(uv_verts_static) + # plt.show() + # pass + + + # offsets = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3] + offsets_static = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3] + body_static = sio.loadmat(os.path.join(garment_dir, 'info_body.mat'))['B'] + 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])]) + 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])] + body_no_orient = np.stack([body_no_orient[i]['mesh_posed_body_no_orient'] for i in range(gt.shape[0])]) + + + verts_static = (body_static + offsets_static) # * mask[:, np.newaxis] + verts_dynamic = gt + + # TODO: check if have to mask out non garment vertices or not? are they the same vertices erpeated many times? if remove, have to interpolate from almost nothing? + + # test(data.f.Vt, data.f.Ft, verts) + # test(uv_map, data.f.Ft, verts) + + start = time.perf_counter() + + + + + # R = smpl.run_rodrigues(np.zeros((1,72), dtype=np.float32)) + # body = smpl.run(R, np.zeros((1,11), dtype=np.float32)) + + # uv_body = convert_to_uv(uv_res, m2uv, uv_map, body[0][0], neighbours=3, mask=None, method='rbf', garment_uv_mask=None).astype(np.float32) + # _, faces = _maskFaces(np.zeros((14475,3)), F, mask) + # x = body[0][0] + # y = body_static + # faces = F + + # R = smpl.run_rodrigues(sio.loadmat(os.path.join(garment_dir, 'info.mat'))['poses'][:, 0][np.newaxis, :].astype(np.float32)) + # R = smpl.run_rodrigues(np.zeros((1,72), dtype=np.float32)) + # betas = np.concatenate((sio.loadmat(os.path.join(garment_dir, 'info.mat'))['shape'], sio.loadmat(os.path.join(garment_dir, 'info.mat'))['gender']), axis=-1).astype(np.float32) + # shaped_body = smpl.run(R, betas)[0][0] + # x = smpl.run_garment(R, betas, x[np.newaxis,:,:])[0] + # y = smpl.run_garment(R, betas, y[np.newaxis,:,:])[0] + + # fig = go.Figure(data=[ + # go.Mesh3d( + # # 8 vertices of a cube + # x=x[:,0], + # y=x[:,1], + # z=x[:,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=y[:,0], + # y=y[:,1], + # z=y[:,2], + # # i, j and k give the vertices of triangles + # i = faces[:,0], + # j = faces[:,1], + # k = faces[:,2], + # ), + # ]) + + # fig.update_layout(scene_aspectmode='manual', + # scene_aspectratio=dict(x=1, y=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() + + uv_verts_static, garment_uv_mask, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, verts_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None) + uv_verts_static = uv_verts_static.astype(np.float32) + # 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_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, 'body_static_uv.npy'), uv_body_static) + # 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]): + # np.savez_compressed(os.path.join(folder, id, 'body_posed_uv', garment + '_' + str(i) + '.npz'), uv_body=uv_body[i, :, :, :]) + + os.makedirs(os.path.join(folder, id, 'uv_posed_body_no_orient'), exist_ok=True) + 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, '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_no_shape'), exist_ok=True) + # for i in range(uv_unposed_no_shape.shape[0]): + # np.savez_compressed(os.path.join(folder, id, 'uv_unposed_no_shape', garment + '_' + str(i) + '.npz'), uv_unposed_no_shape=uv_unposed_no_shape[i, :, :, :]) + + + # os.makedirs(os.path.join(folder, id, 'uv_shaped_body'), exist_ok=True) + # np.savez_compressed(os.path.join(folder, id, 'uv_shaped_body.npz'), uv_shaped_body=uv_shaped_body) + + print(time.perf_counter() - start) + + # mesh = uv_to_mesh(uv_verts_static, m2uv) + # verify_reconstruction(verts_static * mask[:, np.newaxis], mesh * mask[:, np.newaxis]) + + # mesh = uv_to_mesh(uv_body[0,:,:], m2uv) + # verify_reconstruction(body[0,:,:] * mask[:, np.newaxis], mesh * mask[:, np.newaxis]) + + pass + diff --git a/custom_data/train_100_no_shape_avg_uv.npy b/custom_data/train_100_no_shape_avg_uv.npy new file mode 100644 index 0000000..23a1a42 Binary files /dev/null and b/custom_data/train_100_no_shape_avg_uv.npy differ diff --git a/custom_utils/fill_uv.py b/custom_utils/fill_uv.py new file mode 100644 index 0000000..fee3e73 --- /dev/null +++ b/custom_utils/fill_uv.py @@ -0,0 +1,103 @@ +import os +import tqdm +import scipy.interpolate as spi +import numpy as np +import plotly.graph_objects as go +import plotly.express as ex + + +def fill(img, mask): + coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0)) + img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0)) + img_data = img[img_idx[:,0], img_idx[:,1], :] + + interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data) + img_2 = interpolator(coords) + res = np.zeros((512,256,3), dtype=np.float32) + res[coords[:,0], coords[:,1], :] = img_2 + res[img_idx[:,0], img_idx[:,1], :] = img_data + + return res + + +def fill_label(img, mask): + coords = np.transpose(np.stack(np.nonzero((1 - mask))), (1,0)) + img_idx = np.transpose(np.stack(np.nonzero((mask))), (1,0)) + img_data = img[img_idx[:,0], img_idx[:,1]] + + interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data) + img_2 = interpolator(coords) + res = np.zeros((512,256), dtype=np.float32) + res[coords[:,0], coords[:,1]] = img_2 + res[img_idx[:,0], img_idx[:,1]] = img_data + + return res + + +def fill_uv_maps_posed_body(folder, type): + os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True) + mask = np.load('custom_utils/mask_256x512.npy') + + frames = os.listdir(os.path.join(folder, type)) + + for frame in tqdm.tqdm(frames): + id, _ = frame.split('_') + # mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask'] + img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient'] + filled = fill(img, mask) + np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_posed_body_no_orient_filled = filled) + + +def fill_uv_maps_curv(folder, type): + os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True) + mask = np.load('custom_utils/mask_256x512.npy') + + frames = os.listdir(os.path.join(folder, type)) + + for frame in tqdm.tqdm(frames): + id, _ = frame.split('_') + mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask'] + img = np.load(os.path.join(folder, type, frame)); img = img['curvature_uv'] + filled = fill_label(img, mask) + np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_curvature_filled = filled) + + + +def fill_uv_maps_unposed(folder, type): + os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True) + mask = np.load('custom_utils/mask_256x512.npy') + + frames = os.listdir(os.path.join(folder, type)) + + for frame in tqdm.tqdm(frames): + id, _ = frame.split('_') + mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask'] + img = np.load(os.path.join(folder, type, frame)); img = img['uv_unposed_no_shape'] + filled = fill(img, mask) + np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_unposed_no_shape_filled = filled) + + + +def fill_uv_maps_static(folder, type): + os.makedirs(os.path.join(folder, type + '_filled'), exist_ok=True) + mask = np.load('custom_utils/mask_256x512.npy') + + frames = os.listdir(os.path.join(folder, type)) + + for frame in tqdm.tqdm(frames): + id, _ = frame.split('.') + mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); mask = mask['uv_mask'] + img = np.load(os.path.join(folder, type, frame)); img = img['uv_static'] + filled = fill(img, mask) + np.savez_compressed(os.path.join(folder, type + '_filled', frame), uv_static_filled = filled) + + + +if __name__ == '__main__': + # fill_uv_maps_static(folder='/home/workspace/local-data/train_100', type='uv_static') + # fill_uv_maps_unposed(folder='/home/workspace/local-data/train_100', type='uv_unposed_no_shape') + # fill_uv_maps_curv(folder='/home/workspace/local-data/train_100', type='uv_curvature') + fill_uv_maps_posed_body(folder='/home/workspace/local-data/train_100', type='uv_posed_body_no_orient') + + + print() \ No newline at end of file diff --git a/custom_utils/supersmpl_io.py b/custom_utils/supersmpl_io.py new file mode 100644 index 0000000..179a6bb --- /dev/null +++ b/custom_utils/supersmpl_io.py @@ -0,0 +1,138 @@ +import numpy as np +from struct import pack, unpack + + +""" +Reads PC2 files, and proposed format PC16 files +Inputs: +- file: path to .pc2/.pc16 file +- float16: False for PC2 files, True for PC16 +Output: +- data: dictionary with .pc2/.pc16 file data +NOTE: 16-bit floats lose precision with high values (positive or negative), + we do not recommend using this format for data outside range [-2, 2] +""" +def readPC2(file, float16=False): + assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format' + data = {} + bytes = 2 if float16 else 4 + dtype = np.float16 if float16 else np.float32 + with open(file, 'rb') as f: + # Header + data['sign'] = f.read(12) + # data['version'] = int.from_bytes(f.read(4), 'little') + data['version'] = unpack('= 0 and isinstance(frame,int), 'Frame must be a positive integer' + bytes = 2 if float16 else 4 + dtype = np.float16 if float16 else np.float32 + with open(file,'rb') as f: + # Num points + f.seek(16) + # nPoints = int.from_bytes(f.read(4), 'little') + nPoints = unpack(' nSamples: + print("Frame index outside size") + print("\tN. frame: " + str(frame)) + print("\tN. samples: " + str(nSamples)) + return + # Read frame + size = nPoints * 3 * bytes + f.seek(size * frame, 1) # offset from current '1' + T = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32) + return T.reshape(nPoints, 3) + +""" +Writes PC2 and PC16 files +Inputs: +- file: path to file (overwrites if exists) +- V: 3D animation data as a three dimensional array (N. Frames x N. Vertices x 3) +- float16: False for writing as PC2 file, True for PC16 +This function assumes 'startFrame' to be 0 and 'sampleRate' to be 1 +NOTE: 16-bit floats lose precision with high values (positive or negative), + we do not recommend using this format for data outside range [-2, 2] +""" +def writePC2(file, V, float16=False): + assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format' + if float16: V = V.astype(np.float16) + else: V = V.astype(np.float32) + with open(file, 'wb') as f: + # Create the header + headerFormat='<12siiffi' + headerStr = pack(headerFormat, b'POINTCACHE2\0', + 1, V.shape[1], 0, 1, V.shape[0]) + f.write(headerStr) + # Write vertices + f.write(V.tobytes()) + +""" +Reads OBJ files +Only handles vertices, faces and UV maps +Input: +- file: path to .obj file +Outputs: +- V: 3D vertices +- F: 3D faces +- Vt: UV vertices +- Ft: UV faces +Correspondence between mesh and UV map is implicit in F to Ft correspondences +If no UV map data in .obj file, it shall return Vt=None and Ft=None +""" +def readOBJ(file): + V, Vt, F, Ft = [], [], [], [] + with open(file, 'r') as f: + T = f.readlines() + for t in T: + # 3D vertex + if t.startswith('v '): + v = [float(n) for n in t.replace('v ','').split(' ')] + V += [v] + # UV vertex + elif t.startswith('vt '): + v = [float(n) for n in t.replace('vt ','').split(' ')] + Vt += [v] + # Face + elif t.startswith('f '): + idx = [n.split('/') for n in t.replace('f ','').split(' ')] + f = [int(n[0]) - 1 for n in idx] + F += [f] + # UV face + if '/' in t: + f = [int(n[1]) - 1 for n in idx] + Ft += [f] + V = np.array(V, np.float32) + Vt = np.array(Vt, np.float32) + if Ft: assert len(F) == len(Ft), 'Inconsistent .obj file, mesh and UV map do not have the same number of faces' + else: Vt, Ft = None, None + return V, F, Vt, Ft \ No newline at end of file diff --git a/cvae.py b/cvae.py index f0f125e..9dfeb02 100644 --- a/cvae.py +++ b/cvae.py @@ -1,5 +1,5 @@ import os -os.environ["CUDA_VISIBLE_DEVICES"] = "3" +os.environ["CUDA_VISIBLE_DEVICES"] = "0" import tensorflow as tf gpus = tf.config.experimental.list_physical_devices('GPU') @@ -35,10 +35,10 @@ 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.cvae import CVAE +from models.cvae_model import CVAE -name='cvae_15' +name='cvae_16' def init_wandb(): run = wandb.init(project='vae', name=name, notes='learn offsets. smaller flattened latent vectors. with 25-85k pretraining. smaller batchsize. fixed balancing of reconst and kl loss for baseline. \ @@ -47,9 +47,11 @@ offsets over average unposed garment. remove global orientation (in posed body i retrained body encoder with deoriented bodies \ use mean l2 loss. small batch. l2 * 1000, balancing to 1 from 0.01\ fixed gender issue and retrained body encoder\ -remove shape from unposed garment') +remove shape from unposed garment\ +training on 100 garments with offsets over the average of these') artifact = wandb.Artifact('code', type='code') artifact.add_file('cvae.py') + artifact.add_file('models/cvae_model.py') artifact.add_file('dataset.py') artifact.add_file('losses.py') artifact.add_file('convnext/models/convnext_tf.py') @@ -121,7 +123,7 @@ def kl_loss(mean, logvar): return kl_loss -# @tf.function +@tf.function def train_step(batch): with tf.GradientTape() as tape: @@ -385,7 +387,7 @@ def sample(folder): print() if __name__ == '__main__': - # init_wandb() + init_wandb() allow_memory_growth() parser = argparse.ArgumentParser(description='') parser.add_argument('--name', type=str) @@ -419,10 +421,10 @@ if __name__ == '__main__': static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static) - model = CVAE(avg_garment_path='custom_data/jumpsuit_20_avg_uv.npy', - latent_dim=128, static_model='checkpoints/ae_static_3', - body_model='checpoints/ae_body_6', - curvature_model='checpoints/ae_curvature_6') + model = CVAE(avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy', + latent_dim=128, static_model='checkpoints/ae_static_4', + body_model='checpoints/ae_body_7', + curvature_model='checpoints/ae_curvature_7') # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False) diff --git a/cvae_curv.py b/cvae_curv.py index 76fe407..146b885 100644 --- a/cvae_curv.py +++ b/cvae_curv.py @@ -35,7 +35,7 @@ from losses import collision_loss as collision_penalty from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict from custom_utils.smpl import SMPL -from models.cvae_curv import CVAE +from models.cvae_curv_model import CVAE name='cvae_curv_2_2_2' @@ -48,6 +48,7 @@ fixed deoriented bodies and need retraining (ae_body_6 with 50k iteration), adde higher starting kl coeff (0.1)') artifact = wandb.Artifact('code', type='code') artifact.add_file('cvae_curv.py') + artifact.add_file('models/cvae_curv_model.py') artifact.add_file('dataset.py') artifact.add_file('losses.py') artifact.add_file('convnext/models/convnext_tf.py') diff --git a/dataset.py b/dataset.py index 4b9107d..2332cfd 100644 --- a/dataset.py +++ b/dataset.py @@ -54,18 +54,22 @@ class DatasetSimple(tf.data.Dataset): class StaticDataset(): - def __init__(self, folder, signature) -> None: + def __init__(self, folder, signature, use_filled=True) -> None: self.store = dict() self.signature = signature self.folder = folder self.possible_keys = ['uv_static', 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_shaped_body'] self.keys = [value for value in signature.keys() if value in self.possible_keys] + filled = '' + if use_filled: + filled = '_filled' + def read(self, id): data = dict() if 'uv_static' in self.signature.keys(): - uv_static = np.load(os.path.join(self.folder, 'uv_static', id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1]); data['uv_static']=uv_static + uv_static = np.load(os.path.join(self.folder, 'uv_static' + filled, id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1]); data['uv_static' + filled]=uv_static if 'mesh_static' in self.signature.keys(): mesh_static = np.load(os.path.join(self.folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static']; data['mesh_static'] = mesh_static if 'mesh_mask' in self.signature.keys(): @@ -127,9 +131,12 @@ class DatasetFactory(): return tf_signature - def get_dataset(self, folder, batch_size, signature, shuffle=True): + def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=True): # static_data = StaticDataset(folder, signature) dataset_idx = DatasetIdx(folder, shuffle) + filled = '' + if use_filled: + filled = '_filled' def read_item(frame): frame_split = tf.strings.split(frame, sep='_') @@ -161,7 +168,7 @@ class DatasetFactory(): # mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient'] # item['mesh_posed_body_no_orient'] = mesh_posed_body_no_orient if 'uv_body_posed' in signature.keys(): - uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; uv_posed_body_no_orient = np.transpose(uv_posed_body_no_orient, [2, 0, 1]) + uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient' + filled, id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled]; uv_posed_body_no_orient = np.transpose(uv_posed_body_no_orient, [2, 0, 1]) item['uv_body_posed'] = uv_posed_body_no_orient if 'mesh_body_posed' in signature.keys(): mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient'] @@ -179,10 +186,14 @@ class DatasetFactory(): mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape', id + '_' + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape'] item['mesh_unposed'] = mesh_unposed_no_shape if 'uv_unposed' in signature.keys(): - uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape', id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']; uv_unposed_no_shape = np.transpose(uv_unposed_no_shape, [2, 0, 1]) + uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + filled, id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]; uv_unposed_no_shape = np.transpose(uv_unposed_no_shape, [2, 0, 1]) item['uv_unposed'] = uv_unposed_no_shape if 'uv_curvature' in signature.keys(): - uv_curvature = np.load(os.path.join(folder, 'uv_curvature', id + '_' + frame_nr + '.npz')); uv_curvature = uv_curvature['curvature_uv'] + uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + filled, id + '_' + frame_nr + '.npz')); + if filled != '': + uv_curvature = uv_curvature['uv_curvature' + filled] + else: + uv_curvature = uv_curvature['curvature_uv'] item['uv_curvature'] = uv_curvature.astype(np.int32) if 'mesh_curvature' in signature.keys(): mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + '_' + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature'] @@ -239,58 +250,60 @@ if __name__ == '__main__': signature_dynamic = { 'id': (), # 'pose_shape': (72 + 11,), # 'uv_posed': (3, 512, 256,), - # 'uv_body_posed': (3, 512, 256,), + 'uv_body_posed': (3, 512, 256,), # 'mesh_body_posed': (14475, 3,), # 'mesh_posed': (14475, 3,), - 'mesh_unposed': (14475, 3,), + # 'mesh_unposed': (14475, 3,), 'uv_unposed': (3, 512, 256,), # 'uv_curvature': (512, 256), # 'mesh_curvature': (14475) } signature_static = { #'uv_static': (3, 512, 256,), # 'mesh_static': (14475, 3,), - 'mesh_mask': (14475,), + # 'mesh_mask': (14475,), 'uv_mask': (512, 256,), # 'uv_shaped_body': (3, 512, 256,), } - dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', - batch_size=1, - signature=signature_dynamic) - static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', - signature=signature_static) + dataset = factory.get_dataset(folder='/home/workspace/local-data/train_100', + batch_size=300, + signature=signature_dynamic, + use_filled=False) + static_dataset = StaticDataset(folder='/home/workspace/local-data/train_100', + signature=signature_static, + use_filled=False) - # sum_uv = np.zeros((3, 512, 256), dtype=np.float32) - # sum_mask = np.zeros((512, 256), dtype=np.float32) - # for i, batch in enumerate(dataset): - # print(i) - # batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) - # sum_uv += batch['uv_unposed'].numpy().sum(axis=0) - # sum_mask += batch['uv_mask'].numpy().sum(axis=0) + sum_uv = np.zeros((3, 512, 256), dtype=np.float32) + sum_mask = np.zeros((512, 256), dtype=np.float32) + for i, batch in enumerate(dataset): + print(i) + batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) + sum_uv += batch['uv_unposed'].numpy().sum(axis=0) + sum_mask += batch['uv_mask'].numpy().sum(axis=0) - # avg_jumpsuit = sum_uv / sum_mask[np.newaxis, :, :] - # sum_mask[sum_mask == 0] = 1 + sum_mask[sum_mask == 0] = 1 + avg_jumpsuit = sum_uv / sum_mask[np.newaxis, :, :] - # x = np.transpose(avg_jumpsuit, [1,2,0]) - # x = uv2mesh(avg_jumpsuit[tf.newaxis,:,:,:])[0] - # go.Figure(data=[ - # go.Mesh3d( - # # 8 vertices of a cube - # x=x[:,0], - # y=x[:,1], - # z=x[:,2], - # # i, j and k give the vertices of triangles - # i = F[:,0], - # j = F[:,1], - # k = F[:,2] - # )]).show() + x = np.transpose(avg_jumpsuit, [1,2,0]) + x = uv2mesh(avg_jumpsuit[tf.newaxis,:,:,:])[0] + go.Figure(data=[ + go.Mesh3d( + # 8 vertices of a cube + x=x[:,0], + y=x[:,1], + z=x[:,2], + # i, j and k give the vertices of triangles + i = F[:,0], + j = F[:,1], + k = F[:,2] + )]).show() - # np.save('custom_data/jumpsuit_20_avg_uv.npy', np.transpose(avg_jumpsuit, [1,2,0])) + np.save('custom_data/train_100_no_shape_avg_uv.npy', np.transpose(avg_jumpsuit, [1,2,0])) - # print() + print() diff --git a/downscaling_test.py b/downscaling_test.py index 163a480..b795755 100644 --- a/downscaling_test.py +++ b/downscaling_test.py @@ -4,6 +4,7 @@ import argparse import wandb import random import tqdm +import cv2 import scipy as sp import scipy.interpolate as spi import numpy as np @@ -36,14 +37,18 @@ uv2mesh = UV2Mesh(res=256) def downscale_1(img): results = [img] for i in range(5): - results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 * (i + 1)), method=tf.image.ResizeMethod.LANCZOS5)) + # print(str(np.array(img.shape[:2]) / (2 ** (i + 1))) + ' ' + str((2 * (i + 1)))) + results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.LANCZOS5)) return results def downscale_2(img): results = [img] for i in range(5): # print(np.array(img.shape[:2]) / (2 ** (i + 1))) - results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR, antialias=True, preserve_aspect_ratio=True)) + # results.append(tf.image.resize(img, size=np.array(img.shape[:2]) / (2 ** (i + 1)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR, antialias=True, preserve_aspect_ratio=True)) + size = (np.array(img.shape[:2]) / (2 ** (i + 1))).astype(np.int32) + size = (size[1], size[0]) + results.append(cv2.resize(img, size, interpolation=cv2.INTER_NEAREST)) return results def upscale_1(imgs): @@ -136,11 +141,13 @@ if __name__ == '__main__': # 'uv_shaped_body': (3, 512, 256,), } - dataset_train = factory.get_dataset(folder='/home/workspace/local-data/stylegan_new/train', + dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=1, - signature=signature_dynamic) - static_dataset_train = StaticDataset(folder='/home/workspace/local-data/stylegan_new/train', - signature=signature_static) + signature=signature_dynamic, + use_filled=True) + static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', + signature=signature_static, + use_filled=True) step = 0 @@ -161,9 +168,18 @@ if __name__ == '__main__': img_2 = res + uv_mask = batch['uv_mask'].numpy()[0] + uv_mask = downscale_2(uv_mask) + # for i in range(0,3): + # uv_mask[i] = cv2.dilate(uv_mask[i].astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)#.astype(bool)[:, :, np.newaxis] + # for i in range(3,6): + # uv_mask[i] = cv2.dilate(uv_mask[i].astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3), (2, 2)), iterations=1)#.astype(bool)[:, :, np.newaxis] + + down_1 = downscale_1(img) down_2 = downscale_1(img_2) up_1 = upscale_1(down_1) + # down_2 = [down_2[i] * uv_mask[i][:,:,np.newaxis] for i in range(len(down_2))] up_2 = upscale_1(down_2) mask = batch['mesh_mask'].numpy().astype(bool)[0] @@ -172,6 +188,7 @@ if __name__ == '__main__': F = np.load("custom_utils/smpl_faces.npy") _, faces = _maskFaces(np.ones((14475, 3)), F, mask) + # mesh_downscaled = mesh_2[-1].numpy()[mask, :] # uv_optimized = optimize(down_2[-1], tf.transpose(batch['uv_posed'][0], [1,2,0])) # uv_optimized = tf.image.resize(uv_optimized, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5) @@ -213,7 +230,7 @@ if __name__ == '__main__': fig = make_subplots(rows=2, cols=6, start_cell="top-left") - for i, d2 in enumerate(down_2): + for i, d2 in enumerate(uv_mask): fig.add_trace(ex.imshow((d2 + 1) * 128).data[0], row=1, col=i+1) for i, u2 in enumerate(up_2): diff --git a/merge_data.py b/merge_data.py new file mode 100644 index 0000000..519b69d --- /dev/null +++ b/merge_data.py @@ -0,0 +1,268 @@ +import os +import sys +import tqdm +import math +import time +import random +import shutil +import argparse +import numpy as np +import scipy as sp +import scipy.io as sio +import cv2 + +from custom_utils.supersmpl_io import readPC2 + + + +def chunks(lst, n): + """Yield successive n-sized chunks from lst.""" + for i in range(0, len(lst), n): + yield lst[i:i + n] + + +def m2uv_array(m2uv): + result = np.zeros((14475, 2), dtype=np.int32) + for k, v in m2uv.items(): + if len(v) == 2: + result[k] = list(v) + elif len(v) == 1: + x = v.pop() + result[k] = [x, x] + else: + assert False + + return result + + +def uv_to_mesh(uv, m2uv): + res = uv.shape[0] + mesh = np.zeros((14475, 3)) + uv_pixels = uv_map_discreet[m2uv] + # counter = np.ones((14475)) + # counter[m2uv[:, 0] == m2uv[:, 1]] = 2 + mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis] + return mesh + pass + + +def unique_uv_map_mask(uv_map_discreet): + mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool) + temp = np.zeros((uv_res * 2, uv_res)) + for p in uv_map_discreet: + temp[p[0], p[1]] += 1 + # temp = temp == 1 + for i, p in enumerate(uv_map_discreet): + if temp[p[0], p[1]] == 1: + mask[i] = True + elif temp[p[0], p[1]] == 2: + mask[i] = True + temp[p[0], p[1]] = 0 + return mask + + +def verify_reconstruction(orig, reconstructed): + error = np.sqrt(np.sum((orig - reconstructed) ** 2, axis=-1)) + colors = np.array([cm.color('jet', 1 - i) for i in error]) + error_sum = np.sum(error) + + + verts_uv = np.zeros((uv_res * 2, uv_res, 3)) + uv_values = np.zeros((uv_map.shape[0], 3)) + uv_values[m2uv[:, 0]] = colors + uv_values[m2uv[:, 1]] += colors + uv_counter = np.zeros((uv_map.shape[0])) + uv_counter[m2uv[:, 0]] = 1 + uv_counter[m2uv[:, 1]] += 1 + uv_values /= uv_counter[:, np.newaxis] + verts_uv[uv_map_discreet[:, 0], uv_map_discreet[:, 1], :] = uv_values + fig = plt.figure(figsize=(15, 15)) + ax = fig.add_subplot(111) + ax.scatter(uv_map[:, 0], uv_map[:, 1], s=5, c=uv_values/255) + plt.show() + + fig = plt.figure(figsize=(15, 15)) + ax = fig.add_subplot(121, projection='3d') + x = orig + ax.set_xlim3d(-1, 1) + ax.set_ylim3d(-1, 1) + ax.set_zlim3d(-1, 1) + ax.set_xlabel("X") + ax.set_ylabel("Y") + ax.set_zlabel("Z") + ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5) + + ax = fig.add_subplot(122, projection='3d') + x = reconstructed#[error > 0.1, :] + ax.set_xlim3d(-1, 1) + ax.set_ylim3d(-1, 1) + ax.set_zlim3d(-1, 1) + ax.set_xlabel("X") + ax.set_ylabel("Y") + ax.set_zlabel("Z") + ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5, c=colors/255) #[error > 0.1, :] + + plt.show() + + return error_sum + + pass + + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("-fi", "--folder_in", help = "end of range") + parser.add_argument("-fo", "--folder_out", help = "end of range") + parser.add_argument("-s", "--start", help = "start of range") + parser.add_argument("-e", "--end", help = "end of range") + args = parser.parse_args() + folder = args.folder_in + output_folder = args.folder_out + + + # # F = np.load("D:\\Work\\Code\\featurelines\\FeatureLines\\superSMPL\\smpl_faces.npy") + # F = np.load("FeatureLines/superSMPL/smpl_faces.npy") + # # data = np.load("Blender/uv_data_1.npz", allow_pickle=True) + # # uv_faces = data.f.Ft + # # folder = 'D:\\Work\\Dataset\\validation_2' + # uv_map = np.load('MeshUtils/stretched_uv_map_3.npy') + # m2uv = np.load('MeshUtils/m2uv.npy', allow_pickle=True).item() + # m2uv = m2uv_array(m2uv) + # uv_mask = np.load('MeshUtils/mask_256x512.npy') + # curvature_label_type = '6_25' + + # uv_res = 256 + # uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis] + # # uv_mask_dilate = cv2.erode(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_CROSS, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)[:, :, np.newaxis] + # # uv_mask_border = uv_mask * np.invert(uv_mask_dilate) + # uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis] + + # uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32) + # uv_map_mask = unique_uv_map_mask(uv_map_discreet) + + # uv_counter = np.zeros((uv_map.shape[0])) + # uv_counter[m2uv[:, 0]] = 1 + # uv_counter[m2uv[:, 1]] += 1 + # uv_counter = uv_counter[uv_map_mask] + + # uv_map_discreet_unique = uv_map_discreet[uv_map_mask, :] + + # rbf_dest = np.ones((uv_res * 2, uv_res)) + # rbf_dest[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = 0 + # rbf_dest[np.invert(uv_mask[:,:,0])] = 0 + + ids = os.listdir(folder)[int(args.start):int(args.end)] + + os.makedirs(os.path.join(output_folder), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'mesh_mask'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'mesh_static'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'uv_static'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'mesh_posed'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'mesh_body_posed'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'uv_posed'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'uv_body_posed'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'pose_shape'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'uv_mask'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'mesh_unposed'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'uv_unposed'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'mesh_unposed_no_shape'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'uv_unposed_no_shape'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'uv_shaped_body'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'uv_curvature'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'mesh_curvature'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'uv_posed_body_no_orient'), exist_ok=True) + os.makedirs(os.path.join(output_folder, 'mesh_posed_body_no_orient'), exist_ok=True) + + + 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] + + # if garment != 'Jumpsuit': + # continue + + start = time.perf_counter() + offsets_static = sio.loadmat(os.path.join(folder, id, garment + '_offsets_static.mat'))['label'][:,:3] + body_static = sio.loadmat(os.path.join(garment_dir, 'info_body.mat'))['B'] + mask = np.squeeze(sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].astype(bool)) + verts_static = (body_static + offsets_static) * mask[:, np.newaxis] + assert verts_static.min() > -1 and verts_static.max() < 1 + + + # np.savez_compressed(os.path.join(output_folder, 'mesh_static', id + '.npz'), mesh_static=verts_static) + + # uv_verts_static = np.load(os.path.join(folder, id, garment + '_offsets_static_uv.npy')) + # np.savez_compressed(os.path.join(output_folder, 'uv_static', id + '.npz'), uv_static=uv_verts_static) + + # garment_folder_gt = os.path.join(folder, id, 'gt_uv') + # garment_folder_body = os.path.join(folder, id, 'body_posed_uv') + # garment_folder_body_mesh = os.path.join(folder, id, 'mesh_body_posed') + + # shutil.copyfile(os.path.join(folder, id, garment + '_mask.npy'), os.path.join(output_folder, 'mesh_mask', id + '.npy')) + + # garment_posed = readPC2(os.path.join(garment_dir, garment + '_registered.pc16'), True)['V'] + # for i in range(garment_posed.shape[0]): + # np.savez_compressed(os.path.join(output_folder, 'mesh_posed', id + '_' + str(i) + '.npz'), mesh_posed=garment_posed[i,:,:] * mask[:, np.newaxis]) + + body_posed = readPC2(os.path.join(garment_dir, 'Body_posed.pc16'), True)['V'] + # for i in range(body_posed.shape[0]): + # np.save(os.path.join(output_folder, 'mesh_body_posed', id + '_' + str(i) + '.npy'), body_posed[i,:,:]) + + # for i in range(body_posed.shape[0]): + # shutil.copyfile(os.path.join(garment_folder_gt, garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_posed', id + '_' + str(i) + '.npz')) + + # for i in range(body_posed.shape[0]): + # shutil.copyfile(os.path.join(garment_folder_body, garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_body_posed', id + '_' + str(i) + '.npz')) + + # info = sio.loadmat(os.path.join(garment_dir, 'info.mat')) + # pose, shape, gender = info['poses'], info['shape'], info['gender'] + # shape = np.squeeze(np.concatenate((shape, gender), axis=1)) + # for i in range(body_posed.shape[0]): + # np.save(os.path.join(output_folder, 'pose_shape', id + '_' + str(i) + '.npy'), np.concatenate((pose[:, i], shape), axis=-1)) + + + # uv_mask = np.load(os.path.join(folder, id, 'mask_uv.npy')) + # np.savez_compressed(os.path.join(output_folder, 'uv_mask', id + '.npz'), uv_mask=uv_mask) + + # for i in range(body_posed.shape[0]): + # shutil.copyfile(os.path.join(garment_dir, 'mesh_unposed', id + '_' + str(i) + '.npz'), os.path.join(output_folder, 'mesh_unposed', id + '_' + str(i) + '.npz')) + + # for i in range(body_posed.shape[0]): + # shutil.copyfile(os.path.join(garment_dir, 'uv_unposed', garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_unposed', id + '_' + str(i) + '.npz')) + + # for i in range(body_posed.shape[0]): + # shutil.copyfile(os.path.join(garment_dir, 'mesh_unposed_no_shape', id + '_' + str(i) + '.npz'), os.path.join(output_folder, 'mesh_unposed_no_shape', id + '_' + str(i) + '.npz')) + + # for i in range(body_posed.shape[0]): + # shutil.copyfile(os.path.join(garment_dir, 'uv_unposed_no_shape', garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_unposed_no_shape', id + '_' + str(i) + '.npz')) + + # for i in range(body_posed.shape[0]): + # shutil.copyfile(os.path.join(garment_dir, 'curvature_uv', garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_curvature', id + '_' + str(i) + '.npz')) + + # mesh_curvature = np.load(os.path.join(garment_dir, 'curvature_label_6_25.npy')) + # for i in range(body_posed.shape[0]): + # np.savez_compressed(os.path.join(output_folder, 'mesh_curvature', id + '_' + str(i) + '.npz'), mesh_curvature=mesh_curvature[i]) + + # shutil.copyfile(os.path.join(folder, id, 'uv_shaped_body.npz'), os.path.join(output_folder, 'uv_shaped_body', id + '.npz')) + + for i in range(body_posed.shape[0]): + shutil.copyfile(os.path.join(garment_dir, 'uv_posed_body_no_orient', garment + '_' + str(i) + '.npz'), os.path.join(output_folder, 'uv_posed_body_no_orient', id + '_' + str(i) + '.npz')) + + for i in range(body_posed.shape[0]): + shutil.copyfile(os.path.join(garment_dir, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz'), os.path.join(output_folder, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz')) + + + pass + diff --git a/models/cvae_curv.py b/models/cvae_curv_model.py similarity index 100% rename from models/cvae_curv.py rename to models/cvae_curv_model.py diff --git a/models/cvae.py b/models/cvae_model.py similarity index 100% rename from models/cvae.py rename to models/cvae_model.py diff --git a/models/cvae_pyramid_model.py b/models/cvae_pyramid_model.py new file mode 100644 index 0000000..f4d60c7 --- /dev/null +++ b/models/cvae_pyramid_model.py @@ -0,0 +1,191 @@ +import os +import time +import random +import tensorflow as tf +import numpy as np +import scipy as sp +# import plotly.graph_objects as go +# import plotly.express as ex + +from custom_utils.custom_utils import load_model +from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block +from auto_encoder_body import AutoEncoder as BodyAE +from auto_encoder_curv import AutoEncoder as CurvatureAE +from auto_encoder_static import AutoEncoder as StaticAE + + +class Bottleneck(tf.keras.Model): + + def __init__(self, dim): + super(Bottleneck, self).__init__() + + self.bottleneck = tf.keras.Sequential( + [Block(dim=dim, drop_path=0, + layer_scale_init_value=1e-6, prefix=f'bottleneck{j}') for j in range(5)] + ) + + def call(self, inputs, training=None): + x = inputs + + x = self.bottleneck(x, training=training) + + return x + + +class Encoder(tf.keras.Model): + + def __init__(self, latent_dim): + super(Encoder, self).__init__() + + self.encoder = ConvNeXt(depths=[3, 3, 9, 3], dims=[32, 64, 128, 256], include_top=False) + self.flatten = tf.keras.layers.Flatten() + self.mean = tf.keras.layers.Dense(latent_dim) + self.log_var = tf.keras.layers.Dense(latent_dim) + + def call(self, inputs, training=None): + x = inputs + + x = self.encoder(x) + x = self.flatten(x) + mean = self.mean(x) + log_var = self.log_var(x) + + return mean, log_var + + +class Decoder(tf.keras.Model): + + def __init__(self): + super(Decoder, self).__init__() + + self.dense = tf.keras.layers.Dense(16 * 8 * 80) + self.reshape = tf.keras.layers.Reshape((16, 8, 80)) + self.decoder = ConvNeXtReverse(depths=list(reversed([3, 3, 9, 3])), dims=list(reversed([256, 128, 64, 16]))) + + + def call(self, inputs, training=None): + x = inputs + + x = self.dense(x) + x = self.reshape(x) + x = self.decoder(x) + + return x + + +class CVAE(tf.keras.Model): + + def __init__(self, avg_garment_path='custom_data/jumpsuit_20_avg_uv.npy', + latent_dim=128, static_model='checkpoints/ae_static_3', + body_model='checpoints/ae_body_6', + curvature_model='checpoints/ae_curvature_6'): + super(CVAE, self).__init__() + self.latent_dim = latent_dim + + model : StaticAE = StaticAE() + load_model(model, static_model) + self.static_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) + self.static_enc.trainable = False + + model : BodyAE = BodyAE() + load_model(model, body_model) + self.body_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) + self.body_enc.trainable = False + + model : CurvatureAE = CurvatureAE() + load_model(model, curvature_model) + self.curv_enc = tf.keras.Sequential([model.encoder, model.bottleneck]) + self.curv_enc.trainable = False + + self.encoder = Encoder(latent_dim) + self.decoder = Decoder() + + self.avg_garment = np.load(avg_garment_path) + + self.cond_encoder = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'), + tf.keras.layers.Dense(256, activation='relu'), + tf.keras.layers.Dense(128, activation='relu'), + tf.keras.layers.Dense(latent_dim),]) + + # self.shuffle = tf.keras.layers.Lambda(lambda x: tf.stack([tf.gather(x[i], tf.random.shuffle(tf.range(x[i].shape[0]))) for i in range(x.shape[0])])) + + # @tf.function + def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, nr_of_samples=4, fix_noise=False, zero_latent=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) + + if zero_latent: + eps = tf.zeros_like(eps) + + 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 reparameterize(self, mean, logvar): + eps = tf.random.normal(shape=mean.shape) + 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 call(self, inputs, training=None, zero_latent=False): + uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_shaped_body = inputs + uv_unposed, uv_template, uv_body_posed = tf.transpose(uv_unposed, [0,2,3,1]), 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) + + x = tf.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], axis=-1) + + mean, logvar = self.encoder(x) + z = self.reparameterize(mean, logvar) + z_cond = tf.concat([z_static, z_body, z_curv], axis=-1) + z_cond = self.cond_encoder(z_cond) + + if zero_latent: + z = tf.zeros_like(z) + + z = tf.concat([z, z_cond], axis=-1) + + # z = self.shuffle(z) + + # z = tf.concat([z, z_curv], axis=-1) + x = self.decoder(z) + + + # x = tf.transpose(x, [0,3,1,2]) + # x = x * uv_mask[:,tf.newaxis,:,:] + # x = tf.image.resize(x, (512, 256), tf.image.ResizeMethod.LANCZOS5) + # x = x + tf.transpose(uv_shaped_body, [0,2,3,1]) + x = x + self.avg_garment[tf.newaxis,:,:,:] + x = x * uv_mask[:,:,:,tf.newaxis] + + return x, mean, logvar, z \ No newline at end of file diff --git a/out_0.png b/out_0.png deleted file mode 100644 index 388f42e..0000000 Binary files a/out_0.png and /dev/null differ diff --git a/preprocessing.py b/preprocessing.py new file mode 100644 index 0000000..46c127a --- /dev/null +++ b/preprocessing.py @@ -0,0 +1,111 @@ +import os +import sys +import shutil +import tqdm +import random +import math +import numpy as np +import scipy as sp +import scipy.io as sio + + +from custom_utils.smpl import SMPL +import plotly.graph_objects as go + + +def create_no_orient_body(folder): + ids = os.listdir(folder) + F = np.load("custom_utils/smpl_faces.npy") + + smpl = SMPL() + for id in tqdm.tqdm(ids): + if not os.path.isdir(os.path.join(folder, id)): + continue + + info = sio.loadmat(os.path.join(folder, id, 'info.mat')) + + pose, shape, gender = info['poses'], info['shape'], info['gender'] + + # if gender == 1: + # continue + + pose = np.transpose(pose, [1, 0]).astype(np.float32) + pose[:, :3] = 0 + beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32) + R = smpl.run_rodrigues(pose) + T, _, _ = smpl.run(R, np.tile(beta[np.newaxis, :], [pose.shape[0], 1]), True) + T = T.numpy() + + # print(gender[0,0]) + + # meshes = [] + # x = T[0, :, :] + # meshes.append(go.Mesh3d( + # # 8 vertices of a cube + # x=x[:,0], + # y=x[:,2], + # z=x[:,1], + # # i, j and k give the vertices of triangles + # i = F[:,0], + # j = F[:,1], + # k = F[:,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() + + # frame_nr = str(0) + # folder_2 = '/home/workspace/local-data/train_100/mesh_posed_body_no_orient' + # mesh_posed_body_no_orient = np.load(os.path.join(folder_2, id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient'] + + # meshes = [] + # x = mesh_posed_body_no_orient + # meshes.append(go.Mesh3d( + # # 8 vertices of a cube + # x=x[:,0], + # y=x[:,2], + # z=x[:,1], + # # i, j and k give the vertices of triangles + # i = F[:,0], + # j = F[:,1], + # k = F[:,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() + + os.makedirs(os.path.join(folder, id, 'mesh_posed_body_no_orient'), exist_ok=True) + for i in range(pose.shape[0]): + np.savez_compressed(os.path.join(folder, id, 'mesh_posed_body_no_orient', id + '_' + str(i) + '.npz'), mesh_posed_body_no_orient = T[i, :, :]) + + + + +if __name__ == '__main__': + os.environ["CUDA_VISIBLE_DEVICES"] = "1" + + src_folder = '/home/workspace/data/dataset/train_skirt' + # src_folder = r'D:\Work\Dataset\validation_2' + dest_folder = '/home/workspace/data/dataset/validation' + + + create_no_orient_body('/home/workspace/data/dataset/train_100') diff --git a/sample.py b/sample.py index 9e38d11..ccac5d6 100644 --- a/sample.py +++ b/sample.py @@ -43,7 +43,7 @@ from auto_encoder_body import AutoEncoder as BodyAE from auto_encoder_curv import AutoEncoder as CurvatureAE from auto_encoder_static import AutoEncoder as StaticAE from cvae_curv import CVAE as CVAECurv -from models.cvae import CVAE as CVAEMain +from models.cvae_model import CVAE as CVAEMain name='cvae_12_6' def init_wandb(): diff --git a/temp.jpg b/temp.jpg deleted file mode 100644 index 25487f1..0000000 Binary files a/temp.jpg and /dev/null differ diff --git a/verify_genders.py b/verify_genders.py new file mode 100644 index 0000000..133f097 --- /dev/null +++ b/verify_genders.py @@ -0,0 +1,57 @@ +import os +import tqdm +import numpy as np +import plotly.graph_objects as go + + +if __name__ == '__main__': + folder = '/home/workspace/local-data/train_100/mesh_posed_body_no_orient' + mask = np.load('custom_utils/mask_256x512.npy') + F = np.load("custom_utils/smpl_faces.npy") + + frames = os.listdir(folder) + visited = set() + meshes = [] + count = 0 + + for frame in tqdm.tqdm(frames): + id, frame_nr = frame.split('_') + frame_nr = frame_nr.split('.')[0] + if id in visited: + continue + visited.add(id) + print(id) + mesh_posed_body_no_orient = np.load(os.path.join(folder, id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient'] + x = mesh_posed_body_no_orient + # x[:, 2] += 1 * (i + 1) + x[:, 2] += 0.5 * count + meshes.append(go.Mesh3d( + # 8 vertices of a cube + x=x[:,0], + y=x[:,2], + z=x[:,1], + # i, j and k give the vertices of triangles + i = F[:,0], + j = F[:,1], + k = F[:,2] + )) + count += 1 + + if count > 40: + break + + + + fig = go.Figure(data=meshes) + fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=count/2-1, z=1)) + fig.update_layout( + scene = dict( + xaxis = dict(nticks=4, range=[-1,1],), + yaxis = dict(nticks=4, range=[-1,count/2+1],), + zaxis = dict(nticks=4, range=[-1,1],),) + ) + + fig.show() + + print() +