TFUNetHelper Module

Utilities and helper functions for building and working with UNet-style models in TensorFlow.

HMB.TFUNetHelper.PreparePredTensorToNumpy(predTensor, doScale2Image=False)

Utility to convert model output tensor after the sigmoid/softmax activation to a numpy array of class indices. It can be used also with the original mask tensor if it is already in the correct format, as it handles squeezing and type conversion.

Short summary:

Takes the raw output tensor from a TensorFlow model (after activation) and processes it to produce a 2D or 3D numpy array of class indices or binary masks. This involves squeezing unnecessary dimensions, converting boolean masks to integers if needed, and ensuring the final output is in the correct format for evaluation or visualization.

Parameters:

• predTensor (tensorflow.Tensor) – The raw output tensor from the model after activation, expected to be of shape [B, H, W, C] or [B, H, W, 1].

• doScale2Image (bool) – If True, applies a threshold to convert probabilities to binary mask and scales to 0..255. Default False.

Returns:

Numpy array of shape [B, H, W] or [B, H, W, C] containing class indices or the scaled image.

Return type:

numpy.ndarray

class HMB.TFUNetHelper.DoubleConv(inChannels, outChannels)

Bases: Layer

Double convolution block used throughout the U-Net encoder and decoder. Two consecutive 3x3 Conv2D layers each followed by normalization and ReLU activation used as a basic encoder/decoder building block.

Initializes a double convolution block with two Conv2D layers, each followed by BatchNormalization and ReLU activation.

Parameters:

• inChannels (int) – Number of input channels.

• outChannels (int) – Number of output channels.

__init__(inChannels, outChannels)

Initializes a double convolution block with two Conv2D layers, each followed by BatchNormalization and ReLU activation.

Parameters:

• inChannels (int) – Number of input channels.

• outChannels (int) – Number of output channels.

call(x, training=False)

Forward pass through the double convolution block.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, H, W, inChannels].

• training (bool) – Whether the model is in training mode (affects BatchNormalization).

Returns:

Output feature map of shape [B, H, W, outChannels].

Return type:

tensorflow.Tensor

class HMB.TFUNetHelper.ConfigConv(inChannels, outChannels, norm='batch', dropout=0.0, residual=False)

Bases: Layer

Configurable double-convolution block with selectable normalization, dropout and residual option. A DoubleConv-like block with optional normalization type, dropout, and an optional residual connection when input and output channels match.

Initializes a configurable convolutional block with two Conv2D layers, optional normalization, dropout, and residual connection.

Parameters:

• inChannels (int) – Number of input channels.

• outChannels (int) – Number of output channels.

• norm (str) – Normalization type: “batch” | “instance” | “none”.

• dropout (float) – Dropout probability applied after the block.

• residual (bool) – Whether to add a residual skip when shapes permit.

__init__(inChannels, outChannels, norm='batch', dropout=0.0, residual=False)

Initializes a configurable convolutional block with two Conv2D layers, optional normalization, dropout, and residual connection.

Parameters:

• inChannels (int) – Number of input channels.

• outChannels (int) – Number of output channels.

• norm (str) – Normalization type: “batch” | “instance” | “none”.

• dropout (float) – Dropout probability applied after the block.

• residual (bool) – Whether to add a residual skip when shapes permit.

call(x, training=False)

Forward pass through the configurable convolutional block.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, H, W, inChannels].

• training (bool) – Whether the model is in training mode (affects normalization and dropout).

Returns:

Output tensor of shape [B, H, W, outChannels], optionally with residual addition.

Return type:

tensorflow.Tensor

class HMB.TFUNetHelper.ResidualBlock(inChannels, outChannels)

Bases: Layer

Residual convolutional block. Two conv->BatchNorm->ReLU layers with an identity skip connection. A 1x1 projection is applied to the identity when the number of channels differs.

Initializes a residual convolutional block with two Conv2D layers, each followed by BatchNormalization and ReLU activation.

Parameters:

• inChannels (int) – Number of input channels.

• outChannels (int) – Number of output channels.

__init__(inChannels, outChannels)

Initializes a residual convolutional block with two Conv2D layers, each followed by BatchNormalization and ReLU activation.

Parameters:

• inChannels (int) – Number of input channels.

• outChannels (int) – Number of output channels.

call(x, training=False)

Forward pass through the residual convolutional block.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, H, W, inChannels].

• training (bool) – Whether the model is in training mode (affects BatchNormalization).

Returns:

Output tensor of shape [B, H, W, outChannels] after residual addition and activation.

Return type:

tensorflow.Tensor

class HMB.TFUNetHelper.AttentionGate(F_g, F_l, F_int)

Bases: Layer

