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 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) -> 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] 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 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_mask' in self.signature.keys(): uv_mask = np.load(os.path.join(self.folder, 'uv_mask', id + '.npz')); uv_mask=uv_mask['uv_mask']; uv_mask = uv_mask.astype(np.float32); data['uv_mask'] = uv_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']; uv_shaped_body = np.transpose(uv_shaped_body, [2, 0, 1]); data['uv_shaped_body'] = uv_shaped_body return data def __getitem__(self, id) -> dict: # id = tf.strings.unicode_decode(id, 'UTF-8') # id = tf.strings.as_string(id) # id = id.numpy() if id not in self.store: self.store[id] = self.read(id) 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): uv_static_folder = os.path.join(folder, 'pose_shape') frames = os.listdir(uv_static_folder) frames = [frame.split('.')[0] for frame in frames] # .split('_')[0] 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)) else: tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32)) return tf_signature def get_dataset(self, folder, batch_size, signature): # static_data = StaticDataset(folder, signature) dataset_idx = DatasetIdx(folder) 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')) 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']; uv_posed = np.transpose(uv_posed, [2, 0, 1]) 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' 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 # 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() # 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, ) # ) os.environ["CUDA_VISIBLE_DEVICES"] = "3" if __name__ == '__main__': from custom_utils.custom_utils import batch_to_dict factory = DatasetFactory() signature_dynamic = { 'id': (), # 'pose_shape': (72 + 11,), # 'uv_posed': (3, 512, 256,), # 'uv_body_posed': (3, 512, 256,), 'mesh_body_posed': (14475, 3,), 'mesh_posed': (14475, 3,), # 'mesh_unposed': (14475, 3,), # 'uv_unposed': (3, 512, 256,), } signature_static = { #'uv_static': (3, 512, 256,), # '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/stylegan_new/train', batch_size=1024, signature=signature_dynamic) static_dataset = StaticDataset(folder='/home/workspace/data/dataset/stylegan_new/train', signature=signature_static) @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))) / tf.reduce_sum(mask) all_error = [] for batch in dataset: batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], batch['mesh_body_posed']) all_error.append(err.numpy()) # 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 from custom_utils.custom_utils import UV2Mesh import matplotlib.pyplot as plt def plot_3d(x, y=None): fig = plt.figure(figsize=(15, 15)) ax = fig.add_subplot(111, projection='3d') 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.grid(False) ax.scatter(x[:, 0], x[:, 1], x[:, 2], s=5) if y is not None: ax.scatter(y[:, 0], y[:, 1], y[:, 2], s=5) plt.savefig('temp.jpg') uv2mesh = UV2Mesh() 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