Implement support for (optionally) taking a heightmap in

aimodel/src/lib/ai/model_rainfallwater_mono.py

@@ -8,7 +8,7 @@ from .components.convnext_inverse import do_convnext_inverse
 from .components.LayerStack2Image import LayerStack2Image
 from .components.LossCrossentropy import LossCrossentropy
 
-def model_rainfallwater_mono(metadata, shape_water_out, model_arch_enc="convnext_xtiny", model_arch_dec="convnext_i_xtiny", feature_dim=512, batch_size=64, water_bins=2, learning_rate=None):
+def model_rainfallwater_mono(metadata, shape_water_out, model_arch_enc="convnext_xtiny", model_arch_dec="convnext_i_xtiny", feature_dim=512, batch_size=64, water_bins=2, learning_rate=None, heightmap_input=False):
 	"""Makes a new rainfall / waterdepth mono model.
 
 	Args:
@@ -19,6 +19,7 @@ def model_rainfallwater_mono(metadata, shape_water_out, model_arch_enc="convnext
 		model_arch_dec (str, optional): The architecture code for the underlying (inverted) ConvNeXt model for the decoder. Defaults to "convnext_i_xtiny".
 		batch_size (int, optional): The batch size. Reduce to save memory. Defaults to 64.
 		water_bins (int, optional): The number of classes that the water depth output oft he segmentation head should be binned into. Defaults to 2.
+		heightmap_input (bool, option): Whether a heightmap is being passed as an input to the model or not. Required to ensure we know how many channels the model will be taking in (the heightmap takes u  an additional input channel). Default: false.
 		learning_rate (float, optional): The (initial) learning rate. YOU DO NOT USUALLY NEED TO CHANGE THIS. For experimental purposes only. Defaults to None, which means it will be determined automatically.
 
 	Returns:
@@ -27,7 +28,10 @@ def model_rainfallwater_mono(metadata, shape_water_out, model_arch_enc="convnext
 	rainfall_channels, rainfall_height, rainfall_width = metadata["rainfallradar"] # shape = [channels, height, weight]
 	# BUG: We somehow *still* have the rainfall radar data transposed incorrectly! I have no idea how this happened. dataset_mono fixes it with (another) transpose
 	
-	print("RAINFALL channels", rainfall_channels, "width", rainfall_width, "height", rainfall_height)
+	if heightmap_input:
+		rainfall_channels += 1
+	
+	print("RAINFALL channels", rainfall_channels, "width", rainfall_width, "height", rainfall_height, "HEIGHTMAP_INPUT", heightmap_input)
 	out_water_width, out_water_height = shape_water_out
 	
 	layer_input = tf.keras.layers.Input(

aimodel/src/lib/dataset/dataset_mono.py

@@ -10,11 +10,12 @@ from lib.dataset.read_metadata import read_metadata
 
 from ..io.readfile import readfile
 from .shuffle import shuffle
+from .parse_heightmap import parse_heightmap
 
 
 
 # TO PARSE:
-def parse_item(metadata, shape_water_desired=[100,100], water_threshold=0.1, water_bins=2):
+def parse_item(metadata, shape_water_desired=[100,100], water_threshold=0.1, water_bins=2, heightmap=None):
 	water_height_source, water_width_source = metadata["waterdepth"]
 	water_height_target, water_width_target = shape_water_desired
 	water_offset_x = math.ceil((water_width_source - water_width_target) / 2)
@@ -25,6 +26,9 @@ def parse_item(metadata, shape_water_desired=[100,100], water_threshold=0.1, wat
 	print("DEBUG DATASET:water_threshold", water_threshold)
 	print("DEBUG DATASET:water_bins", water_bins)
 	
+	if heightmap is not None:
+		heightmap = tf.expand_dims(heightmap, axis=-1)
+	
 	def parse_item_inner(item):
 		parsed = tf.io.parse_single_example(item, features={
 			"rainfallradar": tf.io.FixedLenFeature([], tf.string),
@@ -46,10 +50,17 @@ def parse_item(metadata, shape_water_desired=[100,100], water_threshold=0.1, wat
 		# I can't believe in this entire project I have yet to get the rotation of the rainfall radar data correct....!
 		# %TRANSPOSE%
 		rainfall = tf.transpose(rainfall, [2, 1, 0])
+		if heightmap is not None:
+			rainfall = tf.concat([rainfall, heightmap], axis=-1)
 		
 		# rainfall = tf.image.resize(rainfall, tf.cast(tf.constant(metadata["rainfallradar"]) / 2, dtype=tf.int32))
 		water = tf.expand_dims(water, axis=-1) # [width, height] → [width, height, channels=1]
-		water = tf.image.crop_to_bounding_box(water, water_offset_x, water_offset_y, water_width_target, water_height_target)
+		water = tf.image.crop_to_bounding_box(water,
+			offset_width=water_offset_x,
+			offset_height=water_offset_y,
+			target_width=water_width_target,
+			target_height=water_height_target
+		)
 		
 		print("DEBUG:dataset BEFORE_SQUEEZE water", water.shape)
 		water = tf.squeeze(water)
@@ -64,17 +75,21 @@ def parse_item(metadata, shape_water_desired=[100,100], water_threshold=0.1, wat
 	
 	return tf.function(parse_item_inner)
 
-def make_dataset(filepaths, compression_type="GZIP", parallel_reads_multiplier=1.5, shuffle_buffer_size=128, batch_size=64, prefetch=True, shuffle=True, **kwargs):
+def make_dataset(filepaths, compression_type="GZIP", parallel_reads_multiplier=1.5, shuffle_buffer_size=128, batch_size=64, prefetch=True, shuffle=True, filepath_heightmap=None **kwargs):
 	if "NO_PREFETCH" in os.environ:
 		logger.info("disabling data prefetching.")
 	
+	heightmap = None
+	if filepath_heightmap is not None:
+		heightmap = parse_heightmap(filepath_heightmap)
+	
 	dataset = tf.data.TFRecordDataset(filepaths,
 		compression_type=compression_type,
 		num_parallel_reads=math.ceil(os.cpu_count() * parallel_reads_multiplier) if parallel_reads_multiplier > 0 else None
 	)
 	if shuffle:
 		dataset = dataset.shuffle(shuffle_buffer_size)
-	dataset = dataset.map(parse_item(**kwargs), num_parallel_calls=tf.data.AUTOTUNE)
+	dataset = dataset.map(parse_item(heightmap=heightmap, **kwargs), num_parallel_calls=tf.data.AUTOTUNE)
 	
 	if batch_size != None:
 		dataset = dataset.batch(batch_size, drop_remainder=True)

aimodel/src/lib/dataset/parse_heightmap.py

@@ -0,0 +1,12 @@
+import json
+
+import tensorflow as tf
+from ..io.readfile import readfile
+
+def parse_heightmap(filepath_heightmap):
+	obj = json.loads(readfile(filepath_heightmap))
+	
+	result = tf.constant(obj.data)
+	result = tf.transpose(result, [1,0]) # [ height, width ] → [ width, height ]
+	
+	return result

aimodel/src/subcommands/train_mono.py

@@ -12,9 +12,10 @@ def parse_args():
 	parser = argparse.ArgumentParser(description="Train an mono rainfall-water model on a directory of .tfrecord.gz rainfall+waterdepth_label files.")
 	# parser.add_argument("--config", "-c", help="Filepath to the TOML config file to load.", required=True)
 	parser.add_argument("--input", "-i", help="Path to input directory containing the .tfrecord.gz files to pretrain with", required=True)
+	parser.add_argument("--heightmap", help="Optional. Filepath to the heightmap to pass as an input to the model along the channel dimension. If not specified, not heightmap will be specified as an input to the model. Default: None (not specified).")
 	parser.add_argument("--output", "-o", help="Path to output directory to write output to (will be automatically created if it doesn't exist)", required=True)
 	parser.add_argument("--batch-size", help="Sets the batch size [default: 64].", type=int)
-	parser.add_argument("--reads-multiplier", help="Optional. The multiplier for the number of files we should read from at once. Defaults to 1.5, which means read ceil(NUMBER_OF_CORES * 1.5) files at once. Set to a higher number of systems with high read latency to avoid starving the GPU of data.")
+	parser.add_argument("--reads-multiplier", help="Optional. The multiplier for the number of files we should read from at once. Defaults to 1.5, which means read ceil(NUMBER_OF_CORES * 1.5) files at once. Set to a higher number of systems with high read latency to avoid starving the GPU of data.", type=float)
 	parser.add_argument("--water-size", help="The width and height of the square of pixels that the model will predict. Smaller values crop the input more [default: 100].", type=int)
 	parser.add_argument("--water-threshold", help="The threshold at which a water cell should be considered water. Water depth values lower than this will be set to 0 (no water). Value unit is metres [default: 0.1].", type=int)
 	parser.add_argument("--bottleneck", help="The size of the bottleneck [default: 512].", type=int)
@@ -22,9 +23,9 @@ def parse_args():
 	parser.add_argument("--arch-dec", help="Next of the underlying decoder convnext model to use [default: convnext_i_xtiny].")
 	parser.add_argument("--learning-rate", help="The initial learning rate. YOU DO NOT USUALLY NEED TO CHANGE THIS. For experimental use only [default: determined automatically].", type=float)
 	
-	
 	return parser
 
+
 def run(args):
 	if (not hasattr(args, "water_size")) or args.water_size == None:
 		args.water_size = 100
@@ -46,6 +47,9 @@ def run(args):
 		args.arch_dec = "convnext_i_xtiny"
 	if (not hasattr(args, "learning_rate")) or args.learning_rate == None:
 		args.learning_rate = None
+	if (not hasattr(args, "heightmap")) or args.heightmap == None:
+		args.heightmap = None
+	
 	
 	
 	# TODO: Validate args here.
@@ -57,7 +61,8 @@ def run(args):
 		dirpath_input=args.input,
 		batch_size=args.batch_size,
 		water_threshold=args.water_threshold,
-		shape_water_desired=[args.water_size, args.water_size]
+		shape_water_desired=[args.water_size, args.water_size],
+		filepath_heightmap=args.heightmap
 	)
 	dataset_metadata = read_metadata(args.input)