sbrl-github/research-rainfallradar
1
0
Fork 0
mirror of https://github.com/sbrl/research-rainfallradar synced 2024-06-30 18:44:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

model_segmentation: document; remove unused args

Browse source
This commit is contained in:
Starbeamrainbowlabs 2022-10-13 17:50:16 +01:00
parent f12e6ab905
commit 5c35c0cee4
Signed by: sbrl
GPG key ID: 1BE5172E637709C2
1 changed files with 6 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
aimodel/src/lib/ai/model_rainfallwater_segmentation.py
View file

@ -6,20 +6,18 @@ import tensorflow as tf
from .components.convnext_inverse import do_convnext_inverse
def model_rainfallwater_segmentation(metadata, shape_water_out, model_arch="convnext_i_xtiny", batch_size=64, summary_file=None, water_bins=2):
def model_rainfallwater_segmentation(metadata, shape_water_out, model_arch="convnext_i_xtiny", batch_size=64, water_bins=2):
"""Makes a new rainfall / waterdepth segmentation head model.
Args:
metadata (dict): A dictionary of metadata about the dataset to use to build the model with.
feature_dim_in (int): The size of the feature dimension
shape_water_out (_type_): _description_
model_arch (str, optional): _description_. Defaults to "convnext_i_xtiny".
batch_size (int, optional): _description_. Defaults to 64.
summary_file (_type_, optional): _description_. Defaults to None.
water_bins (int, optional): _description_. Defaults to 2.
shape_water_out (int[]): The width and height (in that order) that should dictate the output shape of the segmentation head. CURRENTLY NOT USED.
model_arch (str, optional): The architecture code for the underlying (inverted) ConvNeXt model. 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.
Returns:
_type_: _description_
tf.keras.Model: The new model, freshly compiled for your convenience! :D
"""
out_water_width, out_water_height = shape_water_out
feature_dim_in = metadata["rainfallradar"][0]