Skip to content

🧠 Models

The four task models and the builders that construct them from a variant name. Each builder takes the n/s/m/l/x variant and a class count and returns a ready nn.Module; the model returns the dual-head outputs that training supervises and decoding reads. Generated from the docstrings in src/lucid_yolo/models/.

Builders

build_detector

build_detector(variant: str, num_classes: int = _DEFAULT_NUM_CLASSES) -> Detector

Build a :class:Detector for a named scale variant.

Parameters:

Name Type Description Default
variant str

Scale name ("n"/"s"/"m"/"l"/"x").

required
num_classes int

Number of object classes. Defaults to 80 (COCO).

_DEFAULT_NUM_CLASSES

Returns:

Type Description
Detector

The assembled :class:Detector module.

Raises:

Type Description
KeyError

If variant is not one of the five published names.

Examples:

>>> model = build_detector("s")
>>> model.variant, model.num_classes
('s', 80)
>>> model.backbone.width
0.5

build_segmenter

build_segmenter(variant: str, num_classes: int = _DEFAULT_NUM_CLASSES, num_coeffs: int = DEFAULT_NUM_COEFFS) -> Segmenter

Build a :class:Segmenter for a named scale variant.

num_coeffs is forwarded rather than fixed. :class:Segmenter has always accepted it, and a builder that dropped it made the class's own third argument unreachable through the function the package exports for building one — a caller wanting a narrower coefficient width had to bypass the builder and instantiate the class, which is precisely the seam builders exist to remove.

The default is unchanged, so every model this function built before is the model it builds now, byte-identical state dict included; the WP-023 parameter and FLOP goldens read the default path and do not move.

Parameters:

Name Type Description Default
variant str

Scale name ("n"/"s"/"m"/"l"/"x").

required
num_classes int

Number of object classes. Defaults to 80 (COCO).

_DEFAULT_NUM_CLASSES
num_coeffs int

Mask-coefficient width K (A14), shared by the head's coefficient stems and the prototype stack — one prototype per coefficient. Defaults to :data:DEFAULT_NUM_COEFFS.

DEFAULT_NUM_COEFFS

Returns:

Type Description
Segmenter

The assembled :class:Segmenter module, its head and prototype stack both

Segmenter

built at num_coeffs.

Raises:

Type Description
KeyError

If variant is not one of the five published names.

Examples:

>>> model = build_segmenter("s")
>>> model.variant, model.num_classes, model.num_coeffs
('s', 80, 32)
>>> model.protonet.num_prototypes  # one prototype per coefficient
32

A narrower coefficient width reaches both halves of Eq. 7 together:

>>> narrow = build_segmenter("n", num_classes=4, num_coeffs=16)
>>> narrow.num_coeffs, narrow.protonet.num_prototypes
(16, 16)

build_obb_detector

build_obb_detector(variant: str, num_classes: int = _DEFAULT_NUM_CLASSES) -> OrientedDetector

Build an :class:OrientedDetector for a named scale variant.

Parameters:

Name Type Description Default
variant str

Scale name ("n"/"s"/"m"/"l"/"x").

required
num_classes int

Number of object classes. Defaults to 80 to match :func:build_detector; DOTA-v1.0 work passes len(DOTA_CLASSES) == 15.

_DEFAULT_NUM_CLASSES

Returns:

Type Description
OrientedDetector

The assembled :class:OrientedDetector module.

Raises:

Type Description
KeyError

If variant is not one of the five published names.

Examples:

>>> model = build_obb_detector("s", num_classes=15)
>>> model.variant, model.num_classes
('s', 15)
>>> model.head.predict_angle
True

build_keypoint_detector

build_keypoint_detector(variant: str, num_classes: int = _DEFAULT_NUM_CLASSES, *, num_keypoints: int) -> KeypointDetector

Build a :class:KeypointDetector for a named scale variant.

Parameters:

Name Type Description Default
variant str

Scale name ("n"/"s"/"m"/"l"/"x").

required
num_classes int

Number of object classes. Defaults to 80 to match :func:build_detector; the COCO person_keypoints path passes 1.

_DEFAULT_NUM_CLASSES
num_keypoints int

