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 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' + 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): 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 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, do_repeat=True, is_cape=False): assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) # static_data = StaticDataset(folder, signature) dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape) separator = '.' folder = folder + '_packed' 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() data = np.load(os.path.join(folder, id + '.' + frame_nr + '.npz'), allow_pickle=True)['arr_0'].item() # print(os.path.join(folder, id + '.' + frame_nr + '.npz')) if 'id' in signature.keys(): item['id'] = id if 'pose_shape' in signature.keys(): pose_shape = data['pose_shape'] ############################## pose_shape[:3] = 0 ############################## item['pose_shape'] = pose_shape if 'uv_posed' in signature.keys(): item['uv_posed'] = data['uv_posed'] if 'mesh_body_posed' in signature.keys(): item['mesh_body_posed'] = data['mesh_posed_body_no_orient'] if 'mesh_posed' in signature.keys(): item['mesh_posed'] = data['mesh_posed'] if 'mesh_unposed' in signature.keys(): item['mesh_unposed'] = data['mesh_unposed_no_shape'] if 'mesh_curvature' in signature.keys(): item['mesh_curvature'] = data['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(): item['uv_curvature' + res_string] = tf.one_hot(data['uv_curvature' + res_string], 7) if 'uv_unposed' + res_string in signature.keys(): item['uv_unposed' + res_string] = data[ 'uv_unposed_no_shape' + res_string] if 'uv_body_posed' + res_string in signature.keys(): item['uv_body_posed' + res_string] = data['uv_posed_body_no_orient' + res_string] as_list = [] for key in signature.keys(): if key in item: as_list.append(item[key]) return as_list 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() return dataset 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': (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': (32, 16, 3), # '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/cape_processed/train/female', batch_size=1, signature=signature_dynamic, use_filled=False, resolutions=[512, 32], do_repeat=False, is_cape=True) static_dataset = StaticDataset(folder='/home/workspace/local-data/cape_processed/train/female', signature=signature_static, use_filled=False, resolutions=[512, 32], is_cape=True) for i, batch in enumerate(dataset): print(i)