571 lines
26 KiB
Python
571 lines
26 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 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, 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' + self.filled + 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, use_one_hot=False):
|
|
# assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
|
|
assert set(resolutions).issubset(set([32, 64, 128, 256, 512]))
|
|
|
|
# static_data = StaticDataset(folder, signature)
|
|
dataset_idx = DatasetIdx(folder, shuffle, is_cape=is_cape)
|
|
filled = ''
|
|
if use_filled:
|
|
filled = '_filled'
|
|
separator = '_'
|
|
if is_cape:
|
|
separator = '.'
|
|
|
|
one_hot = ''
|
|
one_hot_2 = ''
|
|
if use_one_hot:
|
|
one_hot = '_one_hot'
|
|
one_hot_2 = '_one_hot_2'
|
|
|
|
|
|
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()
|
|
|
|
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 + separator + frame_nr + '.npy'))
|
|
|
|
##############################
|
|
pose_shape[:3] = 0
|
|
##############################
|
|
|
|
item['pose_shape'] = pose_shape
|
|
if 'uv_posed' in signature.keys():
|
|
uv_posed = np.load(os.path.join(folder, 'uv_posed', id + separator + frame_nr + '.npz')); uv_posed = uv_posed['uv_gt']
|
|
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_body_no_orient' in signature.keys():
|
|
# mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + '_' + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
|
|
# item['mesh_posed_body_no_orient'] = mesh_posed_body_no_orient
|
|
|
|
if 'mesh_body_posed' in signature.keys():
|
|
mesh_posed_body_no_orient = np.load(os.path.join(folder, 'mesh_posed_body_no_orient', id + separator + frame_nr + '.npz')); mesh_posed_body_no_orient = mesh_posed_body_no_orient['mesh_posed_body_no_orient']
|
|
item['mesh_body_posed'] = mesh_posed_body_no_orient
|
|
if 'mesh_posed' in signature.keys():
|
|
mesh_posed = np.load(os.path.join(folder, 'mesh_posed', id + separator + 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
|
|
if 'mesh_unposed' in signature.keys():
|
|
# print('one_hot')
|
|
mesh_unposed_no_shape = np.load(os.path.join(folder, 'mesh_unposed_no_shape' + one_hot, id + separator + frame_nr + '.npz')); mesh_unposed_no_shape = mesh_unposed_no_shape['mesh_unposed_no_shape']
|
|
item['mesh_unposed'] = mesh_unposed_no_shape
|
|
|
|
|
|
if 'mesh_curvature' in signature.keys():
|
|
mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + separator + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
|
|
item['mesh_curvature'] = 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():
|
|
uv_curvature = np.load(os.path.join(folder, 'uv_curvature' + filled + res_string, id + separator + frame_nr + '.npz'));
|
|
if filled != '' or res_string != '':
|
|
uv_curvature = uv_curvature['uv_curvature' + filled]
|
|
else:
|
|
uv_curvature = uv_curvature['curvature_uv']
|
|
uv_curvature = tf.one_hot(uv_curvature, 7)
|
|
item['uv_curvature' + res_string] = uv_curvature
|
|
|
|
if 'uv_unposed' + res_string in signature.keys():
|
|
uv_unposed_no_shape = np.load(os.path.join(folder, 'uv_unposed_no_shape' + one_hot_2 + res_string, id + separator + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled]
|
|
item['uv_unposed' + res_string] = uv_unposed_no_shape
|
|
|
|
if 'uv_body_posed' + res_string in signature.keys():
|
|
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient_2' + filled + res_string, id + separator + frame_nr + '.npz'))
|
|
uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient' + filled]
|
|
item['uv_body_posed' + res_string] = uv_posed_body_no_orient
|
|
|
|
|
|
# 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()
|
|
|
|
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()
|
|
|
|
# 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, )
|
|
# )
|
|
|
|
|
|
|
|
|
|
|
|
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': (3, 32, 16,),
|
|
# '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/train',
|
|
batch_size=300,
|
|
signature=signature_dynamic,
|
|
use_filled=False,
|
|
resolutions=[512],
|
|
do_repeat=False,
|
|
use_one_hot=True)
|
|
static_dataset = StaticDataset(folder='/home/workspace/local-data/train',
|
|
signature=signature_static,
|
|
use_filled=False,
|
|
resolutions=[512])
|
|
|
|
|
|
# for i, batch in enumerate(dataset):
|
|
# print(i)
|
|
|
|
|
|
sum_uv = np.zeros((512, 256, 3), dtype=np.float32)
|
|
sum_mask = np.zeros((512, 256), dtype=np.float32)
|
|
for i, batch in enumerate(dataset):
|
|
print(i)
|
|
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
|
sum_uv += batch['uv_unposed'].numpy().sum(axis=0)
|
|
sum_mask += batch['uv_mask'].numpy().sum(axis=0)
|
|
|
|
sum_mask[sum_mask == 0] = 1
|
|
avg_garment = sum_uv / sum_mask[:, :, np.newaxis]
|
|
|
|
# x = np.transpose(avg_garment, [1,2,0])
|
|
x = uv2mesh(avg_garment[tf.newaxis,:,:,:])[0]
|
|
go.Figure(data=[
|
|
go.Mesh3d(
|
|
# 8 vertices of a cube
|
|
x=x[:,0],
|
|
y=x[:,1],
|
|
z=x[:,2],
|
|
# i, j and k give the vertices of triangles
|
|
i = F[:,0],
|
|
j = F[:,1],
|
|
k = F[:,2]
|
|
)]).show()
|
|
|
|
np.save('custom_data/train_no_shape_one_hot_avg_uv.npy', avg_garment)
|
|
|
|
# print()
|
|
|
|
# pass
|
|
|
|
|
|
# uv_res = 256
|
|
|
|
# def m2uv_array(m2uv):
|
|
# result = np.zeros((14475, 2), dtype=np.int32)
|
|
# for k, v in m2uv.items():
|
|
# if len(v) == 2:
|
|
# result[k] = list(v)
|
|
# elif len(v) == 1:
|
|
# x = v.pop()
|
|
# result[k] = [x, x]
|
|
# else:
|
|
# assert False
|
|
|
|
# return result
|
|
|
|
# def unique_uv_map_mask(uv_map_discreet):
|
|
# mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
|
|
# temp = np.zeros((uv_res * 2, uv_res))
|
|
# for p in uv_map_discreet:
|
|
# temp[p[0], p[1]] += 1
|
|
# # temp = temp == 1
|
|
# for i, p in enumerate(uv_map_discreet):
|
|
# if temp[p[0], p[1]] == 1:
|
|
# mask[i] = True
|
|
# elif temp[p[0], p[1]] == 2:
|
|
# mask[i] = True
|
|
# temp[p[0], p[1]] = 0
|
|
# return mask
|
|
|
|
# def uv_to_mesh(uv, m2uv):
|
|
# res = uv.shape[0]
|
|
# mesh = np.zeros((14475, 3))
|
|
# uv_pixels = uv_map_discreet[m2uv]
|
|
# # counter = np.ones((14475))
|
|
# # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
|
|
# mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
|
|
# return mesh
|
|
|
|
|
|
# uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
|
|
# m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
|
|
# m2uv = m2uv_array(m2uv)
|
|
# uv_mask = np.load('custom_utils/mask_256x512.npy')
|
|
# curvature_label_type = '6_25'
|
|
|
|
# uv_res = 256
|
|
# uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
|
|
# uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
|
|
|
|
# uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
|
|
# uv_map_mask = unique_uv_map_mask(uv_map_discreet)
|
|
# uv_counter = np.zeros((uv_map.shape[0]))
|
|
# uv_counter[m2uv[:, 0]] = 1
|
|
# uv_counter[m2uv[:, 1]] += 1
|
|
# uv_counter = uv_counter[uv_map_mask]
|
|
|
|
|
|
|
|
# for i, batch in enumerate(dataset):
|
|
# print(i)
|
|
# batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
|
|
|
# uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
|
|
# mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
|
|
# mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
|
|
|
|
|
|
# values = mesh_unposed
|
|
# uv_values = np.zeros((uv_map.shape[0], 3))
|
|
# uv_values[m2uv[:, 0]] = values
|
|
# uv_values[m2uv[:, 1]] += values
|
|
# uv_counter = np.zeros((uv_map.shape[0]))
|
|
# uv_counter[m2uv[:, 0]] = 1
|
|
# uv_counter[m2uv[:, 1]] += 1
|
|
# uv_values /= uv_counter[:, np.newaxis]
|
|
# x = uv_values
|
|
# # verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
|
|
# im = ((x - x.min()) / (x.max() - x.min())) * 255
|
|
|
|
|
|
# x = np.transpose(uv_unposed, [1,2,0])
|
|
# x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :]
|
|
# _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
|
|
# go.Figure(data=[
|
|
# go.Mesh3d(
|
|
# # 8 vertices of a cube
|
|
# x=x[:,0],
|
|
# y=x[:,1],
|
|
# z=x[:,2],
|
|
# # i, j and k give the vertices of triangles
|
|
# i = faces[:,0],
|
|
# j = faces[:,1],
|
|
# k = faces[:,2]
|
|
# )]).show()
|
|
|
|
|
|
# print()
|
|
|
|
|
|
|
|
|
|
# # @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)), axis=-1) / tf.reduce_sum(mask, axis=-1)
|
|
# return tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1))
|
|
|
|
# all_error = []
|
|
# for batch in dataset:
|
|
# batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
|
|
# mesh_reconstruct = uv2mesh(batch['uv_posed'])
|
|
# err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], uv2mesh(batch['uv_posed']))
|
|
|
|
# i = 296
|
|
# err_normed = (err[i].numpy() - err[i].numpy().min()) / (err[i].numpy().max() - err[i].numpy().min())
|
|
# colors = np.array([cm.color('jet', 1 - i) for i in err_normed])
|
|
|
|
# _, faces = _maskFaces(batch['mesh_posed'][i].numpy(), F, batch['mesh_mask'][i].numpy().astype(bool))
|
|
# x = batch['mesh_posed'][i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
|
|
# y = mesh_reconstruct[i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
|
|
# colors = colors[batch['mesh_mask'][i].numpy().astype(bool), :]
|
|
# fig = go.Figure(data=[
|
|
# go.Mesh3d(
|
|
# # 8 vertices of a cube
|
|
# x=x[:,0],
|
|
# y=x[:,1],
|
|
# z=x[:,2],
|
|
# # i, j and k give the vertices of triangles
|
|
# i = faces[:,0],
|
|
# j = faces[:,1],
|
|
# k = faces[:,2]
|
|
# ),
|
|
# go.Mesh3d(
|
|
# # 8 vertices of a cube
|
|
# x=y[:,0],
|
|
# y=y[:,1],
|
|
# z=y[:,2],
|
|
# # i, j and k give the vertices of triangles
|
|
# i = faces[:,0],
|
|
# j = faces[:,1],
|
|
# k = faces[:,2],
|
|
# )])
|
|
# fig.show()
|
|
|
|
# fig = ex.scatter_3d(x=x[:,0], y=x[:,1], z=x[:,2])
|
|
# fig.update_traces(marker={'color': colors, 'size': 5})
|
|
# fig.show()
|
|
|
|
# all_error.append(err.numpy().sum())
|
|
# # 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
|
|
|
|
|
|
# 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
|
|
|