From db3ea8ad48d3df99e1a20229f7cb85dcaa3986a0 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 29 Oct 2022 20:08:54 +0000 Subject: [PATCH] traied packing dataset, didn't work --- dataset_packed.py | 256 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 256 insertions(+) create mode 100644 dataset_packed.py diff --git a/dataset_packed.py b/dataset_packed.py new file mode 100644 index 0000000..4aa28a0 --- /dev/null +++ b/dataset_packed.py @@ -0,0 +1,256 @@ +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)