Attention gating module for skip connections. Lightweight gating unit that computes attention coefficients for encoder skip features using a gating signal from the decoder, improving focus on relevant spatial locations.

Initializes an attention gate with convolutional layers to compute attention coefficients.

Parameters:

• F_g (int) – Number of channels in the gating signal from the decoder.

• F_l (int) – Number of channels in the skip connection from the encoder.

• F_int (int) – Number of intermediate channels used for attention computation.

__init__(F_g, F_l, F_int)

Initializes an attention gate with convolutional layers to compute attention coefficients.

Parameters:

• F_g (int) – Number of channels in the gating signal from the decoder.

• F_l (int) – Number of channels in the skip connection from the encoder.

• F_int (int) – Number of intermediate channels used for attention computation.

call(g, x, training=False)

Forward pass through the attention gate.

Parameters:

• g (tensorflow.Tensor) – Gating signal from the decoder of shape [B, H, W, F_g].

• x (tensorflow.Tensor) – Skip connection from the encoder of shape [B, H, W, F_l].

• training (bool) – Whether the model is in training mode (affects BatchNormalization).

Returns:

Reweighted skip features of shape [B, H, W, F_l] after applying attention coefficients.

Return type:

tensorflow.Tensor

class HMB.TFUNetHelper.DepthwiseSeparableConv(inChannels, outChannels)

Bases: Layer

Depthwise separable convolution block. Implements a depthwise convolution followed by a pointwise convolution, used to reduce parameter counts and computation in lightweight models.

Initializes a depthwise separable convolution block with a SeparableConv2D layer, BatchNormalization, and ReLU activation.

Parameters:

• inChannels (int) – Number of input channels.

• outChannels (int) – Number of output channels.

__init__(inChannels, outChannels)

Initializes a depthwise separable convolution block with a SeparableConv2D layer, BatchNormalization, and ReLU activation.

Parameters:

• inChannels (int) – Number of input channels.

• outChannels (int) – Number of output channels.

call(x, training=False)

Forward pass through the depthwise separable convolution block.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, H, W, inChannels].

• training (bool) – Whether the model is in training mode (affects BatchNormalization).

class HMB.TFUNetHelper.SEBlock(channels, reduction=16)

Bases: Layer

Squeeze-and-Excitation (SE) block. Performs global channel-wise pooling followed by a small bottleneck MLP that produces per-channel scaling weights applied to the input.

Initializes a Squeeze-and-Excitation block with global average pooling and two dense layers for channel recalibration.

Parameters:

• channels (int) – Number of input/output channels.

• reduction (int) – Reduction ratio for the bottleneck MLP. Default 16.

__init__(channels, reduction=16)

Initializes a Squeeze-and-Excitation block with global average pooling and two dense layers for channel recalibration.

Parameters:

• channels (int) – Number of input/output channels.

• reduction (int) – Reduction ratio for the bottleneck MLP. Default 16.

call(x)

Forward pass through the Squeeze-and-Excitation block.

Parameters:

x (tensorflow.Tensor) – Input tensor of shape [B, H, W, channels].

Returns:

Output tensor of shape [B, H, W, channels] after channel-wise recalibration.

Return type:

tensorflow.Tensor

class HMB.TFUNetHelper.MultiResBlock(inChans, outChans, alpha=1.67)

Bases: Layer

Multi-resolution convolution block capturing features at multiple receptive fields. Parallel convolutions with different receptive fields followed by channel weighting to capture multi-scale context within a single block while controlling parameter growth.

Initializes a MultiRes block with three parallel convolutional paths (3x3, 5x5, 7x7) and a shortcut connection.

Parameters:

• inChans (int) – Number of input channels.

• outChans (int) – Total number of output channels across all paths.

• alpha (float) – Scaling factor to determine channel allocation for each path. Default 1.67.

__init__(inChans, outChans, alpha=1.67)

Initializes a MultiRes block with three parallel convolutional paths (3x3, 5x5, 7x7) and a shortcut connection.

Parameters:

• inChans (int) – Number of input channels.

• outChans (int) – Total number of output channels across all paths.

• alpha (float) – Scaling factor to determine channel allocation for each path. Default 1.67.

call(x, training=False)

Forward pass through the MultiRes block.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, H, W, inChans].

• training (bool) – Whether the model is in training mode (affects BatchNormalization).

Returns:

Output tensor of shape [B, H, W, outChans].

Return type:

tensorflow.Tensor

class HMB.TFUNetHelper.DenseBlock(inChans, numLayers=4, growthRate=32, bnSize=4)

Bases: Layer

Dense connectivity block with bottleneck layers for parameter efficiency. Stacks multiple bottleneck layers where each layer receives feature maps from all preceding layers as input, promoting feature reuse and gradient flow.