Point count K, keyword-only and required — the point count is a property of the annotation schema, not something a model may pick (see :class:KeypointDetector).

required

Returns:

Type Description
KeypointDetector

The assembled :class:KeypointDetector module.

Raises:

Type Description
KeyError

If variant is not one of the five published names.

TypeError

If num_keypoints is omitted.

Examples:

>>> model = build_keypoint_detector("s", num_classes=1, num_keypoints=17)
>>> model.variant, model.num_classes, model.num_keypoints
('s', 1, 17)
>>> model.head.o2o.keypoint_stems is None
False

Task models

Detector

Bases: Module

Composite YOLO26 detector: backbone -> neck -> dual detection head.

A single typed module wiring the three WP-020…WP-022 stages for one scale variant. :meth:forward takes an image batch and returns the head's :class:~lucid_yolo.models.heads.DualHeadOutput (both the one-to-one and one-to-many dense predictions).

Parameters:

Name Type Description Default
variant str

Scale name ("n"/"s"/"m"/"l"/"x") resolved through :func:~lucid_yolo.models.registry.scale_spec.

required
num_classes int

Number of object classes the head predicts.

_DEFAULT_NUM_CLASSES

Examples:

>>> import torch
>>> model = Detector("n", num_classes=80).eval()
>>> with torch.no_grad():
...     out = model(torch.zeros(1, 3, 640, 640))
>>> out.o2o_cls.shape, out.o2o_box.shape
(torch.Size([1, 8400, 80]), torch.Size([1, 8400, 4]))

forward

forward(image: Tensor) -> DualHeadOutput

Run the backbone, neck, and dual head over an image batch.

Parameters:

Name Type Description Default
image Tensor

Input image batch of shape (N, 3, H, W) with H and W divisible by 32.

required

Returns:

Type Description
DualHeadOutput

The head's :class:~lucid_yolo.models.heads.DualHeadOutput with dense

DualHeadOutput

class logits and raw ltrb distances for both branches.

deploy

deploy() -> nn.Module

Return the NMS-free inference model: backbone -> neck -> one-to-one head.

The one-to-many branch is training-only (R6) and never runs at E2E inference, so the deployed model executes a single detection branch. This is the module whose FLOPs match the R1 Table 7 GFLOP column (see the module docstring); its parameters are a strict subset of the full detector's, so model size is still reported from the whole :class:Detector. The returned module shares this detector's parameters (no copy).

Returns:

Name Type Description
A Module

class:torch.nn.Module whose forward maps an image batch to the

Module

one-to-one branch's (cls_logits, ltrb) tensor pair.

Examples:

>>> import torch
>>> deployed = Detector("n").deploy().eval()
>>> with torch.no_grad():
...     cls, box = deployed(torch.zeros(1, 3, 640, 640))
>>> cls.shape, box.shape
(torch.Size([1, 8400, 80]), torch.Size([1, 8400, 4]))

Segmenter

Bases: Module

Composite YOLO26 segmentation model: detector plus the mask and aux branches.

Mirrors :class:Detector — same backbone and neck for a given variant — and adds the three WP-047…WP-050 segmentation stages on top: the detection head now also emits per-anchor mask coefficients, :class:~lucid_yolo.models.heads.ProtoFusion collapses the neck's three levels into the Eq. 8 fused feature, :class:~lucid_yolo.models.heads.ProtoNet turns that feature into K prototype maps, and :class:~lucid_yolo.models.heads.SemanticAux attaches the training-only auxiliary classifier to the same fused feature.

num_coeffs is threaded into both the head's coefficient stems and the prototype count, one prototype per coefficient: Eq. 7 contracts the two against each other, so a single constructor argument is what stops the two sides from drifting apart.

The fused feature is computed exactly once and shared by the prototype and auxiliary branches — recomputing it would silently double that subgraph's cost in the FLOP measurement.

Parameters:

Name Type Description Default
variant str

Scale name ("n"/"s"/"m"/"l"/"x") resolved through :func:~lucid_yolo.models.registry.scale_spec.

required
num_classes int

Number of object classes the head and the auxiliary branch predict.

_DEFAULT_NUM_CLASSES
num_coeffs int

Mask-coefficient width K (A14), also the prototype count.

DEFAULT_NUM_COEFFS

