diff --git a/dataset.py b/dataset.py index 322bada..4eb15d7 100644 --- a/dataset.py +++ b/dataset.py @@ -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 \ No newline at end of file