Initializes a Dense block with a specified number of bottleneck layers, growth rate, and bottleneck size.

Parameters:

• inChans (int) – Number of input channels.

• numLayers (int) – Number of bottleneck layers in the block. Default 4.

• growthRate (int) – Number of channels added per layer. Default 32.

• bnSize (int) – Bottleneck expansion factor for intermediate channels. Default 4.

__init__(inChans, numLayers=4, growthRate=32, bnSize=4)

Initializes a Dense block with a specified number of bottleneck layers, growth rate, and bottleneck size.

Parameters:

• inChans (int) – Number of input channels.

• numLayers (int) – Number of bottleneck layers in the block. Default 4.

• growthRate (int) – Number of channels added per layer. Default 32.

• bnSize (int) – Bottleneck expansion factor for intermediate channels. Default 4.

call(x, training=False)

Forward pass through the Dense block, progressively concatenating features from all previous layers.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, H, W, inChans].

• training (bool) – Whether the model is in training mode (affects BatchNormalization).

Returns:

Output tensor of shape [B, H, W, inChans + numLayers * growthRate].

Return type:

tensorflow.Tensor

class HMB.TFUNetHelper.RecurrentConvLayer(channels, t=2)

Bases: Layer

Recurrent convolutional layer with internal state feedback. Applies convolution repeatedly for T timesteps where each step receives feedback from its previous output, enabling iterative refinement of spatial features.

Initializes a recurrent convolutional layer with shared convolutional weights and batch normalization.

Parameters:

• channels (int) – Number of input/output channels.

• t (int) – Number of recurrent iterations. Default 2.

__init__(channels, t=2)

Initializes a recurrent convolutional layer with shared convolutional weights and batch normalization.

Parameters:

• channels (int) – Number of input/output channels.

• t (int) – Number of recurrent iterations. Default 2.

call(x, training=False)

Forward pass through the recurrent convolutional layer, iteratively refining features over T timesteps.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, H, W, channels].

• training (bool) – Whether the model is in training mode (affects BatchNormalization).

Returns:

Output tensor of shape [B, H, W, channels] after T recurrent iterations.

Return type:

tensorflow.Tensor

class HMB.TFUNetHelper.ASPP(inChans, outChans, dilations=(1, 6, 12, 18))

Bases: Layer

Atrous Spatial Pyramid Pooling for multi-scale context aggregation. Parallel dilated convolutions at multiple rates plus image pooling to capture objects at different scales within a single feature map.

Initializes the ASPP module with parallel convolutional branches and image pooling.

Parameters:

• inChans (int) – Number of input channels.

• outChans (int) – Number of output channels after concatenation and projection.

• dilations (Tuple[int]) – Dilation rates for the 3x3 convolution branches. Default (1, 6, 12, 18).

__init__(inChans, outChans, dilations=(1, 6, 12, 18))

Initializes the ASPP module with parallel convolutional branches and image pooling.

Parameters:

• inChans (int) – Number of input channels.

• outChans (int) – Number of output channels after concatenation and projection.

• dilations (Tuple[int]) – Dilation rates for the 3x3 convolution branches. Default (1, 6, 12, 18).

call(x, training=False)

Forward pass through the ASPP module, aggregating multi-scale features and projecting to output channels.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, H, W, inChans].

• training (bool) – Whether the model is in training mode (affects BatchNormalization and Dropout).

Returns:

Output tensor of shape [B, H, W, outChans] after multi-scale feature aggregation.

Return type:

tensorflow.Tensor

class HMB.TFUNetHelper.PatchEmbedding(inChans, embedDim=256, patchSize=4)

Bases: Layer

2D image to patch embedding with optional positional encoding. Converts input image into non-overlapping patches via convolutional projection and flattens them into a sequence for transformer processing.

Initializes a patch embedding layer that projects an input image into a sequence of patch embeddings using a Conv2D layer.

Parameters:

• inChans (int) – Number of input channels.

• embedDim (int) – Dimension of the embedding for each patch. Default 256.

• patchSize (int) – Size of each patch (height and width). Default 4.

__init__(inChans, embedDim=256, patchSize=4)

Initializes a patch embedding layer that projects an input image into a sequence of patch embeddings using a Conv2D layer.

Parameters:

• inChans (int) – Number of input channels.

• embedDim (int) – Dimension of the embedding for each patch. Default 256.

• patchSize (int) – Size of each patch (height and width). Default 4.

call(x)

Forward pass through the patch embedding layer, projecting the input image into a sequence of flattened patch embeddings.

Parameters:

x (tensorflow.Tensor) – Input tensor of shape [B, H, W, inChans].

