Skip to content

API Reference: Models & Components

Architectural components including Backbones, Necks, Heads, and Detector compositions.


Object Detectors

corecv.models.detector.CoreObjectDetector(backbone, neck, head)

Bases: Module

End-to-end object detection model.

Composes a backbone, optional neck, and detection head into a single module. The forward pass flows through each component sequentially:

input -> backbone -> (neck) -> head -> output

Parameters:

Name Type Description Default
backbone Module

Feature extractor implementing :class:~corecv.core.contract.BaseBackbone.

required
neck Module | None

Optional feature pyramid network (FPN, PANet, etc.). Pass None to skip.

required
head Module

Detection head producing class logits and box predictions.

required

Initialise the detector.

Parameters:

Name Type Description Default
backbone Module

Feature extractor.

required
neck Module | None

Optional neck (None to skip).

required
head Module

Detection head.

required
Source code in src/corecv/models/detector.py
def __init__(
    self,
    backbone: nn.Module,
    neck: nn.Module | None,
    head: nn.Module,
) -> None:
    """Initialise the detector.

    Args:
        backbone: Feature extractor.
        neck: Optional neck (``None`` to skip).
        head: Detection head.
    """
    super().__init__()
    self.backbone = backbone
    self.neck = neck
    self.head = head

forward(x)

Run the full detection pipeline.

Parameters:

Name Type Description Default
x Tensor

Input tensor of shape (B, 3, H, W).

required

Returns:

Type Description
object

Output from the detection head (typically a dict containing

object

"cls_logits" and "pred_boxes" or "reg_pred").

Source code in src/corecv/models/detector.py
def forward(self, x: torch.Tensor) -> object:
    """Run the full detection pipeline.

    Args:
        x: Input tensor of shape ``(B, 3, H, W)``.

    Returns:
        Output from the detection head (typically a ``dict`` containing
        ``"cls_logits"`` and ``"pred_boxes"`` or ``"reg_pred"``).
    """
    features = self.backbone(x)
    if self.neck is not None:
        features = self.neck(features)
    return self.head(features)

Backbones

ResNet

corecv.models.backbones.resnet

ResNet backbone with multi-scale feature extraction.

Wraps :mod:torchvision.models ResNet variants (18, 34, 50, 101) and extracts intermediate feature maps from layer1 through layer4 at strides 4, 8, 16, and 32 respectively.

Extracted feature levels (verified against torchvision output shapes):

.. list-table:: :header-rows: 1 :widths: 20 20 20 15 15 15 15

    • Model
    • Level
    • Stride
    • ResNet-18
    • ResNet-34
    • ResNet-50
    • ResNet-101
    • layer1
    • stride4
    • 4
    • 64
    • 64
    • 256
    • 256
    • layer2
    • stride8
    • 8
    • 128
    • 128
    • 512
    • 512
    • layer3
    • stride16
    • 16
    • 256
    • 256
    • 1024
    • 1024
    • layer4
    • stride32
    • 32
    • 512
    • 512
    • 2048
    • 2048
Example

from corecv.models.backbones.resnet import ResNet50Backbone backbone = ResNet50Backbone(pretrained=False) backbone.feature_info.channels {'stride4': 256, 'stride8': 512, 'stride16': 1024, 'stride32': 2048}

ResNet101Backbone(pretrained=True, **kwargs)

Bases: _ResNetBackbone

ResNet-101 backbone.

Feature levels
  • stride4: 256 channels
  • stride8: 512 channels
  • stride16: 1024 channels
  • stride32: 2048 channels
Source code in src/corecv/models/backbones/resnet.py
def __init__(self, pretrained: bool = True, **kwargs: object) -> None:
    """Initialise the ResNet backbone.

    Args:
        pretrained: If ``True``, load ImageNet-1K pretrained weights.
        **kwargs: Additional keyword arguments forwarded to the
            underlying ResNet factory function.
    """
    super().__init__()
    cfg = _RESNET_VARIANTS[self._variant_key]
    weights = cfg["weights"].IMAGENET1K_V1 if pretrained else None
    self._model = cfg["factory"](weights=weights, **kwargs)

ResNet18Backbone(pretrained=True, **kwargs)

Bases: _ResNetBackbone

ResNet-18 backbone.

Feature levels
  • stride4: 64 channels
  • stride8: 128 channels
  • stride16: 256 channels
  • stride32: 512 channels
Source code in src/corecv/models/backbones/resnet.py
def __init__(self, pretrained: bool = True, **kwargs: object) -> None:
    """Initialise the ResNet backbone.

    Args:
        pretrained: If ``True``, load ImageNet-1K pretrained weights.
        **kwargs: Additional keyword arguments forwarded to the
            underlying ResNet factory function.
    """
    super().__init__()
    cfg = _RESNET_VARIANTS[self._variant_key]
    weights = cfg["weights"].IMAGENET1K_V1 if pretrained else None
    self._model = cfg["factory"](weights=weights, **kwargs)

ResNet34Backbone(pretrained=True, **kwargs)

Bases: _ResNetBackbone

ResNet-34 backbone.

Feature levels
  • stride4: 64 channels
  • stride8: 128 channels
  • stride16: 256 channels
  • stride32: 512 channels
Source code in src/corecv/models/backbones/resnet.py
def __init__(self, pretrained: bool = True, **kwargs: object) -> None:
    """Initialise the ResNet backbone.

    Args:
        pretrained: If ``True``, load ImageNet-1K pretrained weights.
        **kwargs: Additional keyword arguments forwarded to the
            underlying ResNet factory function.
    """
    super().__init__()
    cfg = _RESNET_VARIANTS[self._variant_key]
    weights = cfg["weights"].IMAGENET1K_V1 if pretrained else None
    self._model = cfg["factory"](weights=weights, **kwargs)

ResNet50Backbone(pretrained=True, **kwargs)

Bases: _ResNetBackbone

ResNet-50 backbone.

Feature levels
  • stride4: 256 channels
  • stride8: 512 channels
  • stride16: 1024 channels
  • stride32: 2048 channels
Source code in src/corecv/models/backbones/resnet.py
def __init__(self, pretrained: bool = True, **kwargs: object) -> None:
    """Initialise the ResNet backbone.

    Args:
        pretrained: If ``True``, load ImageNet-1K pretrained weights.
        **kwargs: Additional keyword arguments forwarded to the
            underlying ResNet factory function.
    """
    super().__init__()
    cfg = _RESNET_VARIANTS[self._variant_key]
    weights = cfg["weights"].IMAGENET1K_V1 if pretrained else None
    self._model = cfg["factory"](weights=weights, **kwargs)

MobileNetV3

corecv.models.backbones.mobilenetv3

MobileNetV3 backbone with multi-scale feature extraction.

Wraps :mod:torchvision.models.mobilenet_v3_small and :mod:torchvision.models.mobilenet_v3_large to expose four feature levels at strides 4, 8, 16, and 32. Feature maps are extracted at the last block of each spatial resolution stage.

Extracted feature levels (verified against torchvision output shapes):

  • MobileNetV3-Small (features Sequential, 224x224 input):

============ ===== ======= ========================= Level Index Spatial Channel Count ============ ===== ======= ========================= stride4 1 56x56 16 stride8 3 28x28 24 stride16 8 14x14 48 stride32 12 7x7 576 ============ ===== ======= =========================

  • MobileNetV3-Large (features Sequential, 224x224 input):

============ ===== ======= ========================= Level Index Spatial Channel Count ============ ===== ======= ========================= stride4 3 56x56 24 stride8 6 28x28 40 stride16 12 14x14 112 stride32 16 7x7 960 ============ ===== ======= =========================

Example

from corecv.models.backbones.mobilenetv3 import MobileNetV3SmallBackbone backbone = MobileNetV3SmallBackbone(pretrained=False) backbone.feature_info.channels {'stride4': 16, 'stride8': 24, 'stride16': 48, 'stride32': 576}

MobileNetV3LargeBackbone(pretrained=True, **kwargs)

Bases: _MobileNetV3Backbone

MobileNetV3-Large backbone with four-level feature extraction.

Wraps :func:torchvision.models.mobilenet_v3_large. Feature maps are extracted from features Sequential at indices 3, 6, 12, and 16, corresponding to strides 4, 8, 16, and 32 respectively.

Feature levels
  • stride4: 24 channels, 56x56 spatial (index 3)
  • stride8: 40 channels, 28x28 spatial (index 6)
  • stride16: 112 channels, 14x14 spatial (index 12)
  • stride32: 960 channels, 7x7 spatial (index 16)
Source code in src/corecv/models/backbones/mobilenetv3.py
def __init__(self, pretrained: bool = True, **kwargs: object) -> None:
    """Initialise the MobileNetV3 backbone.

    Args:
        pretrained: If ``True``, load ImageNet-1K pretrained weights.
        **kwargs: Additional keyword arguments forwarded to the
            underlying ``mobilenet_v3_*`` factory.
    """
    super().__init__()
    weights = self._weights_enum.IMAGENET1K_V1 if pretrained else None
    self._model = self._model_factory(weights=weights, **kwargs)

MobileNetV3SmallBackbone(pretrained=True, **kwargs)

Bases: _MobileNetV3Backbone

MobileNetV3-Small backbone with four-level feature extraction.

Wraps :func:torchvision.models.mobilenet_v3_small. Feature maps are extracted from features Sequential at indices 1, 3, 8, and 12, corresponding to strides 4, 8, 16, and 32 respectively.

Feature levels
  • stride4: 16 channels, 56x56 spatial (index 1)
  • stride8: 24 channels, 28x28 spatial (index 3)
  • stride16: 48 channels, 14x14 spatial (index 8)
  • stride32: 576 channels, 7x7 spatial (index 12)
