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

add repeat option for dataset

Browse Source
This commit is contained in:
luni
2022-10-12 21:07:00 +00:00
parent 00a885c99f
commit a9e8072f90
1 changed files with 159 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dataset.py
+159 -153
View File
@@ -146,7 +146,7 @@ class DatasetFactory():
return tf_signature
def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512], in_memory=False):
def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512], in_memory=False, do_repeat=True):
# 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]))
@@ -248,12 +248,15 @@ class DatasetFactory():
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()
), 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).repeat()
), 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()
@@ -308,12 +311,13 @@ if __name__ == '__main__':
# 'uv_shaped_body': (3, 512, 256,),
}
dataset = factory.get_dataset(folder='/home/workspace/local-data/train_skirt_25_2',
dataset = factory.get_dataset(folder='/home/workspace/local-data/train',
batch_size=300,
signature=signature_dynamic,
use_filled=False,
resolutions=[512])
static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt_25_2',
resolutions=[512],
do_repeat=False)
static_dataset = StaticDataset(folder='/home/workspace/local-data/train',
signature=signature_static,
use_filled=False,
resolutions=[512])
@@ -330,10 +334,10 @@ if __name__ == '__main__':
sum_mask += batch['uv_mask'].numpy().sum(axis=0)
sum_mask[sum_mask == 0] = 1
avg_jumpsuit = sum_uv / sum_mask[:, :, np.newaxis]
avg_garment = sum_uv / sum_mask[:, :, np.newaxis]
x = np.transpose(avg_jumpsuit, [1,2,0])
x = uv2mesh(avg_jumpsuit[tf.newaxis,:,:,:])[0]
# 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
@@ -346,184 +350,186 @@ if __name__ == '__main__':
k = F[:,2]
)]).show()
np.save('custom_data/train_skirt_25_2_no_shape_avg_uv.npy', np.transpose(avg_jumpsuit, [1,2,0]))
np.save('custom_data/train_no_shape_avg_uv.npy', avg_garment)
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
pass
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_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]
# 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
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]
# 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))
# 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)
# 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
# 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()
# 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()
# 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))
# # @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']))
# 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])
# 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()
# _, 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()
# 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))
# 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)
# 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 losses import collision_loss as collision_penalty
# from losses import collision_penalty as collision_penalty_2
for d in dataset:
d
# for d in dataset:
# d
v = uv2mesh(d[0])
b = d[3]
m = d[4]
# v = uv2mesh(d[0])
# b = d[3]
# m = d[4]
b2 = tf.concat((b[1][tf.newaxis,:,:], b[0][tf.newaxis,:,:]), axis=0)
# 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(v, b, m)
# c2 = collision_penalty(v, b2, m)
c_ = collision_penalty_2(v, b, m)
c2_ = collision_penalty_2(v, b2, m)
# c_ = collision_penalty_2(v, b, m)
# c2_ = collision_penalty_2(v, b2, m)
print()
pass
# print()
# pass