Returns:

Patch sequence of shape [B, N, embedDim] and tuple (patchH, patchW) representing the grid size of patches.

Return type:

Tuple[tensorflow.Tensor, Tuple[int,int]]

class HMB.TFUNetHelper.TransformerBlock(embedDim, numHeads=8, mlpRatio=4.0, dropout=0.1)

Bases: Layer

Standard transformer encoder block with pre-normalization. Multi-head self-attention followed by position-wise feed-forward network with layer normalization and residual connections before each sub-layer.

Initializes a transformer encoder block with multi-head self-attention and feed-forward MLP.

Parameters:

• embedDim (int) – Dimension of the input embeddings.

• numHeads (int) – Number of attention heads. Default 8.

• mlpRatio (float) – Ratio to determine the hidden dimension of the MLP relative to embedDim. Default 4.0.

• dropout (float) – Dropout rate applied after attention and MLP layers. Default 0.1.

__init__(embedDim, numHeads=8, mlpRatio=4.0, dropout=0.1)

Initializes a transformer encoder block with multi-head self-attention and feed-forward MLP.

Parameters:

• embedDim (int) – Dimension of the input embeddings.

• numHeads (int) – Number of attention heads. Default 8.

• mlpRatio (float) – Ratio to determine the hidden dimension of the MLP relative to embedDim. Default 4.0.

• dropout (float) – Dropout rate applied after attention and MLP layers. Default 0.1.

call(x, training=False)

Forward pass through the transformer encoder block, applying pre-normalization, multi-head self-attention, and feed-forward MLP with residual connections.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, N, embedDim].

• training (bool) – Whether the model is in training mode (affects Dropout).

Returns:

Output tensor of shape [B, N, embedDim] after attention and MLP processing.

Return type:

tensorflow.Tensor

class HMB.TFUNetHelper.ChannelAttn(channels, reduction=16)

Bases: Layer

Channel attention branch of CBAM using squeeze-and-excitation. Global average and max pooling followed by shared MLP to compute channel-wise attention weights that rescale feature channels.

Initializes a channel attention block with global pooling and a two-layer MLP for channel recalibration.

Parameters:

• channels (int) – Number of input/output channels.

• reduction (int) – Reduction ratio for the bottleneck MLP. Default 16.

__init__(channels, reduction=16)

Initializes a channel attention block with global pooling and a two-layer MLP for channel recalibration.

Parameters:

• channels (int) – Number of input/output channels.

• reduction (int) – Reduction ratio for the bottleneck MLP. Default 16.

call(x)

Forward pass through the channel attention block, computing channel-wise attention weights and applying them to the input feature map.

Parameters:

x (tensorflow.Tensor) – Input tensor of shape [B, H, W, channels].

class HMB.TFUNetHelper.SpatialAttn(kernelSize=7)

Bases: Layer

Spatial attention branch of CBAM using channel aggregation. Channel-wise average and max pooling followed by convolution to compute spatial attention map that highlights important regions in feature maps.

Initializes a spatial attention block with a convolutional layer to compute spatial attention weights.

Parameters:

kernelSize (int) – Size of the convolution kernel for spatial attention. Default 7.

__init__(kernelSize=7)

Initializes a spatial attention block with a convolutional layer to compute spatial attention weights.

Parameters:

kernelSize (int) – Size of the convolution kernel for spatial attention. Default 7.

call(x)

Forward pass through the spatial attention block, computing a spatial attention map and applying it to the input feature map.

Parameters:

x (tensorflow.Tensor) – Input tensor of shape [B, H, W, channels].

class HMB.TFUNetHelper.CBAM(channels, reduction=16)

Bases: Layer

Convolutional Block Attention Module (channel + spatial attention). Sequential channel attention followed by spatial attention to adaptively refine feature maps along both channel and spatial dimensions.

Initializes a CBAM module with channel and spatial attention branches.

Parameters:

• channels (int) – Number of input/output channels.

• reduction (int) – Reduction ratio for the channel attention MLP. Default 16.

__init__(channels, reduction=16)

Initializes a CBAM module with channel and spatial attention branches.

Parameters:

• channels (int) – Number of input/output channels.

• reduction (int) – Reduction ratio for the channel attention MLP. Default 16.

call(x)

Forward pass through the CBAM module, applying channel attention followed by spatial attention.

Parameters:

x (tensorflow.Tensor) – Input tensor of shape [B, H, W, channels].

Returns:

Output tensor of shape [B, H, W, channels] after attention refinement.

Return type:

tensorflow.Tensor

class HMB.TFUNetHelper.UNet(inputChannels=3, numClasses=2, baseChannels=64, useConvTranspose2d=True)

Bases: Model

