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):
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: