272 lines
13 KiB
Python
272 lines
13 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
|
|
|
|
|
|
|
|
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
|
|
|