Source code in src/corecv/models/backbones/mobilenetv3.py
def __init__(self, pretrained: bool = True, **kwargs: object) -> None:
    """Initialise the MobileNetV3 backbone.

    Args:
        pretrained: If ``True``, load ImageNet-1K pretrained weights.
        **kwargs: Additional keyword arguments forwarded to the
            underlying ``mobilenet_v3_*`` factory.
    """
    super().__init__()
    weights = self._weights_enum.IMAGENET1K_V1 if pretrained else None
    self._model = self._model_factory(weights=weights, **kwargs)

ConvNeXt

corecv.models.backbones.convnext

ConvNeXt backbone with multi-scale feature extraction.

Wraps :mod:torchvision.models ConvNeXt variants (Tiny, Small, Base, Large) and extracts intermediate feature maps at four spatial scales. Each ConvNeXt features Sequential contains eight :class:torch.nn.Sequential stages paired into four resolution groups; the last stage of each group is used as the extracted feature level.

Extracted feature levels (verified against torchvision output shapes, 224x224 input):

.. list-table:: :header-rows: 1 :widths: 16 12 10 14 14 14 14

    • Level
    • Index
    • Stride
    • Tiny
    • Small
    • Base
    • Large
    • stride4
    • 1
    • 4
    • 96
    • 96
    • 128
    • 192
    • stride8
    • 3
    • 8
    • 192
    • 192
    • 256
    • 384
    • stride16
    • 5
    • 16
    • 384
    • 384
    • 512
    • 768
    • stride32
    • 7
    • 32
    • 768
    • 768
    • 1024
    • 1536
Example

from corecv.models.backbones.convnext import ConvNeXtTinyBackbone backbone = ConvNeXtTinyBackbone(pretrained=False) backbone.feature_info.channels {'stride4': 96, 'stride8': 192, 'stride16': 384, 'stride32': 768}

ConvNeXtBaseBackbone(pretrained=True, **kwargs)

Bases: _ConvNeXtBackbone

ConvNeXt-Base backbone.

Feature levels
  • stride4: 128 channels
  • stride8: 256 channels
  • stride16: 512 channels
  • stride32: 1024 channels
Source code in src/corecv/models/backbones/convnext.py
def __init__(self, pretrained: bool = True, **kwargs: object) -> None:
    """Initialise the ConvNeXt backbone.

    Args:
        pretrained: If ``True``, load ImageNet-1K pretrained weights.
        **kwargs: Additional keyword arguments forwarded to the
            underlying ``convnext_*`` factory.
    """
    super().__init__()
    cfg = _CONVNEXT_VARIANTS[self._variant_key]
    weights = cfg["weights"].IMAGENET1K_V1 if pretrained else None
    self._model = cfg["factory"](weights=weights, **kwargs)

ConvNeXtLargeBackbone(pretrained=True, **kwargs)

Bases: _ConvNeXtBackbone

ConvNeXt-Large backbone.

Feature levels
  • stride4: 192 channels
  • stride8: 384 channels
  • stride16: 768 channels
  • stride32: 1536 channels
Source code in src/corecv/models/backbones/convnext.py
def __init__(self, pretrained: bool = True, **kwargs: object) -> None:
    """Initialise the ConvNeXt backbone.

    Args:
        pretrained: If ``True``, load ImageNet-1K pretrained weights.
        **kwargs: Additional keyword arguments forwarded to the
            underlying ``convnext_*`` factory.
    """
    super().__init__()
    cfg = _CONVNEXT_VARIANTS[self._variant_key]
    weights = cfg["weights"].IMAGENET1K_V1 if pretrained else None
    self._model = cfg["factory"](weights=weights, **kwargs)

ConvNeXtSmallBackbone(pretrained=True, **kwargs)

Bases: _ConvNeXtBackbone

ConvNeXt-Small backbone.

Feature levels
  • stride4: 96 channels
  • stride8: 192 channels
  • stride16: 384 channels
  • stride32: 768 channels
Source code in src/corecv/models/backbones/convnext.py
def __init__(self, pretrained: bool = True, **kwargs: object) -> None:
    """Initialise the ConvNeXt backbone.

    Args:
        pretrained: If ``True``, load ImageNet-1K pretrained weights.
        **kwargs: Additional keyword arguments forwarded to the
            underlying ``convnext_*`` factory.
    """
    super().__init__()
    cfg = _CONVNEXT_VARIANTS[self._variant_key]
    weights = cfg["weights"].IMAGENET1K_V1 if pretrained else None
    self._model = cfg["factory"](weights=weights, **kwargs)

ConvNeXtTinyBackbone(pretrained=True, **kwargs)

Bases: _ConvNeXtBackbone

ConvNeXt-Tiny backbone.

Feature levels
  • stride4: 96 channels
  • stride8: 192 channels
  • stride16: 384 channels
  • stride32: 768 channels
Source code in src/corecv/models/backbones/convnext.py
def __init__(self, pretrained: bool = True, **kwargs: object) -> None:
    """Initialise the ConvNeXt backbone.

    Args:
        pretrained: If ``True``, load ImageNet-1K pretrained weights.
        **kwargs: Additional keyword arguments forwarded to the
            underlying ``convnext_*`` factory.
    """
    super().__init__()
    cfg = _CONVNEXT_VARIANTS[self._variant_key]
    weights = cfg["weights"].IMAGENET1K_V1 if pretrained else None
    self._model = cfg["factory"](weights=weights, **kwargs)

Vision Transformer (ViT)

corecv.models.backbones.vit

Vision Transformer (ViT) backbone with multi-scale pyramid adapter.

Implements ViT-Tiny, ViT-Small, and ViT-Base variants from scratch to enable intermediate feature extraction. A :class:SimplePyramidAdapter converts the flat patch-token sequence output into a multi-scale feature pyramid with strides 8, 16, and 32, making the ViT transparently comply with the :class:~corecv.core.contract.BaseBackbone contract.

Architecture overview::

Image (B, 3, 224, 224)
  |
  v
PatchEmbedding -> (B, N, embed_dim)   N = (H/P)*(W/P) = 14*14 = 196
  |
  v
TransformerBlock x depth
  |
  v
LayerNorm -> (B, 1+N, embed_dim)
  |
  v  (drop CLS token)
patch_tokens: (B, N, embed_dim)
  |
  v
SimplePyramidAdapter -> [f8, f16, f32]
  - f8:  (B, embed_dim, 28, 28)  stride  8
  - f16: (B, embed_dim, 14, 14)  stride 16
  - f32: (B, embed_dim,  7,  7)  stride 32

Variant configurations:

.. list-table:: :header-rows: 1

    • Variant
    • embed_dim
    • depth
    • num_heads
    • mlp_ratio
    • ViT-Tiny
    • 192
    • 12
    • 3
    • 4.0
    • ViT-Small
    • 384
    • 12
    • 6
    • 4.0
    • ViT-Base
    • 768
    • 12
    • 12
    • 4.0
Example

from corecv.models.backbones.vit import ViTBaseBackbone backbone = ViTBaseBackbone(pretrained=False) backbone.feature_info.channels {'stride8': 768, 'stride16': 768, 'stride32': 768}

SimplePyramidAdapter(in_channels, out_channels)

Bases: Module

Convert flat ViT patch tokens into a multi-scale feature pyramid.

Takes the 2D-reshaped patch tokens from the ViT encoder and produces three feature maps at strides 8, 16, and 32 using transposed convolution (upsample), pointwise projection (identity), and strided convolution (downsample).

For a 224x224 input with patch_size=16, the spatial grid is 14x14. The adapter produces: - stride 8: 28x28 via transposed convolution (2x upsample) - stride 16: 14x14 via pointwise projection - stride 32: 7x7 via strided convolution (2x downsample)

Parameters:

Name Type Description Default
in_channels int

Number of input channels (ViT embed_dim).

required
out_channels int

Number of output channels for each pyramid level.

required

Initialise the pyramid adapter.

Parameters:

Name Type Description Default
in_channels int

Number of input channels (ViT embed_dim).

required
out_channels int

Number of output channels for each pyramid level.

required
Source code in src/corecv/models/backbones/vit.py
def __init__(self, in_channels: int, out_channels: int) -> None:
    """Initialise the pyramid adapter.

    Args:
        in_channels: Number of input channels (ViT embed_dim).
        out_channels: Number of output channels for each pyramid level.
    """
    super().__init__()
    self.out_channels = out_channels

    # Stride 8: upsample 14x14 -> 28x28
    self.up = nn.Sequential(
        nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
    )
    # Stride 16: identity projection at 14x14
    self.lateral = nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
    )
    # Stride 32: downsample 14x14 -> 7x7
    self.down = nn.Sequential(
        nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=3,
            stride=2,
            padding=1,
        ),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
    )
forward(patch_tokens, grid_size)

Reshape patch tokens to 2D and create the feature pyramid.

Parameters:

Name Type Description Default
patch_tokens Tensor

Patch token sequence of shape (B, num_patches, in_channels).

required
grid_size int

Spatial side length of the patch grid (num_patches = grid_size ** 2).

required

Returns:

Type Description
list[Tensor]

A list of three feature tensors at strides 8, 16, and 32.

