257 lines
11 KiB
Python
257 lines
11 KiB
Python
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': (), 'frame_nr': (),
|
|
'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)
|