diff --git a/aimodel/src/lib/ai/model_rainfallwater_segmentation.py b/aimodel/src/lib/ai/model_rainfallwater_segmentation.py
index 35880f1..127605c 100644
--- a/aimodel/src/lib/ai/model_rainfallwater_segmentation.py
+++ b/aimodel/src/lib/ai/model_rainfallwater_segmentation.py
@@ -6,6 +6,8 @@ import tensorflow as tf
 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):
+	out_water_width, out_water_height, out_water_channels = shape_water_out
+	
 	
 	layer_input = tf.keras.layers.Input(
 		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")
 	
+	# 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
 	
 	model = tf.keras.Model(