Examples:

>>> import torch
>>> model = Segmenter("n", num_classes=4).eval()
>>> with torch.no_grad():
...     out = model(torch.zeros(1, 3, 128, 128))
>>> out.detect.o2o_cls.shape, out.detect.o2o_coeff.shape
(torch.Size([1, 336, 4]), torch.Size([1, 336, 32]))
>>> out.prototypes.shape
torch.Size([1, 32, 32, 32])
>>> out.semantic is None  # training-only branch (A17)
True

forward

forward(image: Tensor) -> SegmentOutput

Run the backbone, neck, dual head, prototype stack, and auxiliary branch.

Parameters:

Name Type Description Default
image Tensor

Input image batch of shape (N, 3, H, W) with H and W divisible by 32.

required

Returns:

Name Type Description
A SegmentOutput

class:SegmentOutput holding the dual head's dense predictions,

SegmentOutput

the raw prototype maps, and the auxiliary semantic logits (None

SegmentOutput

outside training mode).

deploy

deploy() -> nn.Module

Return the inference model: backbone -> neck -> one-to-one head + prototypes.

Two branches of the trained model are training-only and are therefore absent from the returned module: the one-to-many detection branch, which never runs on the NMS-free E2E path (R6), and the auxiliary semantic branch, which exists only to shape the shared prototype features (A17). Neither is disabled or skipped — neither is held at all. The returned module shares this segmenter's parameters (no copy), so it is a view rather than a second model.

Returns:

Name Type Description
A Module

class:torch.nn.Module whose forward maps an image batch to the

Module

(cls, box, coeff, prototypes) tuple needed to assemble masks.

Examples:

>>> import torch
>>> deployed = Segmenter("n", num_classes=4).deploy().eval()
>>> with torch.no_grad():
...     cls, box, coeff, prototypes = deployed(torch.zeros(1, 3, 128, 128))
>>> cls.shape, box.shape
(torch.Size([1, 336, 4]), torch.Size([1, 336, 4]))
>>> coeff.shape, prototypes.shape
(torch.Size([1, 336, 32]), torch.Size([1, 32, 32, 32]))

OrientedDetector

Bases: Module

Composite YOLO26 oriented detector: the detector plus the A20 angle branch.

Structurally :class:Detector with predict_angle=True: identical backbone, neck, and dual head for a given variant, with a third per-level stem on each head branch emitting the single orientation scalar of R1 sec. 3.4.3. The angle is predicted raw (R1 Eq. 13, theta_hat = z); the oriented box is assembled and normalized by :func:~lucid_yolo.models.heads.obb.decode_rboxes (A23), not by the model.

The angle stems sit on the head's two branches rather than in a fourth top-level module, which is what makes the one-to-many angle stems training-only in the same sense the one-to-many box and class stems already are: they are dropped by :meth:deploy with the branch that owns them, rather than needing a rule of their own.

Parameters:

Name Type Description Default
variant str

Scale name ("n"/"s"/"m"/"l"/"x") resolved through :func:~lucid_yolo.models.registry.scale_spec.

required
num_classes int

Number of object classes the head predicts. Defaults to 80 for signature consistency with :class:Detector; the R1 Table S11 fidelity gate and the DOTA path pass 15 explicitly.

_DEFAULT_NUM_CLASSES

Examples:

>>> import torch
>>> model = OrientedDetector("n", num_classes=15).eval()
>>> with torch.no_grad():
...     out = model(torch.zeros(1, 3, 128, 128))
>>> out.o2o_cls.shape, out.o2o_box.shape, out.o2o_angle.shape
(torch.Size([1, 336, 15]), torch.Size([1, 336, 4]), torch.Size([1, 336, 1]))

forward

forward(image: Tensor) -> DualHeadOutput

Run the backbone, neck, and dual head with the angle branch active.

Parameters:

Name Type Description Default
image Tensor

Input image batch of shape (N, 3, H, W) with H and W divisible by 32.

required

Returns:

Type Description
DualHeadOutput

The head's :class:~lucid_yolo.models.heads.DualHeadOutput, whose

DualHeadOutput

angle fields are populated for both branches.

deploy

deploy() -> nn.Module