Source code in src/corecv/models/backbones/vit.py
def forward(
    self,
    patch_tokens: torch.Tensor,
    grid_size: int,
) -> list[torch.Tensor]:
    """Reshape patch tokens to 2D and create the feature pyramid.

    Args:
        patch_tokens: Patch token sequence of shape
            ``(B, num_patches, in_channels)``.
        grid_size: Spatial side length of the patch grid
            (``num_patches = grid_size ** 2``).

    Returns:
        A list of three feature tensors at strides 8, 16, and 32.
    """
    B, _n, C = patch_tokens.shape
    # Reshape to spatial: (B, C, grid_size, grid_size)
    feat = patch_tokens.transpose(1, 2).reshape(B, C, grid_size, grid_size)

    f8 = self.up(feat)          # (B, out_c, 2*grid, 2*grid)
    f16 = self.lateral(feat)    # (B, out_c, grid, grid)
    f32 = self.down(feat)       # (B, out_c, grid//2, grid//2)

    return [f8, f16, f32]

ViTBaseBackbone(pretrained=False, img_size=224, patch_size=16, in_channels=3, drop=0.0, attn_drop=0.0)

Bases: _ViTBackbone

ViT-Base backbone (embed_dim=768, depth=12, heads=12).

Feature levels (via :class:SimplePyramidAdapter): - stride8: 768 channels, 28x28 spatial - stride16: 768 channels, 14x14 spatial - stride32: 768 channels, 7x7 spatial

Source code in src/corecv/models/backbones/vit.py
def __init__(  # noqa: PLR0913
    self,
    pretrained: bool = False,
    img_size: int = 224,
    patch_size: int = 16,
    in_channels: int = 3,
    drop: float = 0.0,
    attn_drop: float = 0.0,
) -> None:
    """Initialise the ViT backbone.

    Args:
        pretrained: Whether to load pretrained weights.  Currently
            unused (always ``False``) as the custom ViT implementation
            does not ship pretrained checkpoints yet.
        img_size: Input image spatial size (assumed square).
        patch_size: Patch side length.
        in_channels: Number of input image channels.
        drop: Dropout probability in Transformer blocks.
        attn_drop: Attention dropout probability.
    """
    super().__init__()
    if pretrained:
        msg = (
            "Pretrained weights are not yet available for the custom "
            "ViT backbone. Initialising with random weights."
        )
        warnings.warn(msg, stacklevel=2)
    cfg = _VIT_VARIANTS[self._variant_key]
    embed_dim: int = cfg["embed_dim"]
    depth: int = cfg["depth"]
    num_heads: int = cfg["num_heads"]
    mlp_ratio: float = cfg["mlp_ratio"]

    self.patch_embed = _PatchEmbedding(img_size, patch_size, in_channels, embed_dim)
    num_patches = self.patch_embed.num_patches

    # Learnable CLS token and positional embeddings.
    self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
    self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))

    # Transformer encoder blocks.
    self.blocks = nn.ModuleList(
        [
            _TransformerBlock(embed_dim, num_heads, mlp_ratio, drop, attn_drop)
            for _ in range(depth)
        ]
    )
    self.norm = nn.LayerNorm(embed_dim)

    # Multi-scale pyramid adapter (stride 8, 16, 32).
    self.adapter = SimplePyramidAdapter(embed_dim, embed_dim)

    self._init_weights()

ViTSmallBackbone(pretrained=False, img_size=224, patch_size=16, in_channels=3, drop=0.0, attn_drop=0.0)

Bases: _ViTBackbone

ViT-Small backbone (embed_dim=384, depth=12, heads=6).

Feature levels (via :class:SimplePyramidAdapter): - stride8: 384 channels, 28x28 spatial - stride16: 384 channels, 14x14 spatial - stride32: 384 channels, 7x7 spatial

Source code in src/corecv/models/backbones/vit.py
def __init__(  # noqa: PLR0913
    self,
    pretrained: bool = False,
    img_size: int = 224,
    patch_size: int = 16,
    in_channels: int = 3,
    drop: float = 0.0,
    attn_drop: float = 0.0,
) -> None:
    """Initialise the ViT backbone.

    Args:
        pretrained: Whether to load pretrained weights.  Currently
            unused (always ``False``) as the custom ViT implementation
            does not ship pretrained checkpoints yet.
        img_size: Input image spatial size (assumed square).
        patch_size: Patch side length.
        in_channels: Number of input image channels.
        drop: Dropout probability in Transformer blocks.
        attn_drop: Attention dropout probability.
    """
    super().__init__()
    if pretrained:
        msg = (
            "Pretrained weights are not yet available for the custom "
            "ViT backbone. Initialising with random weights."
        )
        warnings.warn(msg, stacklevel=2)
    cfg = _VIT_VARIANTS[self._variant_key]
    embed_dim: int = cfg["embed_dim"]
    depth: int = cfg["depth"]
    num_heads: int = cfg["num_heads"]
    mlp_ratio: float = cfg["mlp_ratio"]

    self.patch_embed = _PatchEmbedding(img_size, patch_size, in_channels, embed_dim)
    num_patches = self.patch_embed.num_patches

    # Learnable CLS token and positional embeddings.
    self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
    self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))

    # Transformer encoder blocks.
    self.blocks = nn.ModuleList(
        [
            _TransformerBlock(embed_dim, num_heads, mlp_ratio, drop, attn_drop)
            for _ in range(depth)
        ]
    )
    self.norm = nn.LayerNorm(embed_dim)

    # Multi-scale pyramid adapter (stride 8, 16, 32).
    self.adapter = SimplePyramidAdapter(embed_dim, embed_dim)

    self._init_weights()

ViTTinyBackbone(pretrained=False, img_size=224, patch_size=16, in_channels=3, drop=0.0, attn_drop=0.0)

Bases: _ViTBackbone

ViT-Tiny backbone (embed_dim=192, depth=12, heads=3).

Feature levels (via :class:SimplePyramidAdapter): - stride8: 192 channels, 28x28 spatial - stride16: 192 channels, 14x14 spatial - stride32: 192 channels, 7x7 spatial

Source code in src/corecv/models/backbones/vit.py
def __init__(  # noqa: PLR0913
    self,
    pretrained: bool = False,
    img_size: int = 224,
    patch_size: int = 16,
    in_channels: int = 3,
    drop: float = 0.0,
    attn_drop: float = 0.0,
) -> None:
    """Initialise the ViT backbone.

    Args:
        pretrained: Whether to load pretrained weights.  Currently
            unused (always ``False``) as the custom ViT implementation
            does not ship pretrained checkpoints yet.
        img_size: Input image spatial size (assumed square).
        patch_size: Patch side length.
        in_channels: Number of input image channels.
        drop: Dropout probability in Transformer blocks.
        attn_drop: Attention dropout probability.
    """
    super().__init__()
    if pretrained:
        msg = (
            "Pretrained weights are not yet available for the custom "
            "ViT backbone. Initialising with random weights."
        )
        warnings.warn(msg, stacklevel=2)
    cfg = _VIT_VARIANTS[self._variant_key]
    embed_dim: int = cfg["embed_dim"]
    depth: int = cfg["depth"]
    num_heads: int = cfg["num_heads"]
    mlp_ratio: float = cfg["mlp_ratio"]

    self.patch_embed = _PatchEmbedding(img_size, patch_size, in_channels, embed_dim)
    num_patches = self.patch_embed.num_patches

    # Learnable CLS token and positional embeddings.
    self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
    self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))

    # Transformer encoder blocks.
    self.blocks = nn.ModuleList(
        [
            _TransformerBlock(embed_dim, num_heads, mlp_ratio, drop, attn_drop)
            for _ in range(depth)
        ]
    )
    self.norm = nn.LayerNorm(embed_dim)

    # Multi-scale pyramid adapter (stride 8, 16, 32).
    self.adapter = SimplePyramidAdapter(embed_dim, embed_dim)

    self._init_weights()

Necks

Feature Pyramid Network (FPN)

corecv.models.necks.fpn

Feature Pyramid Network (FPN) neck implementation.

Implements the FPN architecture from "Feature Pyramid Networks for Object Detection" (Lin et al., 2017). The FPN augments a backbone's multi-scale feature maps with a top-down pathway and lateral connections, producing pyramid feature maps at every scale with a uniform channel dimension.

Architecture

Given backbone feature levels [C2, C3, C4, C5] at strides [4, 8, 16, 32]:

  1. Lateral 1x1 convolutions project each level to out_channels.
  2. Top-down pathway merges coarse features into finer scales via nearest-neighbour upsampling and element-wise addition.
  3. Output 3x3 convolutions reduce aliasing artifacts from upsampling.

The result is [P2, P3, P4, P5] -- one feature map per backbone level, all with out_channels channels.

Dynamic Channel Projection

The neck accepts any :class:~corecv.core.contract.FeatureInfo and automatically instantiates the correct number of 1x1 lateral convolutions with input channels derived from the backbone metadata. This means the same :class:FPN class works with ResNet, MobileNetV3, ConvNeXt, ViT, or any future backbone that implements :class:BaseBackbone.

Example

from corecv.models.backbones.resnet import ResNet50Backbone from corecv.models.necks.fpn import FPN backbone = ResNet50Backbone(pretrained=False) neck = FPN(feature_info=backbone.feature_info, out_channels=256) features = backbone(torch.randn(1, 3, 224, 224)) pyramid = neck(features) tuple(f.shape for f in pyramid) (torch.Size([1, 256, 56, 56]), torch.Size([1, 256, 28, 28]), torch.Size([1, 256, 14, 14]), torch.Size([1, 256, 7, 7]))

FPN(feature_info, out_channels=256)

Bases: Module

Feature Pyramid Network neck with dynamic channel projection.

Produces a multi-scale feature pyramid from a backbone's intermediate feature maps. The number of lateral convolutions, their input channels, and the interpolation targets are all derived at construction time from the provided :class:FeatureInfo, making this neck backbone-agnostic.

Attributes:

Name Type Description
out_channels

Uniform channel count for all output feature levels.

levels

Sorted list of (level_name, stride, in_channels) tuples.

lateral_convs

