From 54a841efe9395bf2c86efabc6ad64b02f7958c0c Mon Sep 17 00:00:00 2001 From: Starbeamrainbowlabs Date: Thu, 24 Nov 2022 17:56:07 +0000 Subject: [PATCH] train_mono_predict: convert to correct format --- aimodel/src/subcommands/train_mono_predict.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/aimodel/src/subcommands/train_mono_predict.py b/aimodel/src/subcommands/train_mono_predict.py index e413f47..709bef0 100644 --- a/aimodel/src/subcommands/train_mono_predict.py +++ b/aimodel/src/subcommands/train_mono_predict.py @@ -111,8 +111,9 @@ def do_png(args, ai, dataset, model_params): for water_predict in water_predict_batch: # [ width, height, softmax_probabilities ] → [ batch, width, height ] water_predict = tf.math.argmax(water_predict, axis=-1) - # [ width, height ] + # [ width, height, bins ] water_actual = tf.squeeze(water[i_batch]) + water_actual = tf.math.argmax(water_actual, axis=-1) segmentation_plot( water_actual, water_predict,