Standard 2D U-Net implemented with tf.keras. A typical encoder-decoder U-Net with four downsampling stages, a bottleneck, and four symmetric upsampling stages. Supports learned transpose convolution upsampling or bilinear upsampling followed by 1x1 projection.

Parameters:

• inputChannels (int) – Number of input image channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• baseChannels (int) – Filters in the first stage. Default 64.

• useConvTranspose2d (bool) – Use Conv2DTranspose when True; otherwise use bilinear upsample.

Variables:

• center (enc1..enc4,) – Encoder and bottleneck blocks.

• dec1..dec4 (up1..up4,) – Upsampling and decoder blocks.

• finalConv (tf.keras.layers.Layer) – 1x1 conv to map to logits.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

call(x, training=False)

class HMB.TFUNetHelper.DynamicUNet(inputChannels=3, numClasses=2, baseChannels=64, depth=4, upMode='transpose', norm='batch', dropout=0.0, residual=False)

Bases: Model

Configurable U-Net implementation with variable depth.

Short summary:

Generalized U-Net that allows customizing depth, normalization type, dropout, residual usage and upsampling mode. Useful for experimenting with architectures without duplicating code.

Parameters:

• inputChannels (int) – Number of input channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• baseChannels (int) – Base filter count. Default 64.

• depth (int) – Number of downsampling stages. Default 4.

• upMode (str) – “transpose” or “bilinear” upsampling. Default “transpose”.

• norm (str) – Normalization type: “batch”|”instance”|”none”. Default “batch”.

• dropout (float) – Dropout probability applied inside config blocks. Default 0.0.

• residual (bool) – Whether to use residual connections inside blocks. Default False.

Variables:

• encs (list) – Encoder blocks.

• pools (list) – Pooling layers.

• ups (list) – Upsampling layers.

• decs (list) – Decoder blocks.

• finalConv (tf.keras.layers.Layer) – 1x1 conv to logits.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

call(x, training=False)

class HMB.TFUNetHelper.AttentionUNet(*args, **kwargs)

Bases: Model

Attention U-Net variant with gating on skip connections.

Short summary:

Applies attention gating to encoder skip features before merging them into the decoder, improving localization and suppressing irrelevant responses in the skip maps.

Parameters:

• inputChannels (int) – Number of input channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• baseChannels (int) – Base filter count. Default 64.

• useConvTranspose2d (bool) – Use Conv2DTranspose for upsampling when True.

Variables:

• center (enc1..enc4,) – Encoder/bottleneck blocks.

• att1..att4 (AttentionGate) – Attention gating modules for skips.

• dec1..dec4 (up1..up4,) – Decoder modules.

• finalConv (tf.keras.layers.Layer) – 1x1 conv mapping to logits.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

call(x, training=False)

class HMB.TFUNetHelper.MobileUNet(*args, **kwargs)

Bases: Model

Lightweight U-Net using depthwise separable convolutions.

Short summary:

Efficient U-Net variant that replaces standard convolutions with depthwise separable convolutions to reduce parameters and compute.

Parameters:

• inputChannels (int) – Number of input channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• baseChannels (int) – Base filter count. Default 32.

• useConvTranspose2d (bool) – Use Conv2DTranspose for upsampling when True.

Variables:

• center (enc1..enc4,) – Encoder and bottleneck sequences.

• dec1..dec4 (up1..up4,) – Decoder modules.

• finalConv (tensorflow.keras.layers.Layer) – 1x1 conv for logits.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

call(x, training=False)

class HMB.TFUNetHelper.ResidualUNet(*args, **kwargs)

Bases: Model

Residual U-Net variant built from ResidualBlock components.

Short summary:

A U-Net where encoder and decoder stages use residual blocks to ease optimization and improve gradient flow for deeper models.

Parameters:

• inputChannels (int) – Number of input channels.

• numClasses (int) – Number of output classes.

• baseChannels (int) – Base number of filters.

• useConvTranspose2d (bool) – Use Conv2DTranspose when True.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

call(x, training=False)

class HMB.TFUNetHelper.SEUNet(*args, **kwargs)

Bases: Model

U-Net with Squeeze-and-Excitation modules applied to encoder and center.

Short summary:

Inserts SE blocks after encoder and center blocks to recalibrate channel-wise responses, improving representational expressiveness.

Parameters:

• inputChannels (int) – Number of input channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• baseChannels (int) – Base number of filters. Default 64.

• useConvTranspose2d (bool) – If True use learned transpose upsampling.

• seReduction (int) – Reduction ratio for SE bottleneck. Default 16.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

call(x, training=False)

class HMB.TFUNetHelper.ResidualAttentionUNet(*args, **kwargs)