Module mapping level name -> 1x1 conv for channel alignment in the top-down pathway.

fpn_convs

Module mapping level name -> 3x3 conv applied after the top-down merge to reduce upsampling aliasing.

Initialise the FPN neck.

Parameters:

Name Type Description Default
feature_info FeatureInfo

Metadata describing the backbone's feature levels. Used to determine the number of levels and their input channel counts for lateral convolution construction.

required
out_channels int

Output channel dimension for every pyramid level. All lateral 1x1 convolutions project to this width.

256
Source code in src/corecv/models/necks/fpn.py
def __init__(
    self,
    feature_info: FeatureInfo,
    out_channels: int = 256,
) -> None:
    """Initialise the FPN neck.

    Args:
        feature_info: Metadata describing the backbone's feature levels.
            Used to determine the number of levels and their input
            channel counts for lateral convolution construction.
        out_channels: Output channel dimension for every pyramid level.
            All lateral 1x1 convolutions project to this width.
    """
    super().__init__()
    self.out_channels = out_channels
    self.levels = _sorted_levels(feature_info)

    # 1x1 lateral convolutions: project backbone channels -> out_channels
    lateral_convs: list[tuple[str, nn.Module]] = []
    # 3x3 output convolutions: suppress upsampling aliasing
    fpn_convs: list[tuple[str, nn.Module]] = []

    for name, _stride, in_ch in self.levels:
        lateral_convs.append(
            (
                name,
                nn.Conv2d(in_ch, out_channels, kernel_size=1),
            )
        )
        fpn_convs.append(
            (
                name,
                nn.Conv2d(
                    out_channels,
                    out_channels,
                    kernel_size=3,
                    padding=1,
                ),
            )
        )

    self.lateral_convs = nn.ModuleDict(OrderedDict(lateral_convs))
    self.fpn_convs = nn.ModuleDict(OrderedDict(fpn_convs))
forward(features)

Apply the FPN top-down pathway to backbone features.

Parameters:

Name Type Description Default
features Sequence[Tensor]

Sequence of backbone feature tensors in stride- ascending order (i.e. highest resolution first). The length must match the number of levels in the :class:FeatureInfo provided at construction time.

required

Returns:

Type Description
list[Tensor]

A list of pyramid feature tensors in stride-ascending order

list[Tensor]

(smallest stride first), each with shape

list[Tensor]

(B, out_channels, H_i, W_i).

Raises:

Type Description
ValueError

If the number of input features does not match the number of registered backbone levels.

Source code in src/corecv/models/necks/fpn.py
def forward(self, features: Sequence[torch.Tensor]) -> list[torch.Tensor]:
    """Apply the FPN top-down pathway to backbone features.

    Args:
        features: Sequence of backbone feature tensors **in stride-
            ascending order** (i.e. highest resolution first).  The
            length must match the number of levels in the
            :class:`FeatureInfo` provided at construction time.

    Returns:
        A list of pyramid feature tensors in stride-ascending order
        (smallest stride first), each with shape
        ``(B, out_channels, H_i, W_i)``.

    Raises:
        ValueError: If the number of input features does not match
            the number of registered backbone levels.
    """
    if len(features) != len(self.levels):
        msg = (
            f"Expected {len(self.levels)} feature maps "
            f"(one per backbone level), received {len(features)}."
        )
        raise ValueError(msg)

    # Step 1 -- Lateral projections (1x1 convs) to uniform channel dim
    laterals: list[torch.Tensor] = []
    for i, (name, _stride, _in_ch) in enumerate(self.levels):
        laterals.append(self.lateral_convs[name](features[i]))

    # Step 2 -- Top-down pathway (deepest -> shallowest)
    # Start from the deepest level (last element) and propagate
    # coarser features into finer scales via nearest-neighbour upsample
    # and element-wise addition.
    for i in range(len(self.levels) - 1, 0, -1):
        target_h = laterals[i - 1].shape[-2]
        target_w = laterals[i - 1].shape[-1]
        upsampled = F.interpolate(
            laterals[i],
            size=(target_h, target_w),
            mode="nearest",
        )
        laterals[i - 1] = laterals[i - 1] + upsampled

    # Step 3 -- Output 3x3 convolutions to reduce aliasing
    outputs: list[torch.Tensor] = []
    for i, (name, _stride, _in_ch) in enumerate(self.levels):
        outputs.append(self.fpn_convs[name](laterals[i]))

    return outputs

Path Aggregation Network (PANet)

corecv.models.necks.panet

Path Aggregation Network (PANet) neck implementation.

Implements the PANet architecture from "Path Aggregation Network for Instance Segmentation" (Liu et al., 2018). PANet extends FPN by adding a bottom-up path augmentation on top of the FPN's top-down pyramid, providing a shorter information path from low-level features to the deepest layers and improving feature fusion across all scales.

Architecture

Given backbone feature levels [C2, C3, C4, C5] at strides [4, 8, 16, 32]:

  1. FPN top-down pathway (identical to :class:~corecv.models.necks.fpn.FPN) produces [P2, P3, P4, P5].
  2. Bottom-up path augmentation merges FPN outputs from shallowest to deepest via stride-2 convolutions for spatial downsampling and element-wise addition, producing [N2, N3, N4, N5].
  3. Output 3x3 convolutions are applied after each bottom-up merge.

The result is [N2, N3, N4, N5] -- one feature map per backbone level, all with out_channels channels.

Dynamic Channel Projection

Like :class:~corecv.models.necks.fpn.FPN, the PANet accepts any :class:~corecv.core.contract.FeatureInfo and dynamically instantiates all 1x1 and 3x3 convolutions based on backbone metadata.

Example

from corecv.models.backbones.resnet import ResNet50Backbone from corecv.models.necks.panet import PANet backbone = ResNet50Backbone(pretrained=False) neck = PANet(feature_info=backbone.feature_info, out_channels=256) features = backbone(torch.randn(1, 3, 224, 224)) pyramid = neck(features) tuple(f.shape for f in pyramid) (torch.Size([1, 256, 56, 56]), torch.Size([1, 256, 28, 28]), torch.Size([1, 256, 14, 14]), torch.Size([1, 256, 7, 7]))

PANet(feature_info, out_channels=256)

Bases: Module

Path Aggregation Network neck with dynamic channel projection.

Combines a top-down FPN pathway with a bottom-up path augmentation, yielding richer multi-scale features for detection and segmentation heads.

Attributes:

Name Type Description
out_channels

Uniform channel count for all output feature levels.

levels

Sorted list of (level_name, stride, in_channels) tuples.

lateral_convs

1x1 convolutions for FPN top-down lateral connections.

fpn_convs

3x3 convolutions after FPN top-down merges.

panet_reduce_convs

3x3 stride-2 convolutions for bottom-up spatial downsampling between adjacent PANet levels.

panet_lateral_convs

1x1 convolutions that project FPN outputs into the PANet bottom-up pathway.

panet_convs

3x3 convolutions applied after each bottom-up addition.

Initialise the PANet neck.

Parameters:

Name Type Description Default
feature_info FeatureInfo

Metadata describing the backbone's feature levels. Used to determine the number of levels and their input channel counts for all convolution construction.

required
out_channels int

Output channel dimension for every pyramid level. All lateral and projection convolutions target this width.

256
Source code in src/corecv/models/necks/panet.py
def __init__(
    self,
    feature_info: FeatureInfo,
    out_channels: int = 256,
) -> None:
    """Initialise the PANet neck.

    Args:
        feature_info: Metadata describing the backbone's feature levels.
            Used to determine the number of levels and their input
            channel counts for all convolution construction.
        out_channels: Output channel dimension for every pyramid level.
            All lateral and projection convolutions target this width.
    """
    super().__init__()
    self.out_channels = out_channels
    self.levels = _sorted_levels(feature_info)
    num_levels = len(self.levels)

    # -----------------------------------------------------------------
    # FPN top-down pathway (same as FPN class)
    # -----------------------------------------------------------------
    lateral_convs: list[tuple[str, nn.Module]] = []
    fpn_convs: list[tuple[str, nn.Module]] = []

    for name, _stride, in_ch in self.levels:
        lateral_convs.append(
            (
                name,
                nn.Conv2d(in_ch, out_channels, kernel_size=1),
            )
        )
        fpn_convs.append(
            (
                name,
                nn.Conv2d(
                    out_channels,
                    out_channels,
                    kernel_size=3,
                    padding=1,
                ),
            )
        )

    self.lateral_convs = nn.ModuleDict(OrderedDict(lateral_convs))
    self.fpn_convs = nn.ModuleDict(OrderedDict(fpn_convs))

    # -----------------------------------------------------------------
    # Bottom-up path augmentation
    # -----------------------------------------------------------------

    # 1x1 lateral convolutions: project FPN outputs into PANet pathway
    panet_lateral_convs: list[tuple[str, nn.Module]] = []

    # 3x3 stride-2 convolutions: downsample from level i to level i+1
    panet_reduce_convs: list[tuple[str, nn.Module]] = []

    # 3x3 output convolutions: applied after each bottom-up addition
    panet_convs: list[tuple[str, nn.Module]] = []

    for _i, (name, _stride, _in_ch) in enumerate(self.levels):
        panet_lateral_convs.append(
            (
                name,
                nn.Conv2d(out_channels, out_channels, kernel_size=1),
            )
        )
        panet_convs.append(
            (
                name,
                nn.Conv2d(
                    out_channels,
                    out_channels,
                    kernel_size=3,
                    padding=1,
                ),
            )
        )

    # Stride-2 conv from level i -> level i+1 exists for all but the
    # deepest level (deepest level is the starting point of the
    # bottom-up path and requires no downsampling).
    for i in range(num_levels - 1):
        src_name = self.levels[i][0]
        panet_reduce_convs.append(
            (
                f"{src_name}_reduce",
                nn.Conv2d(
                    out_channels,
                    out_channels,
                    kernel_size=3,
                    stride=2,
                    padding=1,
                ),
            )
        )

    self.panet_lateral_convs = nn.ModuleDict(OrderedDict(panet_lateral_convs))
    self.panet_reduce_convs = nn.ModuleDict(OrderedDict(panet_reduce_convs))
    self.panet_convs = nn.ModuleDict(OrderedDict(panet_convs))
