Skip to content

fasterrcnn

orchard.architectures.fasterrcnn

Faster R-CNN Detection Architecture.

Wraps torchvision.models.detection.fasterrcnn_resnet50_fpn_v2 with automatic head replacement for custom class counts. The +1 for the background class is handled internally — callers pass only the number of object categories.

build_fasterrcnn(num_classes, pretrained=True)

Build a Faster R-CNN with ResNet-50-FPN v2 backbone.

Loads the torchvision pre-built model and replaces the box predictor head to match the target number of classes. Background is added automatically (num_classes + 1).

Unlike classification builders, this does not accept in_channels because Faster R-CNN always uses RGB (3-channel) input internally.

Parameters:

Name Type Description Default
num_classes int

Number of object categories (excluding background).

required
pretrained bool

If True, load COCO-pretrained weights for the backbone.

True

Returns:

Type Description
Module

Faster R-CNN model with custom class head (device placement handled by factory).

Source code in orchard/architectures/fasterrcnn.py
def build_fasterrcnn(
    num_classes: int,
    pretrained: bool = True,  # pragma: no mutate
) -> nn.Module:
    """
    Build a Faster R-CNN with ResNet-50-FPN v2 backbone.

    Loads the torchvision pre-built model and replaces the box predictor
    head to match the target number of classes. Background is added
    automatically (``num_classes + 1``).

    Unlike classification builders, this does not accept ``in_channels``
    because Faster R-CNN always uses RGB (3-channel) input internally.

    Args:
        num_classes: Number of object categories (excluding background).
        pretrained: If True, load COCO-pretrained weights for the backbone.

    Returns:
        Faster R-CNN model with custom class head (device placement handled by factory).
    """
    weights = "DEFAULT" if pretrained else None  # pragma: no mutate
    model = cast(FasterRCNN, fasterrcnn_resnet50_fpn_v2(weights=weights))  # pragma: no mutate

    # Replace the classification head for the target number of classes
    # (+1 for background, which torchvision expects)
    in_features: int = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes + 1)

    return cast("nn.Module", model)  # pragma: no mutate