From 575e0b3d283fd6d1266b275e2abb107b104a4044 Mon Sep 17 00:00:00 2001 From: Santiago Martinez Balvanera Date: Wed, 15 Jul 2026 16:55:29 +0100 Subject: [PATCH] Add vertical mean block --- src/batdetect2/models/blocks.py | 85 +++++++++++++++++++++++++++++ src/batdetect2/models/bottleneck.py | 85 +++++++++++++++++++++++++++-- uv.lock | 2 +- 3 files changed, 166 insertions(+), 6 deletions(-) diff --git a/src/batdetect2/models/blocks.py b/src/batdetect2/models/blocks.py index ccce61c..b3fe820 100644 --- a/src/batdetect2/models/blocks.py +++ b/src/batdetect2/models/blocks.py @@ -474,6 +474,91 @@ class VerticalConv(Block): ) +class VerticalMeanConfig(BaseConfig): + """Configuration for a ``VerticalMean`` block. + + Attributes + ---------- + name : str + Discriminator field; always ``"VerticalMean"``. + """ + + name: Literal["VerticalMean"] = "VerticalMean" + """Discriminator field indicating the block type.""" + + channels: int + """Number of output channels.""" + + +class VerticalMean(Block): + """Mean pooling block operating along the height dimension. + + Applies a 2D mean pooling operation, followed by a 2D convolution, + followed by a batch normalization and ReLU activation. + + Sequence: Mean Pool -> Conv -> BN -> ReLU. + + Parameters + ---------- + in_channels : int + Number of channels in the input tensor. + out_channels : int + Number of output channels after the mean pooling. + input_height : int + The height (H dimension) of the input tensor. The convolutional kernel + will be sized `(1, 1)`. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + input_height: int, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.input_height = input_height + self.conv = nn.Conv2d( + in_channels, + out_channels, + kernel_size=(1, 1), + padding=0, + ) + self.batch_norm = nn.BatchNorm2d(out_channels) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Apply avg pooling -> Conv -> BN -> ReLU. + + Parameters + ---------- + x : torch.Tensor + Input tensor, shape `(B, C_in, H, W)`. + + Returns + ------- + torch.Tensor + Output tensor, shape `(B, C_out, 1, W)`. + """ + x = x.mean(dim=2, keepdim=True) + x = self.conv(x) + return F.relu(self.batch_norm(x), inplace=True) + + @block_registry.register(VerticalMeanConfig) + @staticmethod + def from_config( + config: VerticalMeanConfig, + input_channels: int, + input_height: int, + ): + return VerticalMean( + in_channels=input_channels, + out_channels=config.channels, + input_height=input_height, + ) + + + class FreqCoordConvDownConfig(BaseConfig): """Configuration for a FreqCoordConvDownBlock.""" diff --git a/src/batdetect2/models/bottleneck.py b/src/batdetect2/models/bottleneck.py index 9b2154a..f148611 100644 --- a/src/batdetect2/models/bottleneck.py +++ b/src/batdetect2/models/bottleneck.py @@ -29,6 +29,9 @@ from batdetect2.models.blocks import ( Block, SelfAttentionConfig, VerticalConv, + VerticalConvConfig, + VerticalMean, + VerticalMeanConfig, build_layer, ) from batdetect2.models.types import BottleneckProtocol @@ -94,6 +97,7 @@ class Bottleneck(Block): in_channels: int, out_channels: int, bottleneck_channels: int | None = None, + frequency_aggregator: Block | None = None, layers: List[torch.nn.Module] | None = None, ) -> None: """Initialise the Bottleneck layer. @@ -125,11 +129,14 @@ class Bottleneck(Block): ) self.layers = nn.ModuleList(layers or []) - self.conv_vert = VerticalConv( - in_channels=in_channels, - out_channels=self.bottleneck_channels, - input_height=input_height, - ) + if frequency_aggregator is None: + frequency_aggregator = VerticalConv( + in_channels=in_channels, + out_channels=self.bottleneck_channels, + input_height=input_height, + ) + + self.conv_vert = frequency_aggregator def forward(self, x: torch.Tensor) -> torch.Tensor: """Process the encoder's bottleneck features. @@ -166,6 +173,13 @@ BottleneckLayerConfig = Annotated[ """Type alias for the discriminated union of block configs usable in the Bottleneck.""" +FrequencyAggregationLayerConfig = Annotated[ + (VerticalConvConfig | VerticalMeanConfig), + Field(discriminator="name"), +] +"""Type alias for the discriminated union of block configs usable in the FrequencyAggregation.""" + + class BottleneckConfig(BaseConfig): """Configuration for the bottleneck component. @@ -182,17 +196,71 @@ class BottleneckConfig(BaseConfig): """ channels: int + frequency_aggregation: FrequencyAggregationLayerConfig = Field( + default_factory=lambda: VerticalConvConfig(channels=256) + ) layers: List[BottleneckLayerConfig] = Field(default_factory=list) DEFAULT_BOTTLENECK_CONFIG: BottleneckConfig = BottleneckConfig( channels=256, + frequency_aggregation=VerticalConvConfig(channels=256), layers=[ SelfAttentionConfig(attention_channels=256), ], ) +def build_frequency_aggregation( + input_height: int, + in_channels: int, + config: FrequencyAggregationLayerConfig, +) -> Block: + """Build a block for aggregating frequency information. + + Parameters + ---------- + input_height : int + Height (number of frequency bins) of the input tensor from the + encoder. Must be positive. + in_channels : int + Number of channels in the input tensor from the encoder. Must be + positive. + config : FrequencyAggregationLayerConfig, optional + Configuration specifying the output channel count and any + additional layers. Uses ``VerticalConvConfig`` if ``None``. + + Returns + ------- + Block + An initialised ``VerticalConv`` module. + + Raises + ------ + AssertionError + If any configured layer changes the height of the feature map + (bottleneck layers must preserve height so that it can be restored + by repetition). + """ + if config.name == "VerticalConv": + return VerticalConv( + in_channels=in_channels, + out_channels=config.channels, + input_height=input_height, + ) + + if config.name == "VerticalMean": + return VerticalMean( + in_channels=in_channels, + out_channels=config.channels, + input_height=input_height, + ) + + raise NotImplementedError( + f"Unknown frequency aggregation layer: {config.name}" + ) + + def build_bottleneck( input_height: int, in_channels: int, @@ -250,9 +318,16 @@ def build_bottleneck( ) layers.append(layer) + frequency_aggregator = build_frequency_aggregation( + input_height=input_height, + in_channels=current_channels, + config=config.frequency_aggregation, + ) + return Bottleneck( input_height=input_height, in_channels=in_channels, out_channels=config.channels, + frequency_aggregator=frequency_aggregator, layers=layers, ) diff --git a/uv.lock b/uv.lock index 0a68e27..c090df2 100644 --- a/uv.lock +++ b/uv.lock @@ -453,7 +453,7 @@ wheels = [ [[package]] name = "batdetect2" -version = "2.0.0b1" +version = "2.0.0b2" source = { editable = "." } dependencies = [ { name = "click" },