forward(features)

Apply FPN top-down then PANet bottom-up pathways.

Parameters:

Name Type Description Default
features Sequence[Tensor]

Sequence of backbone feature tensors in stride- ascending order (highest resolution first). The length must match the number of levels in the :class:FeatureInfo provided at construction time.

required

Returns:

Type Description
list[Tensor]

A list of feature tensors in stride-ascending order (smallest

list[Tensor]

stride first), each with shape

list[Tensor]

(B, out_channels, H_i, W_i).

Raises:

Type Description
ValueError

If the number of input features does not match the number of registered backbone levels.

Source code in src/corecv/models/necks/panet.py
def forward(self, features: Sequence[torch.Tensor]) -> list[torch.Tensor]:
    """Apply FPN top-down then PANet bottom-up pathways.

    Args:
        features: Sequence of backbone feature tensors **in stride-
            ascending order** (highest resolution first).  The length
            must match the number of levels in the
            :class:`FeatureInfo` provided at construction time.

    Returns:
        A list of feature tensors in stride-ascending order (smallest
        stride first), each with shape
        ``(B, out_channels, H_i, W_i)``.

    Raises:
        ValueError: If the number of input features does not match
            the number of registered backbone levels.
    """
    if len(features) != len(self.levels):
        msg = (
            f"Expected {len(self.levels)} feature maps "
            f"(one per backbone level), received {len(features)}."
        )
        raise ValueError(msg)

    num_levels = len(self.levels)

    # -----------------------------------------------------------------
    # Stage 1 -- FPN top-down pathway
    # -----------------------------------------------------------------

    # Lateral projections to uniform channel dimension
    laterals: list[torch.Tensor] = []
    for i, (name, _stride, _in_ch) in enumerate(self.levels):
        laterals.append(self.lateral_convs[name](features[i]))

    # Top-down merge: deepest -> shallowest
    for i in range(num_levels - 1, 0, -1):
        target_h = laterals[i - 1].shape[-2]
        target_w = laterals[i - 1].shape[-1]
        upsampled = F.interpolate(
            laterals[i],
            size=(target_h, target_w),
            mode="nearest",
        )
        laterals[i - 1] = laterals[i - 1] + upsampled

    # FPN output convolutions
    fpn_out: list[torch.Tensor] = []
    for i, (name, _stride, _in_ch) in enumerate(self.levels):
        fpn_out.append(self.fpn_convs[name](laterals[i]))

    # -----------------------------------------------------------------
    # Stage 2 -- Bottom-up path augmentation
    # -----------------------------------------------------------------

    # Project each FPN output into the PANet pathway
    panet_features: list[torch.Tensor] = []
    for i, (name, _stride, _in_ch) in enumerate(self.levels):
        panet_features.append(self.panet_lateral_convs[name](fpn_out[i]))

    # Bottom-up merge: shallowest -> deepest
    # The deepest level (index num_levels - 1) is the starting point
    # and only needs its lateral projection (already done above).
    for i in range(1, num_levels):
        src_name = self.levels[i - 1][0]
        downsampled = self.panet_reduce_convs[f"{src_name}_reduce"](
            panet_features[i - 1],
        )
        panet_features[i] = panet_features[i] + downsampled

    # PANet output convolutions
    outputs: list[torch.Tensor] = []
    for i, (name, _stride, _in_ch) in enumerate(self.levels):
        outputs.append(self.panet_convs[name](panet_features[i]))

    return outputs

Heads

Classification Head

corecv.models.heads.classification

Classification head modules for CoreCV.

Provides a simple linear classification head that consumes the coarsest feature level from a backbone or neck and produces per-class logits.

Available classification heads

.. list-table:: :header-rows: 1 :widths: 25 30 45

    • Class
    • Registry Key
    • Description
    • :class:LinearClassificationHead
    • linear_classification
    • Global average pooling + fully-connected layer for image-level classification.
Example

from corecv.models.heads.classification import LinearClassificationHead from corecv.core.contract import FeatureInfo fi = FeatureInfo(channels={"feat": 2048}, strides={"feat": 32}) head = LinearClassificationHead(feature_info=fi, num_classes=1000)

LinearClassificationHead(feature_info, num_classes, dropout=0.0)

Bases: Module

Image-level classification head with global average pooling.

Consumes the coarsest (last) feature map from a backbone or neck, applies adaptive average pooling to 1x1, and passes through a 1x1 convolution (equivalent to nn.Linear on a 1x1 spatial map) to produce per-class logits.

A 1x1 convolution is used instead of nn.Linear to ensure compatibility with device='meta' shape propagation.

The head dynamically inspects the :class:FeatureInfo to determine the input channel count, making it compatible with any backbone that implements :class:~corecv.core.contract.BaseBackbone.

Parameters:

Name Type Description Default
feature_info FeatureInfo

Metadata from the upstream backbone, used to determine input channel count from the coarsest feature level.

required
num_classes int

Number of output classes.

required
dropout float

Optional dropout probability applied before the classifier (default 0.0 i.e. no dropout).

0.0

Initialise the classification head.

Parameters:

Name Type Description Default
feature_info FeatureInfo

Backbone feature metadata.

required
num_classes int

Number of output classes.

required
dropout float

Dropout probability (default 0.0).

0.0
Source code in src/corecv/models/heads/classification/linear_head.py
def __init__(
    self,
    feature_info: FeatureInfo,
    num_classes: int,
    dropout: float = 0.0,
) -> None:
    """Initialise the classification head.

    Args:
        feature_info: Backbone feature metadata.
        num_classes: Number of output classes.
        dropout: Dropout probability (default 0.0).
    """
    super().__init__()
    self.num_classes = num_classes

    # Sort levels by stride (ascending = finest first).
    sorted_levels: list[str] = sorted(
        feature_info.strides.keys(),
        key=lambda k: feature_info.strides[k],  # type: ignore[arg-type]
    )
    if not sorted_levels:
        msg = "feature_info must contain at least one feature level."
        raise ValueError(msg)

    self._feature_info = feature_info
    self._sorted_levels = sorted_levels
    # Use the coarsest (last) level for classification.
    coarsest_level: str = sorted_levels[-1]
    in_channels: int = feature_info.channels[coarsest_level]

    self.pool = nn.AdaptiveAvgPool2d(1)
    self.dropout = nn.Dropout(p=dropout) if dropout > 0 else nn.Identity()
    # 1x1 conv equivalent to nn.Linear(in_channels, num_classes).
    self.fc = nn.Conv2d(in_channels, num_classes, kernel_size=1)
feature_info property

Return the feature metadata used at construction time.

Returns:

Name Type Description
The FeatureInfo

class:FeatureInfo instance passed to __init__.

forward(features)

Produce classification logits from multi-scale feature maps.

Parameters:

Name Type Description Default
features dict[str, Tensor] | Sequence[Tensor]

Feature maps from a backbone or neck. Can be a dict of level-name -> tensor or a Sequence (list/tuple) of tensors ordered from finest to coarsest.

required

Returns:

Type Description
Tensor

Logit tensor of shape (batch_size, num_classes).

Source code in src/corecv/models/heads/classification/linear_head.py
def forward(
    self, features: dict[str, Tensor] | Sequence[Tensor]
) -> Tensor:
    """Produce classification logits from multi-scale feature maps.

    Args:
        features: Feature maps from a backbone or neck.  Can be a
            ``dict`` of level-name -> tensor or a ``Sequence``
            (list/tuple) of tensors ordered from finest to coarsest.

    Returns:
        Logit tensor of shape ``(batch_size, num_classes)``.
    """
    if isinstance(features, dict):
        # Dict: pick the last level in sorted (by stride) order.
        x: Tensor = features[self._sorted_levels[-1]]
    else:
        # Sequence: last element is the coarsest.
        x = features[-1]

    x = self.pool(x)  # (B, C, 1, 1)
    x = self.fc(x)    # (B, num_classes, 1, 1)
    x = self.dropout(x)
    return x.flatten(1)  # (B, num_classes)

Detection Head

corecv.models.heads.detection

Detection head modules for CoreCV.

Provides :class:~corecv.core.contract.FeatureInfo-aware detection heads for anchor-free and query-based object detection.

Available heads

.. list-table:: :header-rows: 1 :widths: 25 20 30

    • Class
    • Registry Key
    • Description
    • :class:DecoupledAnchorFreeHead
    • decoupled_anchor_free
    • FCOS/YOLOX-style pure-conv head with decoupled cls/reg branches.
    • :class:QueryDetectionHead
    • query_detection
    • RT-DETR/D-FINE-style transformer-decoder head with learnable queries and NMS-free inference.
Example

from corecv.models.heads.detection import ( ... DecoupledAnchorFreeHead, ... QueryDetectionHead, ... )

DecoupledAnchorFreeHead(feature_info, num_classes, feat_channels=256, num_convs=4)

Bases: Module

FCOS/YOLOX-style decoupled anchor-free detection head.

Dynamically adapts to an upstream :class:FeatureInfo by constructing per-level branches whose input channels and count match the feature map metadata. The head is fully stride-aware: a per-level stride buffer is registered so that downstream loss and post-processing can map spatial predictions back to the original image coordinates.

