update

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_posed': (14475, 3,),
-                          'mesh_unposed': (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)
+    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)
+    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()

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
-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 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)

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
+

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()

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('<i', f.read(4))[0]
+		# Num points
+		# data['nPoints'] = int.from_bytes(f.read(4), 'little')
+		data['nPoints'] = unpack('<i', f.read(4))[0]
+		# Start frame
+		data['startFrame'] = unpack('f', f.read(4))
+		# Sample rate
+		data['sampleRate'] = unpack('f', f.read(4))
+		# Number of samples
+		# data['nSamples'] = int.from_bytes(f.read(4), 'little')
+		data['nSamples'] = unpack('<i', f.read(4))[0]
+		# Animation data
+		size = data['nPoints']*data['nSamples']*3*bytes
+		data['V'] = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
+		data['V'] = data['V'].reshape(data['nSamples'], data['nPoints'], 3)
+		
+	return data
+	
+"""
+Reads an specific frame of PC2/PC16 files
+Inputs:
+- file: path to .pc2/.pc16 file
+- frame: number of the frame to read
+- float16: False for PC2 files, True for PC16
+Output:
+- T: mesh vertex data at specified frame
+"""
+def readPC2Frame(file, frame, float16=False):
+	assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
+	assert frame >= 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('<i', f.read(4))[0]
+		# Number of samples
+		f.seek(28)
+		# nSamples = int.from_bytes(f.read(4), 'little')
+		nSamples = unpack('<i', f.read(4))[0]
+		if frame > 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

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',
+    model = CVAE(avg_garment_path='custom_data/train_100_no_shape_avg_uv.npy',
-                 latent_dim=128, static_model='checkpoints/ae_static_3',
+                 latent_dim=128, static_model='checkpoints/ae_static_4',
-                 body_model='checpoints/ae_body_6',
+                 body_model='checpoints/ae_body_7',
-                 curvature_model='checpoints/ae_curvature_6')
+                 curvature_model='checpoints/ae_curvature_7')
 
 
     # lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(0.0001, decay_steps=30000, decay_rate=0.1, staircase=False)

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')

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',
+    dataset = factory.get_dataset(folder='/home/workspace/local-data/train_100',
-                                  batch_size=1,
+                                  batch_size=300,
-                                  signature=signature_dynamic)
+                                  signature=signature_dynamic,
-    static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
+                                  use_filled=False)
-                                   signature=signature_static)
+    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_uv = np.zeros((3, 512, 256), dtype=np.float32)
-    # sum_mask = np.zeros((512, 256), dtype=np.float32)
+    sum_mask = np.zeros((512, 256), dtype=np.float32)
-    # for i, batch in enumerate(dataset):
+    for i, batch in enumerate(dataset):
-    #     print(i)
+        print(i)
-    #     batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
+        batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
-    #     sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
+        sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
-    #     sum_mask += batch['uv_mask'].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 = np.transpose(avg_jumpsuit, [1,2,0])
-    # x = uv2mesh(avg_jumpsuit[tf.newaxis,:,:,:])[0]
+    x = uv2mesh(avg_jumpsuit[tf.newaxis,:,:,:])[0]
-    # go.Figure(data=[
+    go.Figure(data=[
-    #         go.Mesh3d(
+            go.Mesh3d(
-    #             # 8 vertices of a cube
+                # 8 vertices of a cube
-    #             x=x[:,0],
+                x=x[:,0],
-    #             y=x[:,1],
+                y=x[:,1],
-    #             z=x[:,2],
+                z=x[:,2],
-    #             # i, j and k give the vertices of triangles
+                # i, j and k give the vertices of triangles
-    #             i = F[:,0],
+                i = F[:,0],
-    #             j = F[:,1],
+                j = F[:,1],
-    #             k = F[:,2]
+                k = F[:,2]
-    #         )]).show()
+            )]).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()
 
 
 

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)
+                                  signature=signature_dynamic,
-    static_dataset_train = StaticDataset(folder='/home/workspace/local-data/stylegan_new/train',
+                                  use_filled=True)
-                                   signature=signature_static)
+    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):

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
+

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

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')

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():

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()
+