🧠 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 a :class:Detector for a named scale variant.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
variant
|
str
|
Scale name ( |
required |
num_classes
|
int
|
Number of object classes. Defaults to 80 (COCO). |
_DEFAULT_NUM_CLASSES
|
Returns:
| Type | Description |
|---|---|
Detector
|
The assembled :class: |
Raises:
| Type | Description |
|---|---|
KeyError
|
If |
Examples:
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 ( |
required |
num_classes
|
int
|
Number of object classes. Defaults to 80 (COCO). |
_DEFAULT_NUM_CLASSES
|
num_coeffs
|
int
|
Mask-coefficient width |
DEFAULT_NUM_COEFFS
|
Returns:
| Type | Description |
|---|---|
Segmenter
|
The assembled :class: |
Segmenter
|
built at |
Raises:
| Type | Description |
|---|---|
KeyError
|
If |
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:
build_obb_detector
¶
Build an :class:OrientedDetector for a named scale variant.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
variant
|
str
|
Scale name ( |
required |
num_classes
|
int
|
Number of object classes. Defaults to 80 to match
:func: |
_DEFAULT_NUM_CLASSES
|
Returns:
| Type | Description |
|---|---|
OrientedDetector
|
The assembled :class: |
Raises:
| Type | Description |
|---|---|
KeyError
|
If |
Examples:
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 ( |
required |
num_classes
|
int
|
Number of object classes. Defaults to 80 to match
:func: |
_DEFAULT_NUM_CLASSES
|
num_keypoints
|
int
|
Point count |
required |
Returns:
| Type | Description |
|---|---|
KeypointDetector
|
The assembled :class: |
Raises:
| Type | Description |
|---|---|
KeyError
|
If |
TypeError
|
If |
Examples:
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 ( |
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
¶
Run the backbone, neck, and dual head over an image batch.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
image
|
Tensor
|
Input image batch of shape |
required |
Returns:
| Type | Description |
|---|---|
DualHeadOutput
|
The head's :class: |
DualHeadOutput
|
class logits and raw ltrb distances for both branches. |
deploy
¶
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: |
Module
|
one-to-one branch's |
Examples:
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 ( |
required |
num_classes
|
int
|
Number of object classes the head and the auxiliary branch predict. |
_DEFAULT_NUM_CLASSES
|
num_coeffs
|
int
|
Mask-coefficient width |
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
¶
Run the backbone, neck, dual head, prototype stack, and auxiliary branch.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
image
|
Tensor
|
Input image batch of shape |
required |
Returns:
| Name | Type | Description |
|---|---|---|
A |
SegmentOutput
|
class: |
SegmentOutput
|
the raw prototype maps, and the auxiliary semantic logits ( |
|
SegmentOutput
|
outside training mode). |
deploy
¶
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: |
Module
|
|
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 ( |
required |
num_classes
|
int
|
Number of object classes the head predicts. Defaults to 80
for signature consistency with :class: |
_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
¶
Run the backbone, neck, and dual head with the angle branch active.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
image
|
Tensor
|
Input image batch of shape |
required |
Returns:
| Type | Description |
|---|---|
DualHeadOutput
|
The head's :class: |
DualHeadOutput
|
angle fields are populated for both branches. |
deploy
¶
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: |
Module
|
one-to-one branch's |
Examples:
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 ( |
required |
num_classes
|
int
|
Number of object classes the head predicts. Defaults to 80
for signature consistency with :class: |
_DEFAULT_NUM_CLASSES
|
num_keypoints
|
int
|
Point count |
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
¶
Run the backbone, neck, and dual head with the point branch active.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
image
|
Tensor
|
Input image batch of shape |
required |
Returns:
| Type | Description |
|---|---|
DualHeadOutput
|
The head's :class: |
DualHeadOutput
|
keypoint and keypoint-sigma fields are populated for both branches. |
deploy
¶
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: |
Module
|
one-to-one branch's |
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 |
required |
num_classes
|
int
|
Number of object classes. |
required |
num_coeffs
|
int | None
|
Optional mask-coefficient count. |
None
|
predict_angle
|
bool
|
Build both branches' orientation stems (A20). |
False
|
num_keypoints
|
int | None
|
Optional generic point count. |
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
¶
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 |
required |
branch
|
BranchName
|
Which branch to run, one of :data: |
required |
Returns:
| Type | Description |
|---|---|
BranchOutput
|
That branch's :class: |
BranchOutput
|
and whichever of the optional coefficient, angle and keypoint tensors the |
BranchOutput
|
head was built with. |
Raises:
| Type | Description |
|---|---|
ValueError
|
If |
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
¶
Run both branches over the neck features.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
features
|
tuple[Tensor, Tensor, Tensor]
|
The neck maps |
required |
Returns:
| Name | Type | Description |
|---|---|---|
A |
DualHeadOutput
|
class: |
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 |
o2m_box |
Tensor
|
One-to-many raw ltrb distances, shape |
o2o_cls |
Tensor
|
One-to-one class logits, shape |
o2o_box |
Tensor
|
One-to-one raw ltrb distances, shape |
o2m_coeff |
Tensor | None
|
One-to-many tanh mask coefficients, shape |
o2o_coeff |
Tensor | None
|
One-to-one tanh mask coefficients, shape |
o2m_angle |
Tensor | None
|
One-to-many raw orientation angles in radians, shape
|
o2o_angle |
Tensor | None
|
One-to-one raw orientation angles in radians, shape
|
o2m_keypoints |
Tensor | None
|
One-to-many raw point-coordinate offsets, shape |
o2o_keypoints |
Tensor | None
|
One-to-one raw point-coordinate offsets, shape |
o2m_keypoint_sigma |
Tensor | None
|
One-to-many raw, unbounded per-axis uncertainty,
shape |
o2o_keypoint_sigma |
Tensor | None
|
One-to-one raw, unbounded per-axis uncertainty,
shape |
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: |
prototypes |
Tensor
|
Raw Eq. 9 prototype maps of shape |
semantic |
Tensor | None
|
Auxiliary per-class logits |