Return the NMS-free inference model: backbone -> neck -> one-to-one head.

The one-to-many branch is training-only (R6) and never runs at E2E inference, so the deployed model executes a single detection branch — including only that branch's angle stems. This is the module whose FLOPs the R1 Table S11 gate reads, at the table's 1024-pixel input. The returned module shares this detector's parameters (no copy).

Returns:

Name Type Description
A Module

class:torch.nn.Module whose forward maps an image batch to the

Module

one-to-one branch's (cls_logits, ltrb, angle) tensor triple.

Examples:

>>> import torch
>>> deployed = OrientedDetector("n", num_classes=15).deploy().eval()
>>> with torch.no_grad():
...     cls, box, angle = deployed(torch.zeros(1, 3, 128, 128))
>>> cls.shape, box.shape, angle.shape
(torch.Size([1, 336, 15]), torch.Size([1, 336, 4]), torch.Size([1, 336, 1]))

KeypointDetector

Bases: Module

Composite YOLO26 keypoint detector: the detector plus the R14 point branch.

Structurally :class:Detector with num_keypoints=K: identical backbone, neck, and dual head for a given variant, with a third per-level stem on each head branch emitting K raw (x, y) coordinate offsets and K raw per-axis sigma values. Coordinates are unbounded offsets composed against the anchor centre and stride by :func:~lucid_yolo.models.heads.keypoint.decode_keypoints (A70), not by the model; sigma stays raw (A65) because the RLE flow (R14 Eq. 12) is what interprets it.

The point stems sit on the head's two branches rather than in a fourth top-level module, which is what makes the one-to-many point stems training-only in the same sense the one-to-many box and class stems already are: they are dropped by :meth:deploy with the branch that owns them, rather than needing a rule of their own. This mirrors :class:OrientedDetector's angle stems.

The task is generic in K: nothing here knows what a point means, so human pose is one instantiation rather than the subject.

Parameters:

Name Type Description Default
variant str

Scale name ("n"/"s"/"m"/"l"/"x") resolved through :func:~lucid_yolo.models.registry.scale_spec.

required
num_classes int

Number of object classes the head predicts. Defaults to 80 for signature consistency with :class:Detector; COCO's person_keypoints path passes 1.

_DEFAULT_NUM_CLASSES
num_keypoints int

Point count K both head branches predict. Keyword-only and required, with no default: unlike num_classes there is nothing project-wide to fall back on, because K is a property of the dataset's annotation schema (COCO person is 17, another pose set is not). A default here would be a silent claim about data this model has never seen — the same refusal :class:~lucid_yolo.ptl.module.DetectionLitModule makes for task="keypoints".

required

Examples:

>>> import torch
>>> model = KeypointDetector("n", num_classes=4, num_keypoints=3).eval()
>>> with torch.no_grad():
...     out = model(torch.zeros(1, 3, 128, 128))
>>> out.o2o_cls.shape, out.o2o_keypoints.shape
(torch.Size([1, 336, 4]), torch.Size([1, 336, 3, 2]))
>>> out.o2o_keypoint_sigma.shape  # raw per-axis uncertainty (A65)
torch.Size([1, 336, 3, 2])

forward

forward(image: Tensor) -> DualHeadOutput

Run the backbone, neck, and dual head with the point branch active.

Parameters:

Name Type Description Default
image Tensor

Input image batch of shape (N, 3, H, W) with H and W divisible by 32.

required

Returns:

Type Description
DualHeadOutput

The head's :class:~lucid_yolo.models.heads.DualHeadOutput, whose

DualHeadOutput

keypoint and keypoint-sigma fields are populated for both branches.

deploy

deploy() -> nn.Module

Return the NMS-free inference model: backbone -> neck -> one-to-one head.

The one-to-many branch is training-only (R6) and never runs at E2E inference, so the deployed model executes a single detection branch — including only that branch's point stems. This is the module the keypoint golden's GFLOPs are measured on, at detection's 640-pixel protocol. The returned module shares this detector's parameters (no copy).

Returns:

Name Type Description
A Module

class:torch.nn.Module whose forward maps an image batch to the

Module

one-to-one branch's (cls_logits, ltrb, keypoints) tensor triple.

Examples:

