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 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 (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
assert set(resolutions).issubset(set([32, 64, 128, 256, 512])) 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, dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
inp=[i], inp=[i],
Tout=self._build_tf_signature(signature) 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: else:
dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item, dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
inp=[i], inp=[i],
Tout=self._build_tf_signature(signature) 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() # 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,), # '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, batch_size=300,
signature=signature_dynamic, signature=signature_dynamic,
use_filled=False, use_filled=False,
resolutions=[512]) resolutions=[512],
static_dataset = StaticDataset(folder='/home/workspace/local-data/train_skirt_25_2', do_repeat=False)
static_dataset = StaticDataset(folder='/home/workspace/local-data/train',
signature=signature_static, signature=signature_static,
use_filled=False, use_filled=False,
resolutions=[512]) resolutions=[512])
@@ -330,10 +334,10 @@ if __name__ == '__main__':
sum_mask += batch['uv_mask'].numpy().sum(axis=0) sum_mask += batch['uv_mask'].numpy().sum(axis=0)
sum_mask[sum_mask == 0] = 1 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 = np.transpose(avg_garment, [1,2,0])
x = uv2mesh(avg_jumpsuit[tf.newaxis,:,:,:])[0] x = uv2mesh(avg_garment[tf.newaxis,:,:,:])[0]
go.Figure(data=[ go.Figure(data=[
go.Mesh3d( go.Mesh3d(
# 8 vertices of a cube # 8 vertices of a cube
@@ -346,184 +350,186 @@ if __name__ == '__main__':
k = F[:,2] k = F[:,2]
)]).show() )]).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() 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') # uv_res = 256
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 # def m2uv_array(m2uv):
uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis] # result = np.zeros((14475, 2), dtype=np.int32)
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] # 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) # return result
uv_map_mask = unique_uv_map_mask(uv_map_discreet)
uv_counter = np.zeros((uv_map.shape[0])) # def unique_uv_map_mask(uv_map_discreet):
uv_counter[m2uv[:, 0]] = 1 # mask = np.zeros((uv_map_discreet.shape[0]), dtype=bool)
uv_counter[m2uv[:, 1]] += 1 # temp = np.zeros((uv_res * 2, uv_res))
uv_counter = uv_counter[uv_map_mask] # 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): # for i, batch in enumerate(dataset):
print(i) # print(i)
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) # batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
uv_unposed = np.squeeze(batch['uv_unposed'].numpy()) # uv_unposed = np.squeeze(batch['uv_unposed'].numpy())
mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy()) # mesh_unposed = np.squeeze(batch['mesh_unposed'].numpy())
mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool) # mask = np.squeeze(batch['mesh_mask'].numpy()).astype(bool)
values = mesh_unposed # values = mesh_unposed
uv_values = np.zeros((uv_map.shape[0], 3)) # uv_values = np.zeros((uv_map.shape[0], 3))
uv_values[m2uv[:, 0]] = values # uv_values[m2uv[:, 0]] = values
uv_values[m2uv[:, 1]] += values # uv_values[m2uv[:, 1]] += values
uv_counter = np.zeros((uv_map.shape[0])) # uv_counter = np.zeros((uv_map.shape[0]))
uv_counter[m2uv[:, 0]] = 1 # uv_counter[m2uv[:, 0]] = 1
uv_counter[m2uv[:, 1]] += 1 # uv_counter[m2uv[:, 1]] += 1
uv_values /= uv_counter[:, np.newaxis] # uv_values /= uv_counter[:, np.newaxis]
x = uv_values # x = uv_values
# verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1) # # verts = np.concatenate((verts, np.zeros((15068, 1))), axis=-1)
im = ((x - x.min()) / (x.max() - x.min())) * 255 # im = ((x - x.min()) / (x.max() - x.min())) * 255
x = np.transpose(uv_unposed, [1,2,0]) # x = np.transpose(uv_unposed, [1,2,0])
x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :] # x = uv2mesh(uv_unposed[tf.newaxis,:,:,:])[0].numpy()[mask, :]
_, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask) # _, faces = _maskFaces(np.zeros((14475, 3), dtype=np.float32), F, mask)
go.Figure(data=[ # go.Figure(data=[
go.Mesh3d( # go.Mesh3d(
# 8 vertices of a cube # # 8 vertices of a cube
x=x[:,0], # x=x[:,0],
y=x[:,1], # y=x[:,1],
z=x[:,2], # z=x[:,2],
# i, j and k give the vertices of triangles # # i, j and k give the vertices of triangles
i = faces[:,0], # i = faces[:,0],
j = faces[:,1], # j = faces[:,1],
k = faces[:,2] # k = faces[:,2]
)]).show() # )]).show()
print() # print()
# @tf.function # # @tf.function
def euclid_error(mask, gt, pred): # 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.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)) # return tf.math.sqrt(tf.reduce_sum(mask[:,:,tf.newaxis] * (gt - pred) ** 2, axis=-1))
all_error = [] # all_error = []
for batch in dataset: # for batch in dataset:
batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic)) # batch = static_dataset.expand_batch(batch_to_dict(batch, signature_dynamic))
mesh_reconstruct = uv2mesh(batch['uv_posed']) # mesh_reconstruct = uv2mesh(batch['uv_posed'])
err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], uv2mesh(batch['uv_posed'])) # err = euclid_error(batch['mesh_mask'], batch['mesh_posed'], uv2mesh(batch['uv_posed']))
i = 296 # i = 296
err_normed = (err[i].numpy() - err[i].numpy().min()) / (err[i].numpy().max() - err[i].numpy().min()) # 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]) # 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)) # _, 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), :] # 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), :] # y = mesh_reconstruct[i].numpy()[batch['mesh_mask'][i].numpy().astype(bool), :]
colors = colors[batch['mesh_mask'][i].numpy().astype(bool), :] # colors = colors[batch['mesh_mask'][i].numpy().astype(bool), :]
fig = go.Figure(data=[ # fig = go.Figure(data=[
go.Mesh3d( # go.Mesh3d(
# 8 vertices of a cube # # 8 vertices of a cube
x=x[:,0], # x=x[:,0],
y=x[:,1], # y=x[:,1],
z=x[:,2], # z=x[:,2],
# i, j and k give the vertices of triangles # # i, j and k give the vertices of triangles
i = faces[:,0], # i = faces[:,0],
j = faces[:,1], # j = faces[:,1],
k = faces[:,2] # k = faces[:,2]
), # ),
go.Mesh3d( # go.Mesh3d(
# 8 vertices of a cube # # 8 vertices of a cube
x=y[:,0], # x=y[:,0],
y=y[:,1], # y=y[:,1],
z=y[:,2], # z=y[:,2],
# i, j and k give the vertices of triangles # # i, j and k give the vertices of triangles
i = faces[:,0], # i = faces[:,0],
j = faces[:,1], # j = faces[:,1],
k = faces[:,2], # k = faces[:,2],
)]) # )])
fig.show() # fig.show()
fig = ex.scatter_3d(x=x[:,0], y=x[:,1], z=x[:,2]) # fig = ex.scatter_3d(x=x[:,0], y=x[:,1], z=x[:,2])
fig.update_traces(marker={'color': colors, 'size': 5}) # fig.update_traces(marker={'color': colors, 'size': 5})
fig.show() # fig.show()
all_error.append(err.numpy().sum()) # all_error.append(err.numpy().sum())
# b = nest_conditions(static_dataset.expand_batch(batch_to_dict(d, signature_dynamic))) # # b = nest_conditions(static_dataset.expand_batch(batch_to_dict(d, signature_dynamic)))
print(len(all_error)) # 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_loss as collision_penalty
from losses import collision_penalty as collision_penalty_2 # from losses import collision_penalty as collision_penalty_2
for d in dataset: # for d in dataset:
d # d
v = uv2mesh(d[0]) # v = uv2mesh(d[0])
b = d[3] # b = d[3]
m = d[4] # 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) # c = collision_penalty(v, b, m)
c2 = collision_penalty(v, b2, m) # c2 = collision_penalty(v, b2, m)
c_ = collision_penalty_2(v, b, m) # c_ = collision_penalty_2(v, b, m)
c2_ = collision_penalty_2(v, b2, m) # c2_ = collision_penalty_2(v, b2, m)
print() # print()
pass # pass