import os import sys import time import cv2 import random import numpy as np import tensorflow as tf from custom_utils.custom_utils import nest_conditions, batch_to_dict # class DatasetSimple(tf.data.Dataset): # def _generator(folder): # # Opening the file # folder = folder.decode("utf-8") # uv_folder = os.path.join(folder, 'uv_posed') # pose_hape_folder = os.path.join(folder, 'poses_shapes') # frames = os.listdir(uv_folder) # random.shuffle(frames) # for frame in frames: # id, frame_nr = frame.split('.')[0].split('_') # uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']; uv_posed = np.transpose(uv_posed, [2, 0, 1]) # uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body']; uv_body_posed = np.transpose(uv_body_posed, [2, 0, 1]) # pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + frame_nr + '.npy')) # uv_static = np.load(os.path.join(folder, 'uv_static', id + '.npz')); uv_static = uv_static['uv_static']; uv_static = np.transpose(uv_static, [2, 0, 1]) # mesh_static = np.load(os.path.join(folder, 'mesh_static', id + '.npz')); mesh_static = mesh_static['mesh_static'] # mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy')) # mesh_mask = np.load(os.path.join(folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32) # uv_mask = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32) # mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'] # mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed'] # uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1]) # # uv_posed, (pose, shape, uv_static), mesh_static, mesh_body_posed, mask # yield uv_posed, (pose_shape[:72], pose_shape[72:], uv_static, uv_mask, uv_body_posed), mesh_static, mesh_body_posed, mesh_mask, mesh_posed, mesh_unposed, uv_unposed # def __new__(cls, folder, res=256): # return tf.data.Dataset.from_generator( # cls._generator, # output_signature = (tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), # (tf.TensorSpec(shape = (72), dtype = tf.float32), tf.TensorSpec(shape = (11), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (res * 2, res,), dtype = tf.float32), tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32) ), # tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), # tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), # tf.TensorSpec(shape = (14475,), dtype = tf.float32), # tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), # tf.TensorSpec(shape = (14475, 3,), dtype = tf.float32), # tf.TensorSpec(shape = (3, res * 2, res,), dtype = tf.float32), ), # args=(folder, ) # ) class StaticDataset(): def __init__(self, folder, signature, use_filled=False, resolutions=[512], in_memory=True, is_cape=False) -> None: self.store = dict() self.signature = signature self.folder = folder self.in_memory = in_memory assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_64', 'uv_static_128', 'uv_static_256', 'uv_mask', 'uv_mask_32', 'uv_mask_64', 'uv_mask_128', 'uv_mask_256', 'mesh_static', 'mesh_mask', 'uv_shaped_body'] self.keys = [value for value in signature.keys() if value in self.possible_keys] if 'uv_static_128' in signature.keys() and 'uv_static_128' not in self.keys: self.keys.append('uv_static_128') self.resolutions = resolutions self.filled = '' self.is_cape = is_cape if use_filled: self.filled = '_filled' if is_cape: self.uv_masks = dict() uv_mask = np.load('custom_data/garment_uv_mask_full.npy').astype(np.int32) self.uv_masks[512] = uv_mask.astype(np.float32) for res in [32, 64, 128, 256]: down = tf.image.resize(uv_mask[:,:,tf.newaxis], size=(res, int(res/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR) down = np.squeeze(down.numpy().astype(np.float32)) self.uv_masks[res] = down def read(self, id): data = dict() for res in self.resolutions: res_string = '' if res != 512: res_string = '_' + str(res) if 'uv_static' + res_string in self.signature.keys(): uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static if 'uv_mask' + res_string in self.signature.keys(): if self.is_cape: data['uv_mask' + res_string] = self.uv_masks[res] else: uv_mask = np.load(os.path.join(self.folder, 'uv_mask' + res_string, id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask' + res_string] = uv_mask 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.astype(np.float32) if 'mesh_mask' in self.signature.keys(): if self.is_cape: data['mesh_mask'] = np.ones((3638), np.float32) else: mesh_mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy')); mesh_mask = mesh_mask.astype(np.float32); data['mesh_mask'] = mesh_mask if 'uv_shaped_body' in self.signature.keys(): uv_shaped_body = np.load(os.path.join(self.folder, 'uv_shaped_body', id + '.npz')); uv_shaped_body = uv_shaped_body['uv_shaped_body']; data['uv_shaped_body'] = uv_shaped_body return data def __getitem__(self, id) -> dict: if id not in self.store: item = self.read(id) if not self.in_memory: return item else: self.store[id] = item return self.store[id] def expand_batch(self, batch): for key in self.keys: batch[key] = tf.zeros(((0,) + self.signature[key])) for i in range(batch['id'].shape[0]): item = self.__getitem__(batch['id'][i].numpy().decode('UTF-8')) for key in self.keys: batch[key] = tf.concat((batch[key], tf.expand_dims(item[key], axis=0)), axis=0) return batch class DatasetIdx(tf.data.Dataset): def __new__(cls, folder, shuffle=True, is_cape=False, testing=False, frames_file=None): if frames_file is None: uv_static_folder = os.path.join(folder, 'pose_shape') frames = sorted(os.listdir(uv_static_folder)) if is_cape: frames = ['.'.join(frame.split('.')[:-1]) for frame in frames] # .split('_')[0] else: frames = [frame.split('.')[0] for frame in frames] # .split('_')[0] if testing: frames = sorted(os.listdir(os.path.join(folder, 'mesh_static') )) frames = [frame.split('.')[0] + '_4' for frame in frames] if frames_file is not None: frames = [] with open(os.path.join(folder, frames_file), 'r') as f: lines = f.readlines() lines = [line.strip() for line in lines] frames = lines if shuffle: random.shuffle(frames) return tf.data.Dataset.from_tensor_slices(frames) # class DatasetFactory(tf.data.Dataset): class DatasetFactory(): def __init__(self): super(DatasetFactory, self).__init__() def _build_tf_signature(self, signature): tf_signature = [] for key in signature.keys(): if key == 'id': tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string)) elif key == 'frame_nr': tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string)) elif key == 'mesh_curvature': tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32)) else: tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32)) return tf_signature def get_dataset(self, folder, batch_size, signature, shuffle=True, resolutions=[512], in_memory=False, do_repeat=True, is_cape=False, use_one_hot=False, use_no_orient=True, testing=False, frames_file=None, curv_neighbourhood=25): assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape, testing=testing, frames_file=frames_file) separator = '_' if is_cape: separator = '.' one_hot = '' one_hot_2 = '' if use_one_hot: one_hot = '_one_hot' one_hot_2 = '_one_hot_2' uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient_2' if is_cape: uv_posed_body_no_orient_folder_name = 'uv_posed_body_no_orient' if use_no_orient == False: assert resolutions == [512], 'Using global orientation is only possible with non-pyramid model!' def read_item(frame): frame_split = tf.strings.split(frame, sep=separator) id, frame_nr = frame_split[0].numpy().decode("utf-8"), frame_split[1].numpy().decode("utf-8") # id, frame_nr = frame_split[0].numpy(), frame_split[1].numpy() # static_info = static_data[id] item = dict() if 'id' in signature.keys(): item['id'] = id if 'frame_nr' in signature.keys(): item['frame_nr'] = frame_nr if 'pose_shape' in signature.keys(): pose_shape = np.load(os.path.join(folder, 'pose_shape', id + separator + frame_nr + '.npy')) if use_no_orient: pose_shape[:3] = 0 item['pose_shape'] = pose_shape if 'mesh_body_posed' in signature.keys(): if use_no_orient: mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + separator + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient'] item['mesh_body_posed'] = mesh_posed_body_no_orient else: mesh_body_posed = np.load(os.path.join(folder, 'mesh_body_posed', id + '_' + frame_nr + '.npy')) item['mesh_body_posed'] = mesh_body_posed if 'mesh_posed' in signature.keys(): if use_no_orient: mesh_posed = np.load(os.path.join(folder, 'mesh_posed_no_orient', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed_no_orient'].astype(np.float32) item['mesh_posed'] = mesh_posed else: mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + separator + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'].astype(np.float32) item['mesh_posed'] = mesh_posed # if 'mesh_unposed' in signature.keys(): # mesh_unposed = np.load(os.path.join(folder, 'mesh_unposed', id + '_' + frame_nr + '.npz')); mesh_unposed = mesh_unposed['mesh_unposed'] # item['mesh_unposed'] = mesh_unposed # if 'uv_unposed' in signature.keys(): # uv_unposed = np.load(os.path.join(folder, 'uv_unposed', id + '_' + frame_nr + '.npz')); uv_unposed = uv_unposed['uv_unposed']; uv_unposed = np.transpose(uv_unposed, [2, 0, 1]) # item['uv_unposed'] = uv_unposed if 'mesh_unposed' in signature.keys(): # print('one_hot') mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape' + one_hot, id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape'] item['mesh_unposed'] = mesh_unposed_no_shape if 'mesh_curvature' in signature.keys(): mesh_curvature = np.load(os.path.join(folder, f'mesh_curvature_{curv_neighbourhood}_2' if curv_neighbourhood != 25 else 'mesh_curvature', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature'] item['mesh_curvature'] = mesh_curvature.astype(np.int32) for res in resolutions: res_string = '' if res != 512: res_string = '_' + str(res) if 'uv_curvature' + res_string in signature.keys(): # print(os.path.join(folder, (f'uv_curvature_{curv_neighbourhood}_2' if curv_neighbourhood != 25 else 'uv_curvature')+ res_string, id + separator + frame_nr + '.npz')) uv_curvature = np.load(os.path.join(folder, (f'uv_curvature_{curv_neighbourhood}_2' if curv_neighbourhood != 25 else 'uv_curvature')+ res_string, id + separator + frame_nr + '.npz')); if res_string != '': uv_curvature = uv_curvature['uv_curvature'] else: uv_curvature = uv_curvature['curvature_uv'] uv_curvature = tf.one_hot(uv_curvature, 7) item['uv_curvature' + res_string] = uv_curvature if 'uv_unposed' + res_string in signature.keys(): uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + one_hot_2 + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape'] item['uv_unposed' + res_string] = uv_unposed_no_shape if 'uv_body_posed' + res_string in signature.keys(): if use_no_orient: uv_posed_body_no_orient = np.load(os.path.join(folder, uv_posed_body_no_orient_folder_name + res_string, id + separator + frame_nr + '.npz')) uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient'] item['uv_body_posed' + res_string] = uv_posed_body_no_orient else: uv_body_posed = np.load(os.path.join(folder, 'uv_body_posed', id + '_' + frame_nr + '.npz')); uv_body_posed = uv_body_posed['uv_body'] item['uv_body_posed' + res_string] = uv_body_posed if 'uv_posed' + res_string in signature.keys(): if use_no_orient: uv_posed = np.load(os.path.join(folder, 'uv_posed_no_orient' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient'] item['uv_posed' + res_string] = uv_posed else: uv_posed = np.load(os.path.join(folder, 'uv_posed' + res_string, id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt'] item['uv_posed' + res_string] = uv_posed # print('here') # item = { **static_info, **item } as_list = [] for key in signature.keys(): if key in item: as_list.append(item[key]) return as_list # dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item, # inp=[i], # Tout=self._build_tf_signature(signature) # ), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat() if in_memory: dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item, inp=[i], Tout=self._build_tf_signature(signature) ), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True) else: dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item, inp=[i], Tout=self._build_tf_signature(signature) ), num_parallel_calls=tf.data.AUTOTUNE).batch(batch_size=batch_size, drop_remainder=True) if do_repeat: dataset = dataset.repeat() # dataset = dataset_idx.map(read_item, num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat() return dataset # def __new__(cls, folder): # return tf.data.Dataset.from_generator( # cls._generator, # # output_signature = (tf.TensorSpec(shape = (), dtype = tf.string), ), # output_types=(tf.string), # args=(folder, ) # ) if __name__ == '__main__': os.environ["CUDA_VISIBLE_DEVICES"] = "6" from custom_utils.custom_utils import batch_to_dict, UV2Mesh, _maskFaces import plotly.graph_objects as go import plotly.express as ex # import cmapy as cm uv2mesh = UV2Mesh(is_skirt=True) F = np.load("custom_utils/smpl_faces.npy") # dataset = DatasetIdx('/home/workspace/data/dataset/stylegan_new/train') # for d in dataset: # print(d) factory = DatasetFactory() signature_dynamic = { 'id': (), 'frame_nr': (), # 'pose_shape': (72 + 11,), # 'uv_posed': (3, 512, 256,), # 'uv_body_posed': (512, 256, 3), # 'uv_body_posed_32': (32, 16, 3), # 'mesh_body_posed': (14475, 3,), # 'mesh_posed': (14475, 3,), # 'mesh_unposed': (14475, 3,), 'uv_unposed': (512, 256, 3), # 'uv_unposed_32': (3, 32, 16,), 'uv_curvature': (512, 256, 7), # 'mesh_curvature': (14475) } signature_static = { #'uv_static': (3, 512, 256,), # 'uv_static_32': (3, 32, 16,), # 'mesh_static': (14475, 3,), # 'mesh_mask': (14475,), 'uv_mask': (512, 256,), # 'uv_shaped_body': (3, 512, 256,), } dataset = factory.get_dataset(folder='/home/workspace/data/dataset/local-cache/train_skirt', batch_size=1, signature=signature_dynamic, resolutions=[512], do_repeat=False, use_one_hot=False, is_cape=False, frames_file='frames_skirt_10p.txt', curv_neighbourhood=25) static_dataset = StaticDataset(folder='/home/workspace/data/dataset/local-cache/train_skirt', signature=signature_static, resolutions=[512], is_cape=False) # for i, batch in enumerate(dataset): # print(i) sum_uv = np.zeros((512, 256, 3), 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)) ex.imshow(batch['uv_curvature'][0].numpy().argmax(-1) / 6.0).show() sum_uv += batch['uv_unposed'].numpy().sum(axis=0) sum_mask += batch['uv_mask'].numpy().sum(axis=0) sum_mask[sum_mask == 0] = 1 avg_garment = sum_uv / sum_mask[:, :, np.newaxis] # x = np.transpose(avg_garment, [1,2,0]) x = uv2mesh(avg_garment[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/train_dress_30k_no_shape_avg_uv.npy', avg_garment) print() # pass # uv_res = 256 # 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 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 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 # 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 = 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] # for i, batch in enumerate(dataset): # print(i) # batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) # uv_unposed = np.squeeze(batch['uv_unposed'].numpy()) # mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy()) # mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool) # values = mesh_unposed # 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 # x = np.transpose(uv_unposed, [1,2,0]) # x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :] # _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask) # 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] # )]).show() # print() # # @tf.function # def euclid_error(mask, gt, pred): # # return tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1)), axis=-1) / tf.reduce_sum(mask, axis=-1) # return tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1)) # all_error = [] # for batch in dataset: # batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) # mesh_reconstruct = uv2mesh(batch['uv_posed']) # err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], uv2mesh(batch['uv_posed'])) # i = 296 # err_normed = (err[i].numpy() - err[i].numpy().min()) / (err[i].numpy().max() - err[i].numpy().min()) # colors = np.array([cm.color('jet', 1 - i) for i in err_normed]) # _, faces = _maskFaces(batch['mesh_posed'][i].numpy(), F, batch['mesh_mask'][i].numpy().astype(bool)) # x = batch['mesh_posed'][i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :] # y = mesh_reconstruct[i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :] # colors = colors[batch['mesh_mask'][i].numpy().astype(bool), :] # 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.show() # fig = ex.scatter_3d(x=x[:,0], y=x[:,1], z=x[:,2]) # fig.update_traces(marker={'color': colors, 'size': 5}) # fig.show() # all_error.append(err.numpy().sum()) # # b = nest_conditions(static_dataset.expand_batch(batch_to_dict(d, signature_dynamic))) # print(len(all_error)) # dataset = DatasetSimple('/home/workspace/data/dataset/stylegan_new/train').batch(2) # from losses import collision_loss as collision_penalty # from losses import collision_penalty as collision_penalty_2 # for d in dataset: # d # v = uv2mesh(d[0]) # b = d[3] # m = d[4] # b2 = tf.concat((b[1][tf.newaxis,:,:], b[0][tf.newaxis,:,:]), axis=0) # c = collision_penalty(v, b, m) # c2 = collision_penalty(v, b2, m) # c_ = collision_penalty_2(v, b, m) # c2_ = collision_penalty_2(v, b2, m) # print() # pass