From eac6472c97ed033d5bfcb7cf6c6449b8bfe19691 Mon Sep 17 00:00:00 2001 From: Starbeamrainbowlabs Date: Tue, 6 Dec 2022 18:55:58 +0000 Subject: [PATCH] Implement support for (optionally) taking a heightmap in --- .../src/lib/ai/model_rainfallwater_mono.py | 8 +++++-- aimodel/src/lib/dataset/dataset_mono.py | 23 +++++++++++++++---- aimodel/src/lib/dataset/parse_heightmap.py | 12 ++++++++++ aimodel/src/subcommands/train_mono.py | 11 ++++++--- 4 files changed, 45 insertions(+), 9 deletions(-) create mode 100644 aimodel/src/lib/dataset/parse_heightmap.py diff --git a/aimodel/src/lib/ai/model_rainfallwater_mono.py b/aimodel/src/lib/ai/model_rainfallwater_mono.py index 45ba01d..32bde65 100644 --- a/aimodel/src/lib/ai/model_rainfallwater_mono.py +++ b/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( diff --git a/aimodel/src/lib/dataset/dataset_mono.py b/aimodel/src/lib/dataset/dataset_mono.py index 5683b43..1f2e67b 100644 --- a/aimodel/src/lib/dataset/dataset_mono.py +++ b/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) diff --git a/aimodel/src/lib/dataset/parse_heightmap.py b/aimodel/src/lib/dataset/parse_heightmap.py new file mode 100644 index 0000000..24861d9 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/aimodel/src/subcommands/train_mono.py b/aimodel/src/subcommands/train_mono.py index 276ccd9..98255d2 100644 --- a/aimodel/src/subcommands/train_mono.py +++ b/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)