sbrl-github/research-rainfallradar
1
0
Fork 0
mirror of https://github.com/sbrl/research-rainfallradar synced 2024-11-22 09:13:01 +00:00
Code Issues Projects Releases Packages Wiki Activity

ai: tweak the segmentation model structure

Browse source
This commit is contained in:
Starbeamrainbowlabs 2022-09-15 19:54:50 +01:00
parent 1bc8a5bf13
commit e3c8277255
Signed by: sbrl
GPG key ID: 1BE5172E637709C2
1 changed files with 6 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -6,6 +6,8 @@ import tensorflow as tf
from .components.convnext_inverse import do_convnext_inverse from .components.convnext_inverse import do_convnext_inverse
def model_rainfallwater_segmentation(metadata, feature_dim_in, shape_water_out, batch_size=64, summary_file=None): def model_rainfallwater_segmentation(metadata, feature_dim_in, shape_water_out, batch_size=64, summary_file=None):
out_water_width, out_water_height, out_water_channels = shape_water_out
layer_input = tf.keras.layers.Input( layer_input = tf.keras.layers.Input(
shape=(feature_dim_in) shape=(feature_dim_in)
@ -18,6 +20,10 @@ def model_rainfallwater_segmentation(metadata, feature_dim_in, shape_water_out,
layer_next = do_convnext_inverse(layer_next, arch_name="convnext_i_tiny") layer_next = do_convnext_inverse(layer_next, arch_name="convnext_i_tiny")
# TODO: An attention layer here instead of a dense layer, with a skip connection?
layer_next = tf.keras.layers.Dense(32)(layer_next)
layer_next = tf.keras.layers.Conv2D(out_water_channels, 7, activation="softmax", padding="same")(layer_next)
# TODO: Implement projection head here # TODO: Implement projection head here
model = tf.keras.Model( model = tf.keras.Model(