Bases: Model

Residual U-Net combined with Attention Gates.

Short summary:

A U-Net variant that composes residual blocks in the encoder/decoder and applies attention gating on the skip connections to focus decoder features.

Parameters:

• inputChannels (int) – Number of input channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• baseChannels (int) – Base number of filters. Default 64.

• useConvTranspose2d (bool) – Use Conv2DTranspose for learned upsampling when True.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

call(x, training=False)

class HMB.TFUNetHelper.MultiResUNet(*args, **kwargs)

Bases: Model

MultiResUNet using MultiResBlock at each stage.

Short summary:

Uses MultiResBlock modules to capture multi-scale features within each encoder/decoder stage while keeping parameter growth under control.

Parameters:

• inputChannels (int) – Number of input channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• baseChannels (int) – Base number of filters. Default 64.

• useConvTranspose2d (bool) – Use Conv2DTranspose when True.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

call(x, training=False)

class HMB.TFUNetHelper.DenseUNet(*args, **kwargs)

Bases: Model

DenseUNet integrating DenseBlocks and transition convolutions.

Short summary:

Encoder built from DenseBlocks with 1x1 transition convolutions to control channel dimensionality, producing feature reuse and improved gradient flow.

Parameters:

• inputChannels (int) – Number of input channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• baseChannels (int) – Base filter count. Default 32.

• useConvTranspose2d (bool) – Use Conv2DTranspose when True.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

call(x, training=False)

class HMB.TFUNetHelper.R2UNet(*args, **kwargs)

Bases: Model

Recurrent Residual U-Net using recurrent convolution layers.

Short summary:

Introduces recurrent convolutional layers (RCL) that iteratively refine features inside each encoder stage, useful for capturing contextual information while keeping model size modest.

Parameters:

• inputChannels (int) – Number of input channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• baseChannels (int) – Base number of filters. Default 64.

• useConvTranspose2d (bool) – Use Conv2DTranspose when True.

• t (int) – Number of recurrent iterations inside RCL. Default 2.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

call(x, training=False)

class HMB.TFUNetHelper.ASPPUNet(*args, **kwargs)

Bases: Model

U-Net with ASPP at the bottleneck.

Short summary:

Incorporates Atrous Spatial Pyramid Pooling (ASPP) at the bottleneck to capture multi-scale context before upsampling in the decoder.

Parameters:

• inputChannels (int) – Number of input channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• baseChannels (int) – Base number of filters. Default 64.

• useConvTranspose2d (bool) – Use Conv2DTranspose when True.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

call(x, training=False)

class HMB.TFUNetHelper.TransUNet(*args, **kwargs)

Bases: Model

Transformer-enhanced U-Net with a small transformer stack at the bottleneck.

Short summary:

Applies a patch embedding on the bottleneck features and processes them with a small transformer encoder stack, then projects back to spatial features for decoding. Useful to capture long-range dependencies.

Parameters:

• inputChannels (int) – Number of input channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• baseChannels (int) – Base number of filters. Default 64.

• embedDim (int) – Embedding dimension for transformer. Default 256.

• numHeads (int) – Number of attention heads. Default 8.

• numEncoders (int) – Number of transformer encoder blocks. Default 2.

• useConvTranspose2d (bool) – Use Conv2DTranspose when True.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

call(x, training=False)

class HMB.TFUNetHelper.CBAMUNet(*args, **kwargs)

Bases: Model

U-Net variant applying CBAM attention to encoder skip features.

Short summary:

Uses the Convolutional Block Attention Module (CBAM) to refine encoder skip features via channel and spatial attention before merging them into the decoder.

Parameters:

• inputChannels (int) – Number of input channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• baseChannels (int) – Base number of filters. Default 64.

• useConvTranspose2d (bool) – Use Conv2DTranspose when True.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

call(x, training=False)

class HMB.TFUNetHelper.EfficientUNet(*args, **kwargs)

Bases: Model

Efficient wrapper around MobileUNet.

Short summary:

A thin wrapper that exposes the MobileUNet as an “efficient” option in the factory while keeping the same API.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

call(x, training=False)

class HMB.TFUNetHelper.BoundaryAwareUNet(*args, **kwargs)

Bases: Model

Boundary-aware U-Net with explicit boundary detection branch.

Short summary:

A dual-branch architecture with a main segmentation decoder and a lighter boundary decoder producing an auxiliary boundary map. Useful for losses that combine segmentation and boundary supervision.

Parameters:

• inputChannels (int) – Number of input channels. Default 3.

• numClasses (int) – Number of segmentation classes. Default 2.

• baseChannels (int) – Base number of filters. Default 64.