Parameters:

Name Type Description Default
feature_info FeatureInfo

Feature metadata from the backbone or neck. The head builds one branch per key in feature_info.channels.

required
num_classes int

Number of foreground object classes (excluding background).

required
feat_channels int

Intermediate channel dimension inside each per-level conv stack. Default 256.

256
num_convs int

Number of 3x3 Conv-BN-ReLU blocks in the shared portion of each branch. Default 4.

4

Initialise the decoupled anchor-free detection head.

Builds per-level branches matching the FeatureInfo metadata and registers per-level stride buffers for coordinate mapping.

Source code in src/corecv/models/heads/detection/decoupled_anchor_free.py
def __init__(
    self,
    feature_info: FeatureInfo,
    num_classes: int,
    feat_channels: int = 256,
    num_convs: int = 4,
) -> None:
    """Initialise the decoupled anchor-free detection head.

    Builds per-level branches matching the ``FeatureInfo`` metadata and
    registers per-level stride buffers for coordinate mapping.
    """
    super().__init__()

    self.num_classes = num_classes
    self.feat_channels = feat_channels
    self._feature_info = feature_info

    # Preserve insertion order so levels go finest -> coarsest
    self._level_names: list[str] = list(feature_info.channels.keys())
    self._level_channels: list[int] = [
        feature_info.channels[k] for k in self._level_names
    ]
    self._level_strides: list[int] = [
        feature_info.strides[k] for k in self._level_names
    ]

    # Per-level branches
    self.level_heads = nn.ModuleList(
        [
            _FeatureLevelHead(
                in_channels=ch,
                feat_channels=feat_channels,
                num_classes=num_classes,
                num_convs=num_convs,
            )
            for ch in self._level_channels
        ]
    )

    # Register stride buffers (non-learnable, move with device)
    for name, stride in zip(self._level_names, self._level_strides, strict=True):
        self.register_buffer(
            f"stride_{name}",
            torch.tensor(stride, dtype=torch.int64, requires_grad=False),
        )

    self._init_weights()
level_names property

Return the ordered list of feature level names.

level_strides property

Return the ordered list of per-level stride values.

num_levels property

Return the number of feature levels this head operates on.

forward(features)

Run detection on all feature levels.

Parameters:

Name Type Description Default
features list[Tensor]

List of feature tensors from the backbone or neck, one per level, ordered from finest to coarsest spatial resolution. len(features) must equal self.num_levels.

required

Returns:

Type Description
dict[str, list[Tensor]]

Dictionary with keys: - "cls_logits": list of (B, num_classes, H_l, W_l) - "reg_pred": list of (B, 4, H_l, W_l) - "centerness": list of (B, 1, H_l, W_l)

Source code in src/corecv/models/heads/detection/decoupled_anchor_free.py
def forward(
    self,
    features: list[Tensor],
) -> dict[str, list[Tensor]]:
    """Run detection on all feature levels.

    Args:
        features: List of feature tensors from the backbone or neck,
            one per level, ordered from finest to coarsest spatial
            resolution.  ``len(features)`` must equal
            ``self.num_levels``.

    Returns:
        Dictionary with keys:
            - ``"cls_logits"``: list of ``(B, num_classes, H_l, W_l)``
            - ``"reg_pred"``: list of ``(B, 4, H_l, W_l)``
            - ``"centerness"``: list of ``(B, 1, H_l, W_l)``
    """
    if len(features) != self.num_levels:
        msg = (
            f"Expected {self.num_levels} feature levels, "
            f"got {len(features)}."
        )
        raise ValueError(msg)

    cls_logits: list[Tensor] = []
    reg_pred: list[Tensor] = []
    centerness: list[Tensor] = []

    for feat, head in zip(features, self.level_heads, strict=True):
        cls, reg, crt = head(feat)
        cls_logits.append(cls)
        reg_pred.append(reg)
        centerness.append(crt)

    return {
        "cls_logits": cls_logits,
        "reg_pred": reg_pred,
        "centerness": centerness,
    }
get_stride_tensors(device)

Return per-level stride tensors on the given device.

Used by downstream loss modules and post-processors that need to map per-cell spatial predictions back to image coordinates.

Parameters:

Name Type Description Default
device device

Target device for the stride tensors.

required

Returns:

Type Description
list[Tensor]

List of scalar int64 tensors, one per level.

Source code in src/corecv/models/heads/detection/decoupled_anchor_free.py
def get_stride_tensors(self, device: torch.device) -> list[Tensor]:
    """Return per-level stride tensors on the given device.

    Used by downstream loss modules and post-processors that need
    to map per-cell spatial predictions back to image coordinates.

    Args:
        device: Target device for the stride tensors.

    Returns:
        List of scalar ``int64`` tensors, one per level.
    """
    return [
        getattr(self, f"stride_{name}").to(device=device)
        for name in self._level_names
    ]

QueryDetectionHead(feature_info, num_classes, d_model=256, num_queries=300, num_decoder_layers=6, num_heads=8, dim_feedforward=2048, dropout=0.1, num_reg_fcs=3, return_intermediate=False)

Bases: Module

RT-DETR / D-FINE style query-based detection head.

Learns a set of fixed object queries that attend to multi-scale features via a transformer decoder. Each query directly regresses a normalised (cx, cy, w, h) bounding box and produces class logits, enabling NMS-free inference.

The head dynamically inspects FeatureInfo.channels to construct per-level input projections, making it transparent to backbone and neck architecture changes.

Parameters:

Name Type Description Default
feature_info FeatureInfo

Feature metadata from the backbone or neck. Per-level projections are built from feature_info.channels.

required
num_classes int

Number of foreground object classes (excluding background).

required
d_model int

Transformer hidden dimension. Default 256.

256
num_queries int

Number of learnable object queries. Default 300.

300
num_decoder_layers int

Depth of the transformer decoder. Default 6.

6
num_heads int

Number of attention heads in the decoder. Default 8.

8
dim_feedforward int

FFN hidden dimension in the decoder. Default 2048.

2048
dropout float

Dropout probability. Default 0.1.

0.1
num_reg_fcs int

Number of FC layers in the box regression MLP. Default 3.

3
return_intermediate bool

If True, forward returns decoder intermediate outputs (useful for auxiliary losses). Default False.

False

Initialise the query-based detection head.

Builds multi-scale feature projections, learnable object queries, a transformer decoder stack, and classification / regression heads.

Source code in src/corecv/models/heads/detection/query_detection_head.py
def __init__(  # noqa: PLR0913
    self,
    feature_info: FeatureInfo,
    num_classes: int,
    d_model: int = 256,
    num_queries: int = 300,
    num_decoder_layers: int = 6,
    num_heads: int = 8,
    dim_feedforward: int = 2048,
    dropout: float = 0.1,
    num_reg_fcs: int = 3,
    return_intermediate: bool = False,
) -> None:
    """Initialise the query-based detection head.

    Builds multi-scale feature projections, learnable object queries, a
    transformer decoder stack, and classification / regression heads.
    """
    super().__init__()

    self.num_classes = num_classes
    self.d_model = d_model
    self.num_queries = num_queries
    self.num_decoder_layers = num_decoder_layers
    self.return_intermediate = return_intermediate

    # --- Feature-level metadata (preserves insertion order) ---
    self._level_names: list[str] = list(feature_info.channels.keys())
    self._level_channels: list[int] = [
        feature_info.channels[k] for k in self._level_names
    ]
    self._level_strides: list[int] = [
        feature_info.strides[k] for k in self._level_names
    ]

    # --- Learnable object queries ---
    self.query_embed = nn.Embedding(num_queries, d_model)

    # --- Multi-scale feature projection ---
    self.ms_projector = _MSFeatureProjector(
        in_channels_per_level=self._level_channels,
        d_model=d_model,
    )

    # --- Positional encoding for queries ---
    # Learned positional bias added to queries at each decoder layer.
    self.query_pos_embed = nn.Embedding(num_queries, d_model)

    # --- Transformer decoder ---
    self.decoder_layers = nn.ModuleList(
        [
            _TransformerDecoderLayer(
                d_model=d_model,
                num_heads=num_heads,
                dim_feedforward=dim_feedforward,
                dropout=dropout,
            )
            for _ in range(num_decoder_layers)
        ]
    )
    self.decoder_norm = nn.LayerNorm(d_model)

    # --- Classification head ---
    self.cls_head = _MLP(
        in_features=d_model,
        hidden_features=d_model,
        out_features=num_classes,
        num_layers=3,
        dropout=dropout,
    )

    # --- Box regression head ---
    self.reg_head = _MLP(
        in_features=d_model,
        hidden_features=d_model,
        out_features=4,
        num_layers=num_reg_fcs,
        dropout=dropout,
    )

    self._init_weights()
level_names property

Return the ordered list of feature level names.

level_strides property

Return the ordered list of per-level stride values.

num_levels property

Return the number of feature levels.

forward(features)

Run query-based detection on multi-scale features.

Parameters:

Name Type Description Default
features list[Tensor]

List of feature tensors from the backbone or neck, one per level, ordered from finest to coarsest spatial resolution. len(features) must equal self.num_levels.

required

Returns:

Name Type Description
dict[str, Tensor | list[Tensor]]

Dictionary with keys: - "cls_logits": (B, num_queries, num_classes) - "pred_boxes": (B, num_queries, 4) normalised (cx, cy, w, h) in [0, 1] range.

dict[str, Tensor | list[Tensor]]

If return_intermediate is True, the dictionary also

contains dict[str, Tensor | list[Tensor]]
  • "intermediate_cls": list of (B, N_q, num_classes)
  • "intermediate_reg": list of (B, N_q, 4)