>>> import torch
>>> deployed = KeypointDetector("n", num_classes=4, num_keypoints=3).deploy().eval()
>>> with torch.no_grad():
...     cls, box, keypoints = deployed(torch.zeros(1, 3, 128, 128))
>>> cls.shape, box.shape, keypoints.shape
(torch.Size([1, 336, 4]), torch.Size([1, 336, 4]), torch.Size([1, 336, 3, 2]))

Heads and outputs

DualDetectionHead

Bases: Module

Anchor-free dual detection head with direct ltrb regression (reg_max=1).

Builds the R1 sec. 3.2.1-3.2.2 / Fig. S2 head: two branches (:attr:o2o for the NMS-free one-to-one path, :attr:o2m for dense one-to-many training supervision) share the neck features (N3, N4, N5) and each own per-level box and class stems over strides 8/16/32. Boxes are predicted as 4 raw ltrb scalars per anchor with no DFL (reg_max = 1).

Parameters:

Name Type Description Default
in_channels tuple[int, int, int]

Neck output channel counts (N3, N4, N5) in stride order (8, 16, 32) — typically :attr:DetectionNeck.channels.

required
num_classes int

Number of object classes.

required
num_coeffs int | None

Optional mask-coefficient count. None preserves the detection-only head exactly; callers explicitly pass :data:DEFAULT_NUM_COEFFS to enable coefficients.

None
predict_angle bool

Build both branches' orientation stems (A20). False preserves the detection-only head exactly; the oriented model passes True.

False
num_keypoints int | None

Optional generic point count. None preserves the existing head exactly; passing K builds coordinate-and-sigma stems on both branches.

None

Examples:

>>> import torch
>>> head = DualDetectionHead(in_channels=(64, 128, 256), num_classes=80).eval()
>>> feats = (torch.zeros(1, 64, 80, 80), torch.zeros(1, 128, 40, 40), torch.zeros(1, 256, 20, 20))
>>> with torch.no_grad():
...     out = head(feats)
>>> out.o2o_cls.shape, out.o2o_box.shape
(torch.Size([1, 8400, 80]), torch.Size([1, 8400, 4]))
>>> out.o2o_angle is None  # the angle branch is opt-in
True

forward_branch

forward_branch(features: tuple[Tensor, Tensor, Tensor], branch: BranchName) -> BranchOutput

Run one named branch over the neck features.

:meth:forward runs both because training needs both: the one-to-many branch supplies the dense assignment's targets and the one-to-one branch the E2E ones. Inference does not. Every decode path reads exactly one branch -- the top-k path the one-to-one, the suppression path the one-to-many -- and the other branch's stems are computed, materialized and then dropped. This is the seam that lets a caller ask for only the half it will read.

branch is validated rather than dispatched on directly, because its violation is otherwise silent: a bare if branch == "o2o" maps "O2O", "o2o " or a typo onto the other branch and returns a well-shaped tensor that no downstream check can recognise as the wrong one.

Parameters:

Name Type Description Default
features tuple[Tensor, Tensor, Tensor]

The neck maps (n3, n4, n5) at strides 8, 16, and 32, with channel counts matching the constructor's in_channels.

required
branch BranchName

Which branch to run, one of :data:BRANCH_NAMES.

required

Returns:

Type Description
BranchOutput

That branch's :class:BranchOutput: dense class logits, raw ltrb distances,

BranchOutput

and whichever of the optional coefficient, angle and keypoint tensors the

BranchOutput

head was built with.

Raises:

Type Description
ValueError

If branch is not one of :data:BRANCH_NAMES.

Examples:

>>> import torch
>>> head = DualDetectionHead((8, 16, 32), num_classes=3).eval()
>>> maps = (torch.zeros(1, 8, 4, 4), torch.zeros(1, 16, 2, 2), torch.zeros(1, 32, 1, 1))
>>> with torch.no_grad():
...     one_to_one = head.forward_branch(maps, "o2o")
>>> one_to_one.cls.shape  # 4*4 + 2*2 + 1*1 anchors, 3 classes
torch.Size([1, 21, 3])
>>> head.forward_branch(maps, "O2O")
Traceback (most recent call last):
    ...
ValueError: branch must be one of ('o2o', 'o2m'); got 'O2O'

