mirror of
https://github.com/sbrl/research-rainfallradar
synced 2024-11-05 01:23:02 +00:00
train_mono_predict: convert to correct format
This commit is contained in:
parent
105dc5bc56
commit
54a841efe9
1 changed files with 2 additions and 1 deletions
|
@ -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,
|
||||
|
|
Loading…
Reference in a new issue