Source code in src/corecv/models/heads/detection/query_detection_head.py
def forward(
    self,
    features: list[Tensor],
) -> dict[str, Tensor | list[Tensor]]:
    """Run query-based detection on multi-scale features.

    Args:
        features: List of feature tensors from the backbone or neck,
            one per level, ordered from finest to coarsest spatial
            resolution.  ``len(features)`` must equal
            ``self.num_levels``.

    Returns:
        Dictionary with keys:
            - ``"cls_logits"``: ``(B, num_queries, num_classes)``
            - ``"pred_boxes"``: ``(B, num_queries, 4)`` normalised
              ``(cx, cy, w, h)`` in ``[0, 1]`` range.

        If ``return_intermediate`` is ``True``, the dictionary also
        contains:
            - ``"intermediate_cls"``: list of ``(B, N_q, num_classes)``
            - ``"intermediate_reg"``: list of ``(B, N_q, 4)``
    """
    if len(features) != self.num_levels:
        msg = (
            f"Expected {self.num_levels} feature levels, "
            f"got {len(features)}."
        )
        raise ValueError(msg)

    B = features[0].shape[0]

    # Flatten and project multi-scale features: (B, S, d_model)
    memory = self.ms_projector(features)

    # Initialise queries: (B, num_queries, d_model)
    queries = self.query_embed.weight.unsqueeze(0).expand(B, -1, -1)
    query_pos = self.query_pos_embed.weight.unsqueeze(0).expand(B, -1, -1)

    # Decode
    intermediate_cls: list[Tensor] = []
    intermediate_reg: list[Tensor] = []

    for layer in self.decoder_layers:
        queries = layer(queries + query_pos, memory)
        if self.return_intermediate:
            intermediate_cls.append(self.cls_head(queries))
            intermediate_reg.append(self.reg_head(queries).sigmoid())

    queries = self.decoder_norm(queries)  # (B, N_q, d_model)

    # Final predictions
    cls_logits = self.cls_head(queries)          # (B, N_q, num_classes)
    pred_boxes = self.reg_head(queries).sigmoid()  # (B, N_q, 4) in [0, 1]

    result: dict[str, Tensor | list[Tensor]] = {
        "cls_logits": cls_logits,
        "pred_boxes": pred_boxes,
    }

    if self.return_intermediate:
        result["intermediate_cls"] = intermediate_cls
        result["intermediate_reg"] = intermediate_reg

    return result
get_reference_points(features)

Generate normalised reference points for each spatial location.

Produces a grid of cx, cy coordinates in [0, 1] for each feature level. This is useful for downstream box refinement (e.g. D-FINE style iterative decoding).

Parameters:

Name Type Description Default
features list[Tensor]

Feature tensors (used only for spatial dimensions).

required

Returns:

Type Description
Tensor

Reference points of shape (B, S_total, 2) where

Tensor

S_total = sum(H_l * W_l).

Source code in src/corecv/models/heads/detection/query_detection_head.py
def get_reference_points(
    self,
    features: list[Tensor],
) -> Tensor:
    """Generate normalised reference points for each spatial location.

    Produces a grid of ``cx, cy`` coordinates in ``[0, 1]`` for each
    feature level.  This is useful for downstream box refinement (e.g.
    D-FINE style iterative decoding).

    Args:
        features: Feature tensors (used only for spatial dimensions).

    Returns:
        Reference points of shape ``(B, S_total, 2)`` where
        ``S_total = sum(H_l * W_l)``.
    """
    reference_points: list[Tensor] = []
    for feat in features:
        H, W = feat.shape[2], feat.shape[3]
        # Generate normalised cy, cx grid
        cy = (torch.arange(H, device=feat.device, dtype=feat.dtype) + 0.5) / H
        cx = (torch.arange(W, device=feat.device, dtype=feat.dtype) + 0.5) / W
        grid_y, grid_x = torch.meshgrid(cy, cx, indexing="ij")
        # (2, H, W) -> (H*W, 2)
        grid = torch.stack([grid_x.flatten(), grid_y.flatten()], dim=-1)
        # Broadcast to batch: (1, H*W, 2)
        reference_points.append(grid.unsqueeze(0))

    # (1, S_total, 2) -> (B, S_total, 2)
    ref = torch.cat(reference_points, dim=1)
    return ref.expand(features[0].shape[0], -1, -1)

Segmentation Head

corecv.models.heads.segmentation

Segmentation head modules for CoreCV.

Provides decoder implementations for semantic segmentation tasks. All heads dynamically adapt to backbone :class:~corecv.core.contract.FeatureInfo metadata and are registered in :func:~corecv.core.registry.HEAD_REGISTRY.

Available segmentation heads

.. list-table:: :header-rows: 1 :widths: 25 30 45

    • Class
    • Registry Key
    • Description
    • :class:ResUNetDecoder
    • resunet_decoder
    • U-Net style decoder with residual blocks and skip connections
    • :class:ASPPDecoder
    • aspp_decoder
    • DeepLabV3+ style decoder with ASPP context module
Example

from corecv.models.backbones.resnet import ResNet50Backbone from corecv.models.heads.segmentation import ASPPDecoder backbone = ResNet50Backbone(pretrained=False) decoder = ASPPDecoder( ... feature_info=backbone.feature_info, ... out_channels=256, ... num_classes=21, ... ) feats = backbone(torch.randn(1, 3, 224, 224)) logits = decoder(feats)

ASPPDecoder(feature_info, out_channels=256, num_classes=21, atrous_rates=(6, 12, 18), dropout=0.5)

Bases: Module

DeepLabV3+ style segmentation decoder with ASPP and lightweight decoder.

Applies ASPP (Atrous Spatial Pyramid Pooling) to the coarsest feature level for multi-scale context aggregation, then fuses with low-level encoder features through a lightweight decoder path.

Input projections are created dynamically from :class:~corecv.core.contract.FeatureInfo metadata, making this decoder compatible with any backbone that exposes multi-scale features.

For backbones with 4 levels (stride 4, 8, 16, 32): - ASPP is applied to stride-32 features. - Low-level features come from stride-8. - Decoder upsamples 32->8 (4x) then 8->2 (4x) for 8x total.

For backbones with 3 levels (stride 8, 16, 32): - ASPP is applied to stride-32 features. - Low-level features come from stride-8 (the finest available). - Decoder upsamples 32->8 (4x) then 8->2 (4x) for 8x total.

Parameters:

Name Type Description Default
feature_info FeatureInfo

Feature metadata from the backbone or neck, containing channel counts and strides for each feature level.

required
out_channels int

Internal channel dimension throughout the decoder.

256
num_classes int

Number of output segmentation classes.

21
atrous_rates Sequence[int]

Dilation rates for the ASPP module.

(6, 12, 18)
dropout float

Dropout probability in the ASPP and classification head.

0.5

Raises:

Type Description
ValueError

If feature_info contains fewer than 2 feature levels.

Initialise the ASPP decoder.

Parameters:

Name Type Description Default
feature_info FeatureInfo

Feature metadata from the backbone or neck.

required
out_channels int

Internal channel dimension.

256
num_classes int

Number of segmentation classes.

21
atrous_rates Sequence[int]

ASPP dilation rates.

(6, 12, 18)
dropout float

Dropout probability.

0.5
Source code in src/corecv/models/heads/segmentation/aspp_decoder.py
def __init__(
    self,
    feature_info: FeatureInfo,
    out_channels: int = 256,
    num_classes: int = 21,
    atrous_rates: Sequence[int] = (6, 12, 18),
    dropout: float = 0.5,
) -> None:
    """Initialise the ASPP decoder.

    Args:
        feature_info: Feature metadata from the backbone or neck.
        out_channels: Internal channel dimension.
        num_classes: Number of segmentation classes.
        atrous_rates: ASPP dilation rates.
        dropout: Dropout probability.
    """
    super().__init__()
    self.num_classes = num_classes

    # ------------------------------------------------------------------
    # Sort feature levels by stride (ascending = finest first).
    # ------------------------------------------------------------------
    sorted_levels = sorted(
        feature_info.strides.keys(),
        key=lambda k: feature_info.strides[k],
    )
    if len(sorted_levels) < _MIN_LEVELS:
        msg = (
            f"ASPPDecoder requires at least {_MIN_LEVELS} feature "
            f"levels, got {len(sorted_levels)}."
        )
        raise ValueError(msg)

    self._sorted_levels = sorted_levels
    in_channels_list = [
        feature_info.channels[lvl] for lvl in sorted_levels
    ]

    # ------------------------------------------------------------------
    # Input projection: map each backbone feature to out_channels.
    # ------------------------------------------------------------------
    self.input_projections = nn.ModuleDict(
        {
            lvl: nn.Sequential(
                nn.Conv2d(ch, out_channels, kernel_size=1, bias=False),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(inplace=True),
            )
            for lvl, ch in zip(
                sorted_levels, in_channels_list, strict=True
            )
        }
    )

    # ------------------------------------------------------------------
    # ASPP on the coarsest level.
    # ------------------------------------------------------------------
    self.aspp = _ASPPModule(
        in_channels=out_channels,
        out_channels=out_channels,
        atrous_rates=atrous_rates,
    )

    # ------------------------------------------------------------------
    # Decoder path: two stages of upsample + low-level fusion.
    # ------------------------------------------------------------------
    # Stage 1: ASPP output (coarsest) -> upsample to second-finest.
    self.decoder_stage1 = _DecoderBlock(
        in_channels=out_channels,
        low_level_channels=out_channels,
        out_channels=out_channels,
    )

    # Stage 2: upsample to finest resolution.
    self.decoder_stage2 = _DecoderBlock(
        in_channels=out_channels,
        low_level_channels=out_channels,
        out_channels=out_channels,
    )

    # ------------------------------------------------------------------
    # Classification head.
    # ------------------------------------------------------------------
    self.classifier = nn.Sequential(
        nn.Dropout2d(p=dropout),
        nn.Conv2d(out_channels, num_classes, kernel_size=1),
    )
