Reference for ultralytics/nn/modules/head.py

class Detect(nn.Module):
    """YOLO Detect head for object detection models.

    This class implements the detection head used in YOLO models for predicting bounding boxes and class probabilities.
    It supports both training and inference modes, with optional end-to-end detection capabilities.

    Attributes:
        dynamic (bool): Force grid reconstruction.
        export (bool): Export mode flag.
        format (str): Export format.
        end2end (bool): End-to-end detection mode.
        max_det (int): Maximum detections per image.
        shape (tuple): Input shape.
        anchors (torch.Tensor): Anchor points.
        strides (torch.Tensor): Feature map strides.
        legacy (bool): Backward compatibility for v3/v5/v8/v9 models.
        xyxy (bool): Output format, xyxy or xywh.
        nc (int): Number of classes.
        nl (int): Number of detection layers.
        reg_max (int): DFL channels.
        no (int): Number of outputs per anchor.
        stride (torch.Tensor): Strides computed during build.
        cv2 (nn.ModuleList): Convolution layers for box regression.
        cv3 (nn.ModuleList): Convolution layers for classification.
        dfl (nn.Module): Distribution Focal Loss layer.
        one2one_cv2 (nn.ModuleList): One-to-one convolution layers for box regression.
        one2one_cv3 (nn.ModuleList): One-to-one convolution layers for classification.

    Methods:
        forward: Perform forward pass and return predictions.
        forward_end2end: Perform forward pass for end-to-end detection.
        bias_init: Initialize detection head biases.
        decode_bboxes: Decode bounding boxes from predictions.
        postprocess: Post-process model predictions.

    Examples:
        Create a detection head for 80 classes
        >>> detect = Detect(nc=80, ch=(256, 512, 1024))
        >>> x = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
        >>> outputs = detect(x)
    """

    dynamic = False  # force grid reconstruction
    export = False  # export mode
    format = None  # export format
    max_det = 300  # max_det
    agnostic_nms = False
    shape = None
    anchors = torch.empty(0)  # init
    strides = torch.empty(0)  # init
    legacy = False  # backward compatibility for v3/v5/v8/v9 models
    xyxy = False  # xyxy or xywh output

    def __init__(self, nc: int = 80, reg_max=16, end2end=False, ch: tuple = ()):
        """Initialize the YOLO detection layer with specified number of classes and channels.

        Args:
            nc (int): Number of classes.
            reg_max (int): Maximum number of DFL channels.
            end2end (bool): Whether to use end-to-end NMS-free detection.
            ch (tuple): Tuple of channel sizes from backbone feature maps.
        """
        super().__init__()
        self.nc = nc  # number of classes
        self.nl = len(ch)  # number of detection layers
        self.reg_max = reg_max  # DFL channels (ch[0] // 16 to scale 4/8/12/16/20 for n/s/m/l/x)
        self.no = nc + self.reg_max * 4  # number of outputs per anchor
        self.stride = torch.zeros(self.nl)  # strides computed during build
        c2, c3 = max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], min(self.nc, 100))  # channels
        self.cv2 = nn.ModuleList(
            nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch
        )
        self.cv3 = (
            nn.ModuleList(nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, self.nc, 1)) for x in ch)
            if self.legacy
            else nn.ModuleList(
                nn.Sequential(
                    nn.Sequential(DWConv(x, x, 3), Conv(x, c3, 1)),
                    nn.Sequential(DWConv(c3, c3, 3), Conv(c3, c3, 1)),
                    nn.Conv2d(c3, self.nc, 1),
                )
                for x in ch
            )
        )
        self.dfl = DFL(self.reg_max) if self.reg_max > 1 else nn.Identity()

        if end2end:
            self.one2one_cv2 = copy.deepcopy(self.cv2)
            self.one2one_cv3 = copy.deepcopy(self.cv3)

@property
def one2many(self):
    """Returns the one-to-many head components, here for v5/v5/v8/v9/11 backward compatibility."""
    return dict(box_head=self.cv2, cls_head=self.cv3)

@property
def one2one(self):
    """Returns the one-to-one head components."""
    return dict(box_head=self.one2one_cv2, cls_head=self.one2one_cv3)