• useConvTranspose2d (bool) – Use Conv2DTranspose when True.

• boundaryWeight (float) – Weighting factor for combining boundary predictions. Default 0.5.

Returns:

(segmentation_logits, boundary_map)

Return type:

Tuple[tensorflow.Tensor, tensorflow.Tensor]

call(x, training=False)

class HMB.TFUNetHelper.VNet(inputSize=(256, 256, 1), kernelInitializer='he_normal', dropoutRatio=0.0, dropoutType='spatial', activation='relu', applyBatchNorm=False, concatenateType='concatenate', noOfLevels=5, numClasses=1)

Bases: Model

V-Net style residual encoder-decoder implemented as a lazily-built Model subclass.

This class preserves the original functional builder behaviour but exposes it as a tf.keras.Model subclass so it matches the other model classes in this module. The internal functional model is constructed on first call using the same building logic as the previous functional VNet builder.

Parameters:

• inputSize (tuple) – Input size (H, W, C). Supports None for dynamic spatial dimensions. Default (256, 256, 1).

• kernelInitializer (str) – Kernel initializer for convolutional layers. Default “he_normal”.

• dropoutRatio (float) – Dropout ratio for residual blocks. Default 0.0 (no dropout).

• dropoutType (str) – Type of dropout to apply (“spatial” or “standard”). Default “spatial”.

• activation (str) – Activation function to use. Default “relu”.

• applyBatchNorm (bool) – Whether to apply batch normalization after convolutions. Default False.

• concatenateType (str) – Method to merge skip connections (“concatenate” or “add”). Default “concatenate”.

• noOfLevels (int) – Number of encoder/decoder levels. Default 5.

• numClasses (int) – Number of output classes for the final segmentation head. Default 1.

Returns:

Logits tensor of shape [B, H, W, numClasses].

Return type:

tensorflow.Tensor

build(input_shape)

call(x, training=False)

class HMB.TFUNetHelper.SegNet(*args, **kwargs)

Bases: Model

SegNet architecture for multi-class semantic segmentation (Class-based Implementation). Memory-efficient encoder-decoder network supporting VGG16, ResNet50, MobileNetV2, or Vanilla backbones. Uses Functional API internally for optimal graph construction and memory usage.

Initialize SegNet with specified configuration.

Parameters:

• inputChannels (int) – Number of input image channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• level (int) – Decoder starting stage (1-4). Higher levels use deeper features.

• encoder (str) – Backbone type: “VGG16”, “ResNet50”, “MobileNet”, or “Vanilla”.

• inputSize (tuple, Optional) – Fixed input shape (H, W, C). None = dynamic.

• useBias (bool) – Whether Conv layers use bias. Default False (BN compatible).

__init__(inputChannels=3, numClasses=2, level=3, encoder='VGG16', inputSize=None, useBias=False)

Initialize SegNet with specified configuration.

Parameters:

• inputChannels (int) – Number of input image channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• level (int) – Decoder starting stage (1-4). Higher levels use deeper features.

• encoder (str) – Backbone type: “VGG16”, “ResNet50”, “MobileNet”, or “Vanilla”.

• inputSize (tuple, Optional) – Fixed input shape (H, W, C). None = dynamic.

• useBias (bool) – Whether Conv layers use bias. Default False (BN compatible).

BuildEncoderLayers()

Initialize encoder backbone layers based on selected type.

Raises:

ValueError – If an unsupported encoder type is specified.

InitVggEncoder()

VGG16-style encoder: 4 blocks returning feature levels at strides 2,4,8,16.

InitResnet50Encoder()

ResNet50 encoder using keras.applications for memory efficiency.

InitMobilenetEncoder()

MobileNetV2 encoder using keras.applications for memory efficiency.

InitVanillaEncoder()

Vanilla encoder with correct channel progression: in→64→128→256→512.

BuildDecoderLayers()

Build decoder layers with unique names to prevent duplication errors.

BuildOutputHead()

Build final classification head.

RunVggEncoder(x, training=False)

Execute VGG encoder and return 4 feature levels.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, H, W, C].

• training (bool) – Whether the model is in training mode (affects BatchNorm). Default False.

Returns:

List of encoder feature levels for decoder input.

Return type:

list of tensorflow.Tensor

RunVanillaEncoder(x, training=False)

Execute Vanilla encoder and return 4 feature levels.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, H, W, C].

• training (bool) – Whether the model is in training mode (affects BatchNorm). Default False.

Returns:

List of encoder feature levels for decoder input.

Return type:

list of tensorflow.Tensor

RunResnetEncoder(x, training=False)

Execute ResNet50 encoder and return 4 feature levels.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, H, W, C].

• training (bool) – Whether the model is in training mode (affects BatchNorm). Default False.

