run-35
This commit is contained in:
+177
-1
@@ -6,6 +6,10 @@ 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):
|
||||
@@ -47,10 +51,182 @@ class DatasetSimple(tf.data.Dataset):
|
||||
args=(folder, )
|
||||
)
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user