From 64264364be5dcb10604d6743117aae7425d6dff1 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Tue, 27 Sep 2022 21:24:41 +0000 Subject: [PATCH] curv by default is one_hot --- auto_encoder_curv.py | 25 +++++++++++-------------- dataset.py | 30 ++++++++++++++++++------------ verify_data.py | 2 +- 3 files changed, 30 insertions(+), 27 deletions(-) diff --git a/auto_encoder_curv.py b/auto_encoder_curv.py index 04f2788..52b0b5a 100644 --- a/auto_encoder_curv.py +++ b/auto_encoder_curv.py @@ -77,17 +77,14 @@ def dice_dense(prediction, ground_truth, weight_map=None): def train_step(batch): with tf.GradientTape() as tape: - uv_curvature = tf.one_hot(batch['uv_curvature'], 7) + pred = model([batch['uv_curvature'], batch['uv_mask']], training=True) - pred = model([uv_curvature, batch['uv_mask']], training=True) - - - ce_uv_loss = scce(batch['uv_curvature'], pred) - ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1)) + ce_uv_loss = cce(batch['uv_curvature'], pred) + ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1)) loss = ce_uv_loss - mean_iou = get_mean_iou(batch['uv_curvature'], tf.math.argmax(pred, axis=-1)) - dice = dice_dense(pred, uv_curvature) + mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1)) + dice = dice_dense(pred, batch['uv_curvature']) gradients = tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables)) @@ -101,7 +98,7 @@ def valid_step(batch): pred = model([uv_curvature, batch['uv_mask']]) - ce_uv_loss = scce(batch['uv_curvature'], pred) + ce_uv_loss = cce(batch['uv_curvature'], pred) ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1)) loss = ce_uv_loss @@ -126,10 +123,9 @@ def train(): to_log = { **train_loss, **valid_loss} if step % 2500 == 0: - uv_curvature = tf.one_hot(data_valid['uv_curvature'], 7) - pred = model([uv_curvature, data_valid['uv_mask']], training=False) + pred = model([data_valid['uv_curvature'], data_valid['uv_mask']], training=False) - orig_img = data_valid['uv_curvature'] + orig_img = tf.math.argmax(data_valid['uv_curvature'], axis=-1, output_type=tf.dtypes.int32) pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32) img = tf.concat((orig_img, pred_img), axis=2).numpy() # imgs = { 'valid/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(4)] } @@ -147,7 +143,8 @@ def train(): def test(): - ckpt.restore(manager.latest_checkpoint) + #TODO + # ckpt.restore(manager.latest_checkpoint) for batch in dataset_train: batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) uv_curvature = tf.one_hot(batch['uv_curvature'], 7) @@ -169,7 +166,7 @@ if __name__ == '__main__': parser.add_argument('--name', type=str) args = vars(parser.parse_args()) - scce = tf.keras.losses.SparseCategoricalCrossentropy() + cce = tf.keras.losses.CategoricalCrossentropy() resolution_modifier = '' if config['resolution'] != 512: diff --git a/dataset.py b/dataset.py index db20d92..0079fcb 100644 --- a/dataset.py +++ b/dataset.py @@ -61,7 +61,7 @@ class StaticDataset(): assert set(resolutions).issubset(set([32, 128, 512])) self.possible_keys = ['uv_static', 'uv_static_32', 'uv_static_128', 'mesh_static', 'mesh_mask', 'uv_mask', 'uv_mask_32', 'uv_mask_128', 'uv_shaped_body'] self.keys = [value for value in signature.keys() if value in self.possible_keys] - if 'uv_static_128' in signature.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 = '' @@ -132,14 +132,14 @@ 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' or key == 'uv_curvature_32' or key == 'uv_curvature_128': + 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]): + def get_dataset(self, folder, batch_size, signature, shuffle=True, use_filled=False, resolutions=[512], in_memory=True): # assert (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1) assert set(resolutions).issubset(set([32, 128, 512])) @@ -212,7 +212,8 @@ class DatasetFactory(): uv_curvature = uv_curvature['uv_curvature' + filled] else: uv_curvature = uv_curvature['curvature_uv'] - item['uv_curvature' + res_string] = uv_curvature.astype(np.int32) + 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' + filled + res_string, id + '_' + frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape' + filled] @@ -235,10 +236,18 @@ class DatasetFactory(): # 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) - ), 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).repeat() + 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() + # dataset = dataset_idx.map(read_item, num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat() return dataset @@ -281,7 +290,7 @@ if __name__ == '__main__': # 'mesh_unposed': (14475, 3,), #'uv_unposed': (3, 512, 256,), 'uv_unposed_32': (3, 32, 16,), - # 'uv_curvature': (512, 256), + 'uv_curvature': (512, 256, 7), # 'mesh_curvature': (14475) } signature_static = { #'uv_static': (3, 512, 256,), @@ -335,9 +344,6 @@ if __name__ == '__main__': print() - - - uv_res = 256 def m2uv_array(m2uv): diff --git a/verify_data.py b/verify_data.py index 199efb1..6275018 100644 --- a/verify_data.py +++ b/verify_data.py @@ -32,7 +32,7 @@ if __name__ == '__main__': # 'mesh_posed': (14475, 3,), # 'mesh_unposed': (14475, 3,), 'uv_unposed': (512, 256, 3), - 'uv_curvature': (512, 256), + 'uv_curvature': (512, 256, 7), # 'mesh_curvature': (14475) } signature_static = { 'uv_static': (512, 256, 3),