forward(features)

Forward pass through the ASPP decoder.

Parameters:

Name Type Description Default
features Sequence[Tensor]

Ordered sequence of feature tensors from the backbone, sorted by ascending stride (finest to coarsest). Each tensor has shape (B, C_i, H_i, W_i).

required

Returns:

Type Description
Tensor

Segmentation logits of shape (B, num_classes, H, W)

Tensor

where H and W are approximately 2x the finest input

Tensor

feature's spatial dimensions (or 4x if a single decoder stage

Tensor

suffices).

Raises:

Type Description
ValueError

If the number of features does not match the number of levels in feature_info.

Source code in src/corecv/models/heads/segmentation/aspp_decoder.py
def forward(self, features: Sequence[Tensor]) -> Tensor:
    """Forward pass through the ASPP decoder.

    Args:
        features: Ordered sequence of feature tensors from the backbone,
            sorted by ascending stride (finest to coarsest).  Each
            tensor has shape ``(B, C_i, H_i, W_i)``.

    Returns:
        Segmentation logits of shape ``(B, num_classes, H, W)``
        where ``H`` and ``W`` are approximately 2x the finest input
        feature's spatial dimensions (or 4x if a single decoder stage
        suffices).

    Raises:
        ValueError: If the number of features does not match the number
            of levels in ``feature_info``.
    """
    if len(features) != len(self._sorted_levels):
        msg = (
            f"Expected {len(self._sorted_levels)} feature maps, "
            f"got {len(features)}."
        )
        raise ValueError(msg)

    # ------------------------------------------------------------------
    # Step 1: Project all encoder features.
    # ------------------------------------------------------------------
    projected: OrderedDict[str, Tensor] = OrderedDict()
    for lvl, feat in zip(
        self._sorted_levels, features, strict=True
    ):
        projected[lvl] = self.input_projections[lvl](feat)

    # ------------------------------------------------------------------
    # Step 2: ASPP on the coarsest level.
    # ------------------------------------------------------------------
    coarsest_lvl = self._sorted_levels[-1]
    x = self.aspp(projected[coarsest_lvl])

    # ------------------------------------------------------------------
    # Step 3: Decoder stage 1 -- upsample to second-finest resolution
    # with low-level skip connection.
    # ------------------------------------------------------------------
    second_finest_lvl = self._sorted_levels[1]
    x = self.decoder_stage1(x, projected[second_finest_lvl])

    # ------------------------------------------------------------------
    # Step 4: Decoder stage 2 -- upsample to finest resolution.
    # ------------------------------------------------------------------
    finest_lvl = self._sorted_levels[0]
    x = self.decoder_stage2(x, projected[finest_lvl])

    # ------------------------------------------------------------------
    # Step 5: Classification head.
    # ------------------------------------------------------------------
    return self.classifier(x)

ResUNetDecoder(feature_info, out_channels=256, num_classes=21, dropout=0.1)

Bases: Module

U-Net style segmentation decoder with residual blocks and skip connections.

Consumes an ordered list of multi-scale feature maps (finest to coarsest) and produces per-pixel class logits at the original input resolution. Input projections are created dynamically from :class:~corecv.core.contract.FeatureInfo metadata.

The decoder expects features sorted by ascending stride (finest first). For backbones with 4 levels (stride 4, 8, 16, 32), the full decoder path is used. For backbones with 3 levels (stride 8, 16, 32), the stride-4 stage is skipped.

Parameters:

Name Type Description Default
feature_info FeatureInfo

Feature metadata from the backbone or neck, containing channel counts and strides for each feature level.

required
out_channels int

Number of channels in the bottleneck and decoder internal feature maps.

256
num_classes int

Number of output segmentation classes.

21
dropout float

Dropout probability applied to the final classification head (default 0.1).

0.1

Raises:

Type Description
ValueError

If feature_info contains fewer than 2 feature levels.

Initialise the ResUNet decoder.

Parameters:

Name Type Description Default
feature_info FeatureInfo

Feature metadata from the backbone or neck.

required
out_channels int

Internal channel dimension.

256
num_classes int

Number of segmentation classes.

21
dropout float

Dropout probability in the classification head.

0.1
Source code in src/corecv/models/heads/segmentation/resunet_decoder.py
def __init__(
    self,
    feature_info: FeatureInfo,
    out_channels: int = 256,
    num_classes: int = 21,
    dropout: float = 0.1,
) -> None:
    """Initialise the ResUNet decoder.

    Args:
        feature_info: Feature metadata from the backbone or neck.
        out_channels: Internal channel dimension.
        num_classes: Number of segmentation classes.
        dropout: Dropout probability in the classification head.
    """
    super().__init__()
    self.num_classes = num_classes

    # ------------------------------------------------------------------
    # Determine the sorted feature levels and their channel counts.
    # ------------------------------------------------------------------
    # Sort by stride (ascending = finest first).
    sorted_levels = sorted(
        feature_info.strides.keys(),
        key=lambda k: feature_info.strides[k],
    )
    if len(sorted_levels) < _MIN_LEVELS:
        msg = (
            f"ResUNetDecoder requires at least {_MIN_LEVELS} feature "
            f"levels, got {len(sorted_levels)}."
        )
        raise ValueError(msg)

    self._sorted_levels = sorted_levels
    in_channels_list = [
        feature_info.channels[lvl] for lvl in sorted_levels
    ]

    # ------------------------------------------------------------------
    # Input projection: map each backbone feature to common channel dim.
    # ------------------------------------------------------------------
    self.input_projections = nn.ModuleDict(
        {
            lvl: nn.Sequential(
                nn.Conv2d(ch, out_channels, kernel_size=1, bias=False),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(inplace=True),
            )
            for lvl, ch in zip(
                sorted_levels, in_channels_list, strict=True
            )
        }
    )

    # ------------------------------------------------------------------
    # Bottleneck: residual blocks at the coarsest scale.
    # ------------------------------------------------------------------
    self.bottleneck = nn.Sequential(
        _ResidualBlock(out_channels),
        _ResidualBlock(out_channels),
    )

    # ------------------------------------------------------------------
    # Decoder path: up-blocks from coarsest to finest.
    # ------------------------------------------------------------------
    num_stages = len(sorted_levels) - 1
    self.decoder_stages = nn.ModuleList()

    for _i in range(num_stages):
        # _i=0: upsample from coarsest to second-coarsest.
        # _i=1: upsample from second-coarsest to third-coarsest, etc.
        # All projected features share the same channel dim.
        self.decoder_stages.append(
            _UpBlock(
                in_channels=out_channels,
                skip_channels=out_channels,
                out_channels=out_channels,
            )
        )

    # ------------------------------------------------------------------
    # Final classification head: 1x1 conv + optional dropout.
    # ------------------------------------------------------------------
    self.classifier = nn.Sequential(
        nn.Dropout2d(p=dropout),
        nn.Conv2d(out_channels, num_classes, kernel_size=1),
    )
forward(features)

Forward pass through the ResUNet decoder.

Parameters:

Name Type Description Default
features Sequence[Tensor]

Ordered sequence of feature tensors from the backbone, sorted by ascending stride (finest to coarsest). Each tensor has shape (B, C_i, H_i, W_i).

required

Returns:

Type Description
Tensor

Segmentation logits of shape (B, num_classes, H, W)

Tensor

where H and W match the spatial dimensions of the

Tensor

finest-scale input feature.

Raises:

Type Description
ValueError

If the number of features does not match the number of levels in feature_info.

Source code in src/corecv/models/heads/segmentation/resunet_decoder.py
def forward(self, features: Sequence[Tensor]) -> Tensor:
    """Forward pass through the ResUNet decoder.

    Args:
        features: Ordered sequence of feature tensors from the backbone,
            sorted by ascending stride (finest to coarsest).  Each
            tensor has shape ``(B, C_i, H_i, W_i)``.

    Returns:
        Segmentation logits of shape ``(B, num_classes, H, W)``
        where ``H`` and ``W`` match the spatial dimensions of the
        finest-scale input feature.

    Raises:
        ValueError: If the number of features does not match the number
            of levels in ``feature_info``.
    """
    if len(features) != len(self._sorted_levels):
        msg = (
            f"Expected {len(self._sorted_levels)} feature maps, "
            f"got {len(features)}."
        )
        raise ValueError(msg)

    # ------------------------------------------------------------------
    # Step 1: Project all encoder features to common channel dimension.
    # ------------------------------------------------------------------
    projected: OrderedDict[str, Tensor] = OrderedDict()
    for lvl, feat in zip(
        self._sorted_levels, features, strict=True
    ):
        projected[lvl] = self.input_projections[lvl](feat)

    # ------------------------------------------------------------------
    # Step 2: Bottleneck at the coarsest scale.
    # ------------------------------------------------------------------
    coarsest_lvl = self._sorted_levels[-1]
    x = self.bottleneck(projected[coarsest_lvl])

    # ------------------------------------------------------------------
    # Step 3: Decoder path with skip connections (coarsest -> finest).
    # ------------------------------------------------------------------
    for i, stage in enumerate(self.decoder_stages):
        skip_idx = len(self._sorted_levels) - 2 - i
        skip_lvl = self._sorted_levels[skip_idx]
        x = stage(x, projected[skip_lvl])

    # ------------------------------------------------------------------
    # Step 4: Classification head.
    # ------------------------------------------------------------------
    return self.classifier(x)