forward

forward(features: tuple[Tensor, Tensor, Tensor]) -> DualHeadOutput

Run both branches over the neck features.

Parameters:

Name Type Description Default
features tuple[Tensor, Tensor, Tensor]

The neck maps (n3, n4, n5) at strides 8, 16, and 32, with channel counts matching the constructor's in_channels.

required

Returns:

Name Type Description
A DualHeadOutput

class:DualHeadOutput with dense class logits, raw ltrb

DualHeadOutput

distances, and — for whichever optional stems were built — tanh mask

DualHeadOutput

coefficients, raw orientation angles, and raw point-coordinate and

DualHeadOutput

uncertainty tensors, for both branches.

DualHeadOutput dataclass

Dense predictions of both detection branches.

Every field is a dense per-anchor tensor with the A axis ordered level-by-level in row-major order (matching :func:~lucid_yolo.assign.grid.make_anchor_points). Class fields are raw logits; box fields are raw ltrb distances in stride units (decode them with :func:decode_ltrb). Coefficient fields, when enabled, are already tanh-activated in the head and lie in [-1, 1]. This deliberate asymmetry from the raw class and box outputs keeps the activation with its regression head; a later mask loss must not move it. Angle fields go the other way and are deliberately unactivated and unbounded: R1 Eq. 13 predicts theta_hat = z, so any range normalization belongs to the decode (:func:~lucid_yolo.models.heads.obb.decode_rboxes, A23), not here. Keypoint-coordinate fields are likewise raw offsets. Keypoint-sigma fields are raw and unbounded: no range or positivity normalization happens anywhere in this head, and WP-123's RLE loss owns that mapping (A65).

Attributes:

Name Type Description
o2m_cls Tensor

One-to-many class logits, shape (B, A, num_classes).

o2m_box Tensor

One-to-many raw ltrb distances, shape (B, A, 4).

o2o_cls Tensor

One-to-one class logits, shape (B, A, num_classes).

o2o_box Tensor

One-to-one raw ltrb distances, shape (B, A, 4).

o2m_coeff Tensor | None

One-to-many tanh mask coefficients, shape (B, A, K), or None when coefficients are disabled.

o2o_coeff Tensor | None

One-to-one tanh mask coefficients, shape (B, A, K), or None when coefficients are disabled.

o2m_angle Tensor | None

One-to-many raw orientation angles in radians, shape (B, A, 1), or None when the angle branch is disabled.

o2o_angle Tensor | None

One-to-one raw orientation angles in radians, shape (B, A, 1), or None when the angle branch is disabled.

o2m_keypoints Tensor | None

One-to-many raw point-coordinate offsets, shape (B, A, K, 2), or None when keypoints are disabled.

o2o_keypoints Tensor | None

One-to-one raw point-coordinate offsets, shape (B, A, K, 2), or None when keypoints are disabled.

o2m_keypoint_sigma Tensor | None

One-to-many raw, unbounded per-axis uncertainty, shape (B, A, K, 2), or None when keypoints are disabled. WP-123's RLE loss owns the positivity mapping.

o2o_keypoint_sigma Tensor | None

One-to-one raw, unbounded per-axis uncertainty, shape (B, A, K, 2), or None when keypoints are disabled. WP-123's RLE loss owns the positivity mapping.

SegmentOutput dataclass

Dense predictions of the segmentation model's three branches.

Groups what a segmentation forward pass produces without flattening it: the detection head's own dataclass is kept whole (so the dual-branch contract of :class:~lucid_yolo.models.heads.DualHeadOutput stays in one place), and the two mask-side tensors sit beside it.

Attributes:

Name Type Description
detect DualHeadOutput

The full :class:~lucid_yolo.models.heads.DualHeadOutput, including both branches' tanh mask coefficients.

prototypes Tensor

Raw Eq. 9 prototype maps of shape (B, K, 2*H3, 2*W3), i.e. twice the P3 grid (A15), unactivated.

semantic Tensor | None

Auxiliary per-class logits (B, num_classes, H3, W3) in training mode, None at eval — the branch is training-only (A17) and :class:~lucid_yolo.models.heads.SemanticAux returns None outside training rather than being skipped by the caller.