curv by default is one_hot
This commit is contained in:
+11
-14
@@ -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:
|
||||
|
||||
+18
-12
@@ -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):
|
||||
|
||||
+1
-1
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user