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) -> 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 = '' if use_filled: self.filled = '_filled' 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(): 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 if 'mesh_mask' in self.signature.keys(): 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): uv_static_folder = os.path.join(folder, 'pose_shape') frames = sorted(os.listdir(uv_static_folder)) frames = [frame.split('.')[0] for frame in frames] # .split('_')[0] 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 == '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, use_filled=False, resolutions=[512], in_memory=False): # assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1) assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) # 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='_') 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 'pose_shape' in signature.keys(): pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + frame_nr + '.npy')) ############################## pose_shape[:3] = 0 ############################## item['pose_shape'] = pose_shape if 'uv_posed' in signature.keys(): uv_posed = np.load(os.path.join(folder, 'uv_posed', id + '_' + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt'] item['uv_posed'] = uv_posed # if 'uv_body_posed' in signature.keys(): # 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]) # item['uv_body_posed'] = uv_body_posed # if 'mesh_body_posed' in signature.keys(): # 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_body_no_orient' 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'] # item['mesh_posed_body_no_orient'] = mesh_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'] item['mesh_body_posed'] = mesh_posed_body_no_orient if 'mesh_posed' in signature.keys(): mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + '_' + frame_nr + '.npz')); mesh_posed = mesh_posed['mesh_posed'] 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(): 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 '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'] 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(): uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + filled + res_string, id + '_' + frame_nr + '.npz')); if filled != '' or res_string != '': uv_curvature = uv_curvature['uv_curvature' + filled] 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' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled] item['uv_unposed' + res_string] = uv_unposed_no_shape if 'uv_body_posed' + res_string in signature.keys(): uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled] item['uv_body_posed' + res_string] = uv_posed_body_no_orient # 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).repeat() 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).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_skirt.npy") # dataset = DatasetIdx('/home/workspace/data/dataset/stylegan_new/train') # for d in dataset: # print(d) factory = DatasetFactory() signature_dynamic = { 'id': (), # 'pose_shape': (72 + 11,), # 'uv_posed': (3, 512, 256,), #'uv_body_posed': (3, 512, 256,), # 'uv_body_posed_32': (3, 32, 16,), # '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/local-data/train_skirt_25_2', batch_size=300, signature=signature_dynamic, use_filled=False, resolutions=[512]) static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt_25_2', signature=signature_static, use_filled=False, resolutions=[512]) 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)) 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_jumpsuit = sum_uv / sum_mask[:, :, np.newaxis] x = np.transpose(avg_jumpsuit, [1,2,0]) x = uv2mesh(avg_jumpsuit[tf.newaxis,:,:,:])[0] go.Figure(data=[ go.Mesh3d( # 8 vertices of a cube x=x[:,0], y=x[:,1], z=x[:,2], # i, j and k give the vertices of triangles i = F[:,0], j = F[:,1], k = F[:,2] )]).show() np.save('custom_data/train_skirt_25_2_no_shape_avg_uv.npy', np.transpose(avg_jumpsuit, [1,2,0])) print() 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