Returns:

List of encoder feature levels for decoder input.

Return type:

list of tensorflow.Tensor

RunMobilenetEncoder(x, training=False)

Execute MobileNetV2 encoder and return 4 feature levels.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, H, W, C].

• training (bool) – Whether the model is in training mode (affects BatchNorm). Default False.

Returns:

List of encoder feature levels for decoder input.

Return type:

list of tensorflow.Tensor

RunEncoder(x, training=False)

Route to appropriate encoder implementation.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, H, W, C].

• training (bool) – Whether the model is in training mode (affects BatchNorm). Default False.

Returns:

List of encoder feature levels for decoder input.

Return type:

list of tensorflow.Tensor

RunDecoder(x, training=False)

Execute decoder pathway from selected level.

Parameters:

• x (tensorflow.Tensor) – Input tensor from encoder level.

• training (bool) – Whether the model is in training mode (affects BatchNorm). Default False.

Returns:

Output tensor after decoder processing.

Return type:

tensorflow.Tensor

r

BuildFunctionalModel()

Build static functional model for fixed input shapes (optimization).

call(x, training=False)

Forward pass: route to functional model if fixed shape, else dynamic.

Parameters:

• x (tensorflow.Tensor) – Input tensor of shape [B, H, W, C].

• training (bool) – Whether the model is in training mode (affects BatchNorm and Dropout). Default False.

Returns:

Output tensor of shape [B, H, W, numClasses] with probabilities or logits depending on configuration.

Return type:

tensorflow.Tensor

build(input_shape=None)

Explicitly build model weights (useful for dynamic shapes).

Parameters:

input_shape (tuple, Optional) – Input shape to build the model. If None, a dynamic shape with known channels is used.

get_config()

Enable model serialization.

Returns:

Configuration dictionary containing model parameters for reproduction.

Return type:

dict

classmethod from_config(config)

Enable model deserialization.

Parameters:

• cls – The class to instantiate (SegNet).

• config (dict) – Configuration dictionary containing model parameters.

Returns:

An instance of SegNet reconstructed from the provided configuration.

Return type:

SegNet

HMB.TFUNetHelper.CreateUNet(inputChannels=3, numClasses=2, baseChannels=64, depth=4, upMode='transpose', norm='batch', dropout=0.0, residual=False, modelType='dynamic')

Extended factory that supports the UNet family and related variants.

Short summary:

Returns an instantiated tf.keras.Model implementing the requested UNet variant selected by the case-insensitive modelType string. This factory exposes the common constructor arguments used by the various implementations and forwards them to the selected model. Models in this file follow TensorFlow’s NHWC convention (batch, height, width, channels).

Parameters:

• inputChannels (int) – Number of input image channels. Default 3.

• numClasses (int) – Number of output classes. Default 2.

• baseChannels (int) – Base filter count used by most architectures. Default 64.

• depth (int) – Depth / number of downsampling stages for dynamic variants. Default 4.

• upMode (str) – “transpose” or “bilinear” upsampling selection. Default “transpose”.

• norm (str) – Normalization mode for configurable blocks: “batch” | “instance” | “none”. Default “batch”.

• dropout (float) – Dropout probability passed to configurable blocks. Default 0.0.

• residual (bool) – Whether to enable residual connections where supported. Default False.

• modelType (str) – Case-insensitive model selection string (e.g. “dynamic”, “UNet”, “ASPPUNet”).

Returns:

Instantiated model ready for training or inference. The returned model accepts NHWC inputs and produces logits in NHWC format (shape [B, H, W, C]).

Return type:

tensorflow.keras.Model

Notes

• If an unrecognized modelType is provided the function falls back to the DynamicUNet configuration. Use the names exposed in AVAILABLE_UNETS or the CamelCase class names when calling GetUNetModel.

HMB.TFUNetHelper.GetUNetModel(modelName, inputChannels, numClasses, baseChannels=64)

Factory helper to instantiate a specific UNet variant by its CamelCase class name.

Short summary:

Convenience function that maps a model class name (exact CamelCase string) to the corresponding tf.keras.Model constructor in this module and returns an instantiated model. This is useful when model names are provided from configuration files or command-line arguments.

Parameters:

• modelName (str) – Exact CamelCase name of the desired model class (e.g. “UNet”, “TransUNet”).

• inputChannels (int) – Number of input channels to construct the model with.

• numClasses (int) – Number of output classes for the model.

• baseChannels (int) – Base channel count used by many model constructors. Default 64.

Returns:

Instantiated model corresponding to modelName.

Return type:

tensorflow.keras.Model

Raises:

ValueError – If modelName is not recognized among the known UNet class names in this module.