luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

lot of things

Browse Source
This commit is contained in:
luni
2022-08-06 12:27:24 +00:00
parent cd498e79f3
commit ee8f1cb68c
28 changed files with 4595 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset.py
+231 -39
View File
@@ -6,7 +6,7 @@ import random
import numpy as np
import tensorflow as tf
from custom_utils.custom_utils import nest_conditions
from custom_utils.custom_utils import nest_conditions, batch_to_dict
@@ -101,7 +101,7 @@ 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 = sorted(os.listdir(uv_static_folder))#[453:503]
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
random.shuffle(frames)
@@ -119,6 +119,8 @@ class DatasetFactory():
for key in signature.keys():
if key == 'id':
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string))
elif key == 'uv_curvature' or 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
@@ -139,16 +141,27 @@ class DatasetFactory():
item['id'] = id
if 'pose_shape' in signature.keys():
pose_shape = np.load(os.path.join(folder, 'pose_shape', id + '_' + 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 + '_' + 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 '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
uv_posed_body_no_orient = np.load(os.path.join(folder, 'uv_posed_body_no_orient', id + '_' + frame_nr + '.npz')); uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']; uv_posed_body_no_orient = np.transpose(uv_posed_body_no_orient, [2, 0, 1])
item['uv_body_posed'] = uv_posed_body_no_orient
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
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_body_posed'] = mesh_posed_body_no_orient
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
@@ -158,7 +171,12 @@ class DatasetFactory():
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 'uv_curvature' in signature.keys():
uv_curvature = np.load(os.path.join(folder, 'uv_curvature', id + '_' + frame_nr + '.npz')); uv_curvature = uv_curvature['curvature_uv']
item['uv_curvature'] = uv_curvature.astype(np.int32)
if 'mesh_curvature' in signature.keys():
mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', id + '_' + frame_nr + '.npz')); mesh_curvature = mesh_curvature['mesh_curvature']
item['mesh_curvature'] = mesh_curvature.astype(np.int32)
# item = { **static_info, **item }
as_list = []
@@ -168,6 +186,10 @@ class DatasetFactory():
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(lambda i: tf.py_function(func=read_item,
inp=[i],
Tout=self._build_tf_signature(signature)
@@ -186,44 +208,232 @@ class DatasetFactory():
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
if __name__ == '__main__':
from custom_utils.custom_utils import batch_to_dict
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()
F = np.load("custom_utils/smpl_faces.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': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,),
# 'uv_unposed': (3, 512, 256,),
# 'mesh_body_posed': (14475, 3,),
# 'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
'uv_unposed': (3, 512, 256,),
# 'uv_curvature': (512, 256),
# 'mesh_curvature': (14475)
}
signature_static = { #'uv_static': (3, 512, 256,),
# 'mesh_static': (14475, 3,),
'mesh_mask': (14475,),
# 'uv_mask': (512, 256,),
'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,
dataset = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train',
batch_size=1,
signature=signature_dynamic)
static_dataset = StaticDataset(folder='/home/workspace/data/dataset/stylegan_new/train',
static_dataset = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train',
signature=signature_static)
@tf.function
# sum_uv = np.zeros((3, 512, 256), 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)
# avg_jumpsuit = sum_uv / sum_mask[np.newaxis, :, :]
# sum_mask[sum_mask == 0] = 1
# x = np.transpose(avg_jumpsuit, [1,2,0])
# x = uv2mesh(avg_jumpsuit[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/jumpsuit_20_avg_uv.npy', np.transpose(avg_jumpsuit, [1,2,0]))
# print()
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))) / tf.reduce_sum(mask)
# 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))
err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], batch['mesh_body_posed'])
all_error.append(err.numpy())
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))
@@ -232,25 +442,7 @@ if __name__ == '__main__':
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