From c89677abd714938d43973b936fad707cd288a086 Mon Sep 17 00:00:00 2001 From: Starbeamrainbowlabs Date: Fri, 2 Sep 2022 16:57:59 +0100 Subject: [PATCH] dataset: explicit reshape --- aimodel/src/lib/dataset/dataset.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/aimodel/src/lib/dataset/dataset.py b/aimodel/src/lib/dataset/dataset.py index 12585a5..ff24c90 100644 --- a/aimodel/src/lib/dataset/dataset.py +++ b/aimodel/src/lib/dataset/dataset.py @@ -22,6 +22,9 @@ def parse_item(metadata): water = tf.io.parse_tensor(parsed["waterdepth"], out_type=tf.float32) # [channels, width, height] → [width, height, channels] - ref ConvNeXt does not support data_format=channels_first + rainfall = tf.reshape(rainfall, tf.constant(metadata["rainfallradar"], dtype=tf.int32)) + water = tf.reshape(water, tf.constant(metadata["waterdepth"], dtype=tf.int32)) + rainfall = tf.transpose(rainfall, [1, 2, 0]) # [width, height] → [width, height, channels] water = tf.expand_dims(water, axis=-1)