@property
def end2end(self):
    """Checks if the model has one2one for v5/v5/v8/v9/11 backward compatibility."""
    return getattr(self, "_end2end", True) and hasattr(self, "one2one")

def _get_decode_boxes(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
    """Get decoded boxes based on anchors and strides."""
    shape = x["feats"][0].shape  # BCHW
    if self.dynamic or self.shape != shape:
        self.anchors, self.strides = (a.transpose(0, 1) for a in make_anchors(x["feats"], self.stride, 0.5))
        self.shape = shape

    dbox = self.decode_bboxes(self.dfl(x["boxes"]), self.anchors.unsqueeze(0)) * self.strides
    return dbox

def _inference(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
    """Decode predicted bounding boxes and class probabilities based on multiple-level feature maps.

    Args:
        x (dict[str, torch.Tensor]): List of feature maps from different detection layers.

    Returns:
        (torch.Tensor): Concatenated tensor of decoded bounding boxes and class probabilities.
    """
    # Inference path
    dbox = self._get_decode_boxes(x)
    return torch.cat((dbox, x["scores"].sigmoid()), 1)

def bias_init(self):
    """Initialize Detect() biases, WARNING: requires stride availability."""
    for i, (a, b) in enumerate(zip(self.one2many["box_head"], self.one2many["cls_head"])):  # from
        a[-1].bias.data[:] = 2.0  # box
        b[-1].bias.data[: self.nc] = math.log(
            5 / self.nc / (640 / self.stride[i]) ** 2
        )  # cls (.01 objects, 80 classes, 640 img)
    if self.end2end:
        for i, (a, b) in enumerate(zip(self.one2one["box_head"], self.one2one["cls_head"])):  # from
            a[-1].bias.data[:] = 2.0  # box
            b[-1].bias.data[: self.nc] = math.log(
                5 / self.nc / (640 / self.stride[i]) ** 2
            )  # cls (.01 objects, 80 classes, 640 img)

def decode_bboxes(self, bboxes: torch.Tensor, anchors: torch.Tensor, xywh: bool = True) -> torch.Tensor:
    """Decode bounding boxes from predictions."""
    return dist2bbox(
        bboxes,
        anchors,
        xywh=xywh and not self.end2end and not self.xyxy,
        dim=1,
    )

@end2end.setter
def end2end(self, value):
    """Override the end-to-end detection mode."""
    self._end2end = value

def forward(
    self, x: list[torch.Tensor]
) -> dict[str, torch.Tensor] | torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]:
    """Concatenates and returns predicted bounding boxes and class probabilities."""
    preds = self.forward_head(x, **self.one2many)
    if self.end2end:
        x_detach = [xi.detach() for xi in x]
        one2one = self.forward_head(x_detach, **self.one2one)
        preds = {"one2many": preds, "one2one": one2one}
    if self.training:
        return preds
    y = self._inference(preds["one2one"] if self.end2end else preds)
    if self.end2end:
        y = self.postprocess(y.permute(0, 2, 1))
    return y if self.export else (y, preds)

def forward_head(
    self, x: list[torch.Tensor], box_head: torch.nn.Module = None, cls_head: torch.nn.Module = None
) -> dict[str, torch.Tensor]:
    """Concatenates and returns predicted bounding boxes and class probabilities."""
    if box_head is None or cls_head is None:  # for fused inference
        return dict()
    bs = x[0].shape[0]  # batch size
    boxes = torch.cat([box_head[i](x[i]).view(bs, 4 * self.reg_max, -1) for i in range(self.nl)], dim=-1)
    scores = torch.cat([cls_head[i](x[i]).view(bs, self.nc, -1) for i in range(self.nl)], dim=-1)
    return dict(boxes=boxes, scores=scores, feats=x)

def fuse(self) -> None:
    """Remove the one2many head for inference optimization."""
    self.cv2 = self.cv3 = None

def get_topk_index(self, scores: torch.Tensor, max_det: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Get top-k indices from scores.

    Args:
        scores (torch.Tensor): Scores tensor with shape (batch_size, num_anchors, num_classes).
        max_det (int): Maximum detections per image.

    Returns:
        (torch.Tensor, torch.Tensor, torch.Tensor): Top scores, class indices, and filtered indices.
    """
    batch_size, anchors, nc = scores.shape  # i.e. shape(16,8400,84)
    # Use max_det directly during export for TensorRT compatibility (requires k to be constant),
    # otherwise use min(max_det, anchors) for safety with small inputs during Python inference
    k = max_det if self.export else min(max_det, anchors)
    if self.agnostic_nms:
        scores, labels = scores.max(dim=-1, keepdim=True)
        scores, indices = scores.topk(k, dim=1)
        labels = labels.gather(1, indices)
        return scores, labels, indices
    ori_index = scores.max(dim=-1)[0].topk(k)[1].unsqueeze(-1)
    scores = scores.gather(dim=1, index=ori_index.repeat(1, 1, nc))
    scores, index = scores.flatten(1).topk(k)
    idx = ori_index[torch.arange(batch_size)[..., None], index // nc]  # original index
    return scores[..., None], (index % nc)[..., None].float(), idx

def postprocess(self, preds: torch.Tensor) -> torch.Tensor:
    """Post-processes YOLO model predictions.

    Args:
        preds (torch.Tensor): Raw predictions with shape (batch_size, num_anchors, 4 + nc) with last dimension
            format [x, y, w, h, class_probs].

    Returns:
        (torch.Tensor): Processed predictions with shape (batch_size, min(max_det, num_anchors), 6) and last
            dimension format [x, y, w, h, max_class_prob, class_index].
    """
    boxes, scores = preds.split([4, self.nc], dim=-1)
    scores, conf, idx = self.get_topk_index(scores, self.max_det)
    boxes = boxes.gather(dim=1, index=idx.repeat(1, 1, 4))
    return torch.cat([boxes, scores, conf], dim=-1)

class Segment(Detect):
    """YOLO Segment head for segmentation models.

    This class extends the Detect head to include mask prediction capabilities for instance segmentation tasks.

    Attributes:
        nm (int): Number of masks.
        npr (int): Number of protos.
        proto (Proto): Prototype generation module.
        cv4 (nn.ModuleList): Convolution layers for mask coefficients.

    Methods:
        forward: Return model outputs and mask coefficients.

    Examples:
        Create a segmentation head
        >>> segment = Segment(nc=80, nm=32, npr=256, ch=(256, 512, 1024))
        >>> x = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
        >>> outputs = segment(x)
    """

    def __init__(self, nc: int = 80, nm: int = 32, npr: int = 256, reg_max=16, end2end=False, ch: tuple = ()):
        """Initialize the YOLO model attributes such as the number of masks, prototypes, and the convolution layers.

        Args:
            nc (int): Number of classes.
            nm (int): Number of masks.
            npr (int): Number of protos.
            reg_max (int): Maximum number of DFL channels.
            end2end (bool): Whether to use end-to-end NMS-free detection.
            ch (tuple): Tuple of channel sizes from backbone feature maps.
        """
        super().__init__(nc, reg_max, end2end, ch)
        self.nm = nm  # number of masks
        self.npr = npr  # number of protos
        self.proto = Proto(ch[0], self.npr, self.nm)  # protos

        c4 = max(ch[0] // 4, self.nm)
        self.cv4 = nn.ModuleList(nn.Sequential(Conv(x, c4, 3), Conv(c4, c4, 3), nn.Conv2d(c4, self.nm, 1)) for x in ch)
        if end2end:
            self.one2one_cv4 = copy.deepcopy(self.cv4)

@property
def one2many(self):
    """Returns the one-to-many head components, here for backward compatibility."""
    return dict(box_head=self.cv2, cls_head=self.cv3, mask_head=self.cv4)

@property
def one2one(self):
    """Returns the one-to-one head components."""
    return dict(box_head=self.one2one_cv2, cls_head=self.one2one_cv3, mask_head=self.one2one_cv4)

def _inference(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
    """Decode predicted bounding boxes and class probabilities, concatenated with mask coefficients."""
    preds = super()._inference(x)
    return torch.cat([preds, x["mask_coefficient"]], dim=1)

def forward(self, x: list[torch.Tensor]) -> tuple | list[torch.Tensor] | dict[str, torch.Tensor]:
    """Return model outputs and mask coefficients if training, otherwise return outputs and mask coefficients."""
    outputs = super().forward(x)
    preds = outputs[1] if isinstance(outputs, tuple) else outputs
    proto = self.proto(x[0])  # mask protos
    if isinstance(preds, dict):  # training and validating during training
        if self.end2end:
            preds["one2many"]["proto"] = proto
            preds["one2one"]["proto"] = proto.detach()
        else:
            preds["proto"] = proto
    if self.training:
        return preds
    return (outputs, proto) if self.export else ((outputs[0], proto), preds)

def forward_head(
    self, x: list[torch.Tensor], box_head: torch.nn.Module, cls_head: torch.nn.Module, mask_head: torch.nn.Module
) -> torch.Tensor:
    """Concatenates and returns predicted bounding boxes, class probabilities, and mask coefficients."""
    preds = super().forward_head(x, box_head, cls_head)
    if mask_head is not None:
        bs = x[0].shape[0]  # batch size
        preds["mask_coefficient"] = torch.cat([mask_head[i](x[i]).view(bs, self.nm, -1) for i in range(self.nl)], 2)
    return preds

def fuse(self) -> None:
    """Remove the one2many head for inference optimization."""
    self.cv2 = self.cv3 = self.cv4 = None

def postprocess(self, preds: torch.Tensor) -> torch.Tensor:
    """Post-process YOLO model predictions.

    Args:
        preds (torch.Tensor): Raw predictions with shape (batch_size, num_anchors, 4 + nc + nm) with last dimension
            format [x, y, w, h, class_probs, mask_coefficient].

    Returns:
        (torch.Tensor): Processed predictions with shape (batch_size, min(max_det, num_anchors), 6 + nm) and last
            dimension format [x, y, w, h, max_class_prob, class_index, mask_coefficient].
    """
    boxes, scores, mask_coefficient = preds.split([4, self.nc, self.nm], dim=-1)
    scores, conf, idx = self.get_topk_index(scores, self.max_det)
    boxes = boxes.gather(dim=1, index=idx.repeat(1, 1, 4))
    mask_coefficient = mask_coefficient.gather(dim=1, index=idx.repeat(1, 1, self.nm))
    return torch.cat([boxes, scores, conf, mask_coefficient], dim=-1)

class Segment26(Segment):
    """YOLO26 Segment head for segmentation models.

    This class extends the Detect head to include mask prediction capabilities for instance segmentation tasks.

    Attributes:
        nm (int): Number of masks.
        npr (int): Number of protos.
        proto (Proto): Prototype generation module.
        cv4 (nn.ModuleList): Convolution layers for mask coefficients.

    Methods:
        forward: Return model outputs and mask coefficients.

    Examples:
        Create a segmentation head
        >>> segment = Segment26(nc=80, nm=32, npr=256, ch=(256, 512, 1024))
        >>> x = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
        >>> outputs = segment(x)
    """

    def __init__(self, nc: int = 80, nm: int = 32, npr: int = 256, reg_max=16, end2end=False, ch: tuple = ()):
        """Initialize the YOLO model attributes such as the number of masks, prototypes, and the convolution layers.

        Args:
            nc (int): Number of classes.
            nm (int): Number of masks.
            npr (int): Number of protos.
            reg_max (int): Maximum number of DFL channels.
            end2end (bool): Whether to use end-to-end NMS-free detection.
            ch (tuple): Tuple of channel sizes from backbone feature maps.
        """
        super().__init__(nc, nm, npr, reg_max, end2end, ch)
        self.proto = Proto26(ch, self.npr, self.nm, nc)  # protos

def forward(self, x: list[torch.Tensor]) -> tuple | list[torch.Tensor] | dict[str, torch.Tensor]:
    """Return model outputs and mask coefficients if training, otherwise return outputs and mask coefficients."""
    outputs = Detect.forward(self, x)
    preds = outputs[1] if isinstance(outputs, tuple) else outputs
    proto = self.proto(x)  # mask protos
    if isinstance(preds, dict):  # training and validating during training
        if self.end2end:
            preds["one2many"]["proto"] = proto
            preds["one2one"]["proto"] = (
                tuple(p.detach() for p in proto) if isinstance(proto, tuple) else proto.detach()
            )
        else:
            preds["proto"] = proto
    if self.training:
        return preds
    return (outputs, proto) if self.export else ((outputs[0], proto), preds)

def fuse(self) -> None:
    """Remove the one2many head and extra part of proto module for inference optimization."""
    super().fuse()
    if hasattr(self.proto, "fuse"):
        self.proto.fuse()

class OBB(Detect):
    """YOLO OBB detection head for detection with rotation models.

    This class extends the Detect head to include oriented bounding box prediction with rotation angles.

    Attributes:
        ne (int): Number of extra parameters.
        cv4 (nn.ModuleList): Convolution layers for angle prediction.
        angle (torch.Tensor): Predicted rotation angles.

    Methods:
        forward: Concatenate and return predicted bounding boxes and class probabilities.
        decode_bboxes: Decode rotated bounding boxes.

    Examples:
        Create an OBB detection head
        >>> obb = OBB(nc=80, ne=1, ch=(256, 512, 1024))
        >>> x = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
        >>> outputs = obb(x)
    """

    def __init__(self, nc: int = 80, ne: int = 1, reg_max=16, end2end=False, ch: tuple = ()):
        """Initialize OBB with number of classes `nc` and layer channels `ch`.

        Args:
            nc (int): Number of classes.
            ne (int): Number of extra parameters.
            reg_max (int): Maximum number of DFL channels.
            end2end (bool): Whether to use end-to-end NMS-free detection.
            ch (tuple): Tuple of channel sizes from backbone feature maps.
        """
        super().__init__(nc, reg_max, end2end, ch)
        self.ne = ne  # number of extra parameters

        c4 = max(ch[0] // 4, self.ne)
        self.cv4 = nn.ModuleList(nn.Sequential(Conv(x, c4, 3), Conv(c4, c4, 3), nn.Conv2d(c4, self.ne, 1)) for x in ch)
        if end2end:
            self.one2one_cv4 = copy.deepcopy(self.cv4)

@property
def one2many(self):
    """Returns the one-to-many head components, here for backward compatibility."""
    return dict(box_head=self.cv2, cls_head=self.cv3, angle_head=self.cv4)

@property
def one2one(self):
    """Returns the one-to-one head components."""
    return dict(box_head=self.one2one_cv2, cls_head=self.one2one_cv3, angle_head=self.one2one_cv4)

def _inference(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
    """Decode predicted bounding boxes and class probabilities, concatenated with rotation angles."""
    # For decode_bboxes convenience
    self.angle = x["angle"]  # TODO: need to test obb
    preds = super()._inference(x)
    return torch.cat([preds, x["angle"]], dim=1)

def decode_bboxes(self, bboxes: torch.Tensor, anchors: torch.Tensor) -> torch.Tensor:
    """Decode rotated bounding boxes."""
    return dist2rbox(bboxes, self.angle, anchors, dim=1)

def forward_head(
    self, x: list[torch.Tensor], box_head: torch.nn.Module, cls_head: torch.nn.Module, angle_head: torch.nn.Module
) -> torch.Tensor:
    """Concatenates and returns predicted bounding boxes, class probabilities, and angles."""
    preds = super().forward_head(x, box_head, cls_head)
    if angle_head is not None:
        bs = x[0].shape[0]  # batch size
        angle = torch.cat(
            [angle_head[i](x[i]).view(bs, self.ne, -1) for i in range(self.nl)], 2
        )  # OBB theta logits
        angle = (angle.sigmoid() - 0.25) * math.pi  # [-pi/4, 3pi/4]
        preds["angle"] = angle
    return preds

def fuse(self) -> None:
    """Remove the one2many head for inference optimization."""
    self.cv2 = self.cv3 = self.cv4 = None

def postprocess(self, preds: torch.Tensor) -> torch.Tensor:
    """Post-process YOLO model predictions.

    Args:
        preds (torch.Tensor): Raw predictions with shape (batch_size, num_anchors, 4 + nc + ne) with last dimension
            format [x, y, w, h, class_probs, angle].

    Returns:
        (torch.Tensor): Processed predictions with shape (batch_size, min(max_det, num_anchors), 7) and last
            dimension format [x, y, w, h, max_class_prob, class_index, angle].
    """
    boxes, scores, angle = preds.split([4, self.nc, self.ne], dim=-1)
    scores, conf, idx = self.get_topk_index(scores, self.max_det)
    boxes = boxes.gather(dim=1, index=idx.repeat(1, 1, 4))
    angle = angle.gather(dim=1, index=idx.repeat(1, 1, self.ne))
    return torch.cat([boxes, scores, conf, angle], dim=-1)

class OBB26(OBB):

def forward_head(
    self, x: list[torch.Tensor], box_head: torch.nn.Module, cls_head: torch.nn.Module, angle_head: torch.nn.Module
) -> torch.Tensor:
    """Concatenates and returns predicted bounding boxes, class probabilities, and raw angles."""
    preds = Detect.forward_head(self, x, box_head, cls_head)
    if angle_head is not None:
        bs = x[0].shape[0]  # batch size
        angle = torch.cat(
            [angle_head[i](x[i]).view(bs, self.ne, -1) for i in range(self.nl)], 2
        )  # OBB theta logits (raw output without sigmoid transformation)
        preds["angle"] = angle
    return preds

class Pose(Detect):
    """YOLO Pose head for keypoints models.

    This class extends the Detect head to include keypoint prediction capabilities for pose estimation tasks.

    Attributes:
        kpt_shape (tuple): Number of keypoints and dimensions (2 for x,y or 3 for x,y,visible).
        nk (int): Total number of keypoint values.
        cv4 (nn.ModuleList): Convolution layers for keypoint prediction.

    Methods:
        forward: Perform forward pass through YOLO model and return predictions.
        kpts_decode: Decode keypoints from predictions.

    Examples:
        Create a pose detection head
        >>> pose = Pose(nc=80, kpt_shape=(17, 3), ch=(256, 512, 1024))
        >>> x = [torch.randn(1, 256, 80, 80), torch.randn(1, 512, 40, 40), torch.randn(1, 1024, 20, 20)]
        >>> outputs = pose(x)
    """

    def __init__(self, nc: int = 80, kpt_shape: tuple = (17, 3), reg_max=16, end2end=False, ch: tuple = ()):
        """Initialize YOLO network with default parameters and Convolutional Layers.

        Args:
            nc (int): Number of classes.
            kpt_shape (tuple): Number of keypoints, number of dims (2 for x,y or 3 for x,y,visible).
            reg_max (int): Maximum number of DFL channels.
            end2end (bool): Whether to use end-to-end NMS-free detection.
            ch (tuple): Tuple of channel sizes from backbone feature maps.
        """
        super().__init__(nc, reg_max, end2end, ch)
        self.kpt_shape = kpt_shape  # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
        self.nk = kpt_shape[0] * kpt_shape[1]  # number of keypoints total

        c4 = max(ch[0] // 4, self.nk)
        self.cv4 = nn.ModuleList(nn.Sequential(Conv(x, c4, 3), Conv(c4, c4, 3), nn.Conv2d(c4, self.nk, 1)) for x in ch)
        if end2end:
            self.one2one_cv4 = copy.deepcopy(self.cv4)

@property
def one2many(self):
    """Returns the one-to-many head components, here for backward compatibility."""
    return dict(box_head=self.cv2, cls_head=self.cv3, pose_head=self.cv4)

@property
def one2one(self):
    """Returns the one-to-one head components."""
    return dict(box_head=self.one2one_cv2, cls_head=self.one2one_cv3, pose_head=self.one2one_cv4)

def _inference(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
    """Decode predicted bounding boxes and class probabilities, concatenated with keypoints."""
    preds = super()._inference(x)
    return torch.cat([preds, self.kpts_decode(x["kpts"])], dim=1)

def forward_head(
    self, x: list[torch.Tensor], box_head: torch.nn.Module, cls_head: torch.nn.Module, pose_head: torch.nn.Module
) -> torch.Tensor:
    """Concatenates and returns predicted bounding boxes, class probabilities, and keypoints."""
    preds = super().forward_head(x, box_head, cls_head)
    if pose_head is not None:
        bs = x[0].shape[0]  # batch size
        preds["kpts"] = torch.cat([pose_head[i](x[i]).view(bs, self.nk, -1) for i in range(self.nl)], 2)
    return preds