.gitignore

@@ -120,6 +120,3 @@ notebooks/lightning_logs
 
 # Intermediate artifacts
 example_data/preprocessed
-
-# Dev notebooks
-notebooks/tmp

src/batdetect2/detector/models.py

@@ -94,10 +94,7 @@ class Net2DFast(nn.Module):
             num_filts // 4, 2, kernel_size=1, padding=0
         )
         self.conv_classes_op = nn.Conv2d(
-            num_filts // 4,
-            self.num_classes + 1,
-            kernel_size=1,
-            padding=0,
+            num_filts // 4, self.num_classes + 1, kernel_size=1, padding=0
         )
 
         if self.emb_dim > 0:

src/batdetect2/models/backbones.py

@@ -26,6 +26,7 @@ from soundevent import data
 from torch import nn
 
 from batdetect2.configs import BaseConfig, load_config
+from batdetect2.models.blocks import ConvBlock
 from batdetect2.models.bottleneck import (
     DEFAULT_BOTTLENECK_CONFIG,
     BottleneckConfig,
@@ -88,6 +89,7 @@ class Backbone(BackboneModel):
     def __init__(
         self,
         input_height: int,
+        out_channels: int,
         encoder: Encoder,
         decoder: Decoder,
         bottleneck: nn.Module,
@@ -114,12 +116,16 @@ class Backbone(BackboneModel):
         """
         super().__init__()
         self.input_height = input_height
+        self.out_channels = out_channels
 
         self.encoder = encoder
         self.decoder = decoder
         self.bottleneck = bottleneck
 
-        self.out_channels = decoder.out_channels
+        self.final_conv = ConvBlock(
+            in_channels=decoder.out_channels,
+            out_channels=out_channels,
+        )
 
         # Down/Up scaling factor. Need to ensure inputs are divisible by
         # this factor in order to be processed by the down/up scaling layers
@@ -158,7 +164,7 @@ class Backbone(BackboneModel):
         # Restore original size
         x = _restore_pad(x, h_pad=h_pad, w_pad=w_pad)
 
-        return x
+        return self.final_conv(x)
 
 
 class BackboneConfig(BaseConfig):
@@ -293,6 +299,7 @@ def build_backbone(config: BackboneConfig) -> BackboneModel:
 
     return Backbone(
         input_height=config.input_height,
+        out_channels=config.out_channels,
         encoder=encoder,
         decoder=decoder,
         bottleneck=bottleneck,

src/batdetect2/train/losses.py

@@ -210,7 +210,7 @@ class FocalLoss(nn.Module):
             pos_loss = pos_loss * torch.tensor(self.class_weights)
 
         if self.mask_zero:
-            valid_mask = (gt.sum(1) > 0).float().unsqueeze(1)
+            valid_mask = gt.any(dim=1, keepdim=True).float()
             pos_loss = pos_loss * valid_mask
             neg_loss = neg_loss * valid_mask
 
@@ -476,7 +476,7 @@ def build_loss(
 
     size_loss_fn = BBoxLoss()
 
-    return LossFunction(  # type: ignore
+    return LossFunction(
         size_loss=size_loss_fn,
         classification_loss=classification_loss_fn,
         detection_loss=detection_loss_fn,