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

curv by default is one_hot

Browse Source
This commit is contained in:
luni
2022-09-27 21:24:41 +00:00
parent a5eff778df
commit 64264364be
3 changed files with 30 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_curv.py
+11 -14
View File
@@ -77,17 +77,14 @@ def dice_dense(prediction, ground_truth, weight_map=None):
def train_step(batch): def train_step(batch):
with tf.GradientTape() as tape: 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 = cce(batch['uv_curvature'], pred)
ms_ssim = tf.reduce_mean(tf.image.ssim(pred, batch['uv_curvature'], max_val=1))
ce_uv_loss = scce(batch['uv_curvature'], pred)
ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
loss = ce_uv_loss loss = ce_uv_loss
mean_iou = get_mean_iou(batch['uv_curvature'], tf.math.argmax(pred, axis=-1)) mean_iou = get_mean_iou(tf.math.argmax(batch['uv_curvature'], axis=-1), tf.math.argmax(pred, axis=-1))
dice = dice_dense(pred, uv_curvature) dice = dice_dense(pred, batch['uv_curvature'])
gradients = tape.gradient(loss, model.trainable_variables) gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, 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']]) 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)) ms_ssim = tf.reduce_mean(tf.image.ssim(pred, uv_curvature, max_val=1))
loss = ce_uv_loss loss = ce_uv_loss
@@ -126,10 +123,9 @@ def train():
to_log = { **train_loss, **valid_loss} to_log = { **train_loss, **valid_loss}
if step % 2500 == 0: if step % 2500 == 0:
uv_curvature = tf.one_hot(data_valid['uv_curvature'], 7) pred = model([data_valid['uv_curvature'], data_valid['uv_mask']], training=False)
pred = model([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) pred_img = tf.math.argmax(pred, axis=-1, output_type=tf.dtypes.int32)
img = tf.concat((orig_img, pred_img), axis=2).numpy() 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)] } # imgs = { 'valid/images': [wandb.Image(img[i,:,:,np.newaxis] / 6.0) for i in range(4)] }
@@ -147,7 +143,8 @@ def train():
def test(): def test():
ckpt.restore(manager.latest_checkpoint) #TODO
# ckpt.restore(manager.latest_checkpoint)
for batch in dataset_train: for batch in dataset_train:
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
uv_curvature = tf.one_hot(batch['uv_curvature'], 7) uv_curvature = tf.one_hot(batch['uv_curvature'], 7)
@@ -169,7 +166,7 @@ if __name__ == '__main__':
parser.add_argument('--name', type=str) parser.add_argument('--name', type=str)
args = vars(parser.parse_args()) args = vars(parser.parse_args())
scce = tf.keras.losses.SparseCategoricalCrossentropy() cce = tf.keras.losses.CategoricalCrossentropy()
resolution_modifier = '' resolution_modifier = ''
if config['resolution'] != 512: if config['resolution'] != 512:
dataset.py
+18 -12
View File
@@ -61,7 +61,7 @@ class StaticDataset():
assert set(resolutions).issubset(set([32, 128, 512])) 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.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] 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.keys.append('uv_static_128')
self.resolutions = resolutions self.resolutions = resolutions
self.filled = '' self.filled = ''
@@ -132,14 +132,14 @@ class DatasetFactory():
for key in signature.keys(): for key in signature.keys():
if key == 'id': if key == 'id':
tf_signature.append(tf.TensorSpec(shape = (), dtype = tf.string)) 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)) tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.int32))
else: else:
tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32)) tf_signature.append(tf.TensorSpec(shape = signature[key], dtype = tf.float32))
return tf_signature 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 (use_filled and len(resolutions) >= 1) or (not use_filled and len(resolutions) == 1)
assert set(resolutions).issubset(set([32, 128, 512])) assert set(resolutions).issubset(set([32, 128, 512]))
@@ -212,7 +212,8 @@ class DatasetFactory():
uv_curvature = uv_curvature['uv_curvature' + filled] uv_curvature = uv_curvature['uv_curvature' + filled]
else: else:
uv_curvature = uv_curvature['curvature_uv'] 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(): 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] 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], # 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).repeat()
dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
inp=[i], if in_memory:
Tout=self._build_tf_signature(signature) dataset = dataset_idx.map(lambda i: tf.py_function(func=read_item,
), num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat() 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() # dataset = dataset_idx.map(read_item, num_parallel_calls=tf.data.AUTOTUNE).cache().batch(batch_size=batch_size, drop_remainder=True).repeat()
return dataset return dataset
@@ -281,7 +290,7 @@ if __name__ == '__main__':
# 'mesh_unposed': (14475, 3,), # 'mesh_unposed': (14475, 3,),
#'uv_unposed': (3, 512, 256,), #'uv_unposed': (3, 512, 256,),
'uv_unposed_32': (3, 32, 16,), 'uv_unposed_32': (3, 32, 16,),
# 'uv_curvature': (512, 256), 'uv_curvature': (512, 256, 7),
# 'mesh_curvature': (14475) # 'mesh_curvature': (14475)
} }
signature_static = { #'uv_static': (3, 512, 256,), signature_static = { #'uv_static': (3, 512, 256,),
@@ -335,9 +344,6 @@ if __name__ == '__main__':
print() print()
uv_res = 256 uv_res = 256
def m2uv_array(m2uv): def m2uv_array(m2uv):
verify_data.py
+1 -1
View File
@@ -32,7 +32,7 @@ if __name__ == '__main__':
# 'mesh_posed': (14475, 3,), # 'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,), # 'mesh_unposed': (14475, 3,),
'uv_unposed': (512, 256, 3), 'uv_unposed': (512, 256, 3),
'uv_curvature': (512, 256), 'uv_curvature': (512, 256, 7),
# 'mesh_curvature': (14475) # 'mesh_curvature': (14475)
} }
signature_static = { 'uv_static': (512, 256, 3), signature_static = { 'uv_static': (512, 256, 3),