Reference for ultralytics/nn/modules/head.py

def bias_init(self):
    """Initialize Detect() biases, WARNING: requires stride availability."""
    m = self  # self.model[-1]  # Detect() module
    # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1
    # ncf = math.log(0.6 / (m.nc - 0.999999)) if cf is None else torch.log(cf / cf.sum())  # nominal class frequency
    for a, b, s in zip(m.cv2, m.cv3, m.stride):  # from
        a[-1].bias.data[:] = 1.0  # box
        b[-1].bias.data[: m.nc] = math.log(5 / m.nc / (640 / s) ** 2)  # cls (.01 objects, 80 classes, 640 img)
    if self.end2end:
        for a, b, s in zip(m.one2one_cv2, m.one2one_cv3, m.stride):  # from
            a[-1].bias.data[:] = 1.0  # box
            b[-1].bias.data[: m.nc] = math.log(5 / m.nc / (640 / s) ** 2)  # cls (.01 objects, 80 classes, 640 img)

def decode_bboxes(self, bboxes, anchors, xywh=True):
    """Decode bounding boxes."""
    return dist2bbox(bboxes, anchors, xywh=xywh and (not self.end2end), dim=1)

def forward(self, x):
    """Concatenates and returns predicted bounding boxes and class probabilities."""
    if self.end2end:
        return self.forward_end2end(x)

    for i in range(self.nl):
        x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
    if self.training:  # Training path
        return x
    y = self._inference(x)
    return y if self.export else (y, x)

def forward_end2end(self, x):
    """
    Performs forward pass of the v10Detect module.

    Args:
        x (List[torch.Tensor]): Input feature maps from different levels.

    Returns:
        (dict | tuple): If in training mode, returns a dictionary containing the outputs of both one2many and
            one2one detections. If not in training mode, returns processed detections or a tuple with
            processed detections and raw outputs.
    """
    x_detach = [xi.detach() for xi in x]
    one2one = [
        torch.cat((self.one2one_cv2[i](x_detach[i]), self.one2one_cv3[i](x_detach[i])), 1) for i in range(self.nl)
    ]
    for i in range(self.nl):
        x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
    if self.training:  # Training path
        return {"one2many": x, "one2one": one2one}

    y = self._inference(one2one)
    y = self.postprocess(y.permute(0, 2, 1), self.max_det, self.nc)
    return y if self.export else (y, {"one2many": x, "one2one": one2one})

@staticmethod
def postprocess(preds: torch.Tensor, max_det: int, nc: int = 80):
    """
    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].
        max_det (int): Maximum detections per image.
        nc (int, optional): Number of classes. Default: 80.

    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].
    """
    batch_size, anchors, _ = preds.shape  # i.e. shape(16,8400,84)
    boxes, scores = preds.split([4, nc], dim=-1)
    index = scores.amax(dim=-1).topk(min(max_det, anchors))[1].unsqueeze(-1)
    boxes = boxes.gather(dim=1, index=index.repeat(1, 1, 4))
    scores = scores.gather(dim=1, index=index.repeat(1, 1, nc))
    scores, index = scores.flatten(1).topk(min(max_det, anchors))
    i = torch.arange(batch_size)[..., None]  # batch indices
    return torch.cat([boxes[i, index // nc], scores[..., None], (index % nc)[..., None].float()], dim=-1)

def forward(self, x):
    """Return model outputs and mask coefficients if training, otherwise return outputs and mask coefficients."""
    p = self.proto(x[0])  # mask protos
    bs = p.shape[0]  # batch size

    mc = torch.cat([self.cv4[i](x[i]).view(bs, self.nm, -1) for i in range(self.nl)], 2)  # mask coefficients
    x = Detect.forward(self, x)
    if self.training:
        return x, mc, p
    return (torch.cat([x, mc], 1), p) if self.export else (torch.cat([x[0], mc], 1), (x[1], mc, p))

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

def forward(self, x):
    """Concatenates and returns predicted bounding boxes and class probabilities."""
    bs = x[0].shape[0]  # batch size
    angle = torch.cat([self.cv4[i](x[i]).view(bs, self.ne, -1) for i in range(self.nl)], 2)  # OBB theta logits
    # NOTE: set `angle` as an attribute so that `decode_bboxes` could use it.
    angle = (angle.sigmoid() - 0.25) * math.pi  # [-pi/4, 3pi/4]
    # angle = angle.sigmoid() * math.pi / 2  # [0, pi/2]
    if not self.training:
        self.angle = angle
    x = Detect.forward(self, x)
    if self.training:
        return x, angle
    return torch.cat([x, angle], 1) if self.export else (torch.cat([x[0], angle], 1), (x[1], angle))

def forward(self, x):
    """Perform forward pass through YOLO model and return predictions."""
    bs = x[0].shape[0]  # batch size
    kpt = torch.cat([self.cv4[i](x[i]).view(bs, self.nk, -1) for i in range(self.nl)], -1)  # (bs, 17*3, h*w)
    x = Detect.forward(self, x)
    if self.training:
        return x, kpt
    pred_kpt = self.kpts_decode(bs, kpt)
    return torch.cat([x, pred_kpt], 1) if self.export else (torch.cat([x[0], pred_kpt], 1), (x[1], kpt))

def kpts_decode(self, bs, kpts):
    """Decodes keypoints."""
    ndim = self.kpt_shape[1]
    if self.export:
        if self.format in {
            "tflite",
            "edgetpu",
        }:  # required for TFLite export to avoid 'PLACEHOLDER_FOR_GREATER_OP_CODES' bug
            # Precompute normalization factor to increase numerical stability
            y = kpts.view(bs, *self.kpt_shape, -1)
            grid_h, grid_w = self.shape[2], self.shape[3]
            grid_size = torch.tensor([grid_w, grid_h], device=y.device).reshape(1, 2, 1)
            norm = self.strides / (self.stride[0] * grid_size)
            a = (y[:, :, :2] * 2.0 + (self.anchors - 0.5)) * norm
        else:
            # NCNN fix
            y = kpts.view(bs, *self.kpt_shape, -1)
            a = (y[:, :, :2] * 2.0 + (self.anchors - 0.5)) * self.strides
        if ndim == 3:
            a = torch.cat((a, y[:, :, 2:3].sigmoid()), 2)
        return a.view(bs, self.nk, -1)
    else:
        y = kpts.clone()
        if ndim == 3:
            y[:, 2::ndim] = y[:, 2::ndim].sigmoid()  # sigmoid (WARNING: inplace .sigmoid_() Apple MPS bug)
        y[:, 0::ndim] = (y[:, 0::ndim] * 2.0 + (self.anchors[0] - 0.5)) * self.strides
        y[:, 1::ndim] = (y[:, 1::ndim] * 2.0 + (self.anchors[1] - 0.5)) * self.strides
        return y

def forward(self, x):
    """Performs a forward pass of the YOLO model on input image data."""
    if isinstance(x, list):
        x = torch.cat(x, 1)
    x = self.linear(self.drop(self.pool(self.conv(x)).flatten(1)))
    if self.training:
        return x
    y = x.softmax(1)  # get final output
    return y if self.export else (y, x)

def bias_init(self):
    """Initialize Detect() biases, WARNING: requires stride availability."""
    m = self  # self.model[-1]  # Detect() module
    # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1
    # ncf = math.log(0.6 / (m.nc - 0.999999)) if cf is None else torch.log(cf / cf.sum())  # nominal class frequency
    for a, b, s in zip(m.cv2, m.cv3, m.stride):  # from
        a[-1].bias.data[:] = 1.0  # box

def forward(self, x, text):
    """Concatenates and returns predicted bounding boxes and class probabilities."""
    for i in range(self.nl):
        x[i] = torch.cat((self.cv2[i](x[i]), self.cv4[i](self.cv3[i](x[i]), text)), 1)
    if self.training:
        return x
    self.no = self.nc + self.reg_max * 4  # self.nc could be changed when inference with different texts
    y = self._inference(x)
    return y if self.export else (y, x)

def conv2linear(self, conv):
    """Convert a 1x1 convolutional layer to a linear layer."""
    assert isinstance(conv, nn.Conv2d) and conv.kernel_size == (1, 1)
    linear = nn.Linear(conv.in_channels, conv.out_channels)
    linear.weight.data = conv.weight.view(conv.out_channels, -1).data
    linear.bias.data = conv.bias.data
    return linear

def forward(self, cls_feat, loc_feat, conf):
    """Process classification and localization features to generate detection proposals."""
    if self.enabled:
        pf_score = self.pf(cls_feat)[0, 0].flatten(0)
        mask = pf_score.sigmoid() > conf
        cls_feat = cls_feat.flatten(2).transpose(-1, -2)
        cls_feat = self.vocab(cls_feat * mask.unsqueeze(-1).int() if not conf else cls_feat[:, mask])
        return (self.loc(loc_feat), cls_feat.transpose(-1, -2)), mask
    else:
        cls_feat = self.vocab(cls_feat)
        loc_feat = self.loc(loc_feat)
        return (loc_feat, cls_feat.flatten(2)), torch.ones(
            cls_feat.shape[2] * cls_feat.shape[3], device=cls_feat.device, dtype=torch.bool
        )

def bias_init(self):
    """Initialize biases for detection heads."""
    m = self  # self.model[-1]  # Detect() module
    # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1
    # ncf = math.log(0.6 / (m.nc - 0.999999)) if cf is None else torch.log(cf / cf.sum())  # nominal class frequency
    for a, b, c, s in zip(m.cv2, m.cv3, m.cv4, m.stride):  # from
        a[-1].bias.data[:] = 1.0  # box
        # b[-1].bias.data[:] = math.log(5 / m.nc / (640 / s) ** 2)  # cls (.01 objects, 80 classes, 640 img)
        b[-1].bias.data[:] = 0.0
        c.bias.data[:] = math.log(5 / m.nc / (640 / s) ** 2)

def forward(self, x, cls_pe, return_mask=False):
    """Process features with class prompt embeddings to generate detections."""
    if hasattr(self, "lrpc"):  # for prompt-free inference
        return self.forward_lrpc(x, return_mask)
    for i in range(self.nl):
        x[i] = torch.cat((self.cv2[i](x[i]), self.cv4[i](self.cv3[i](x[i]), cls_pe)), 1)
    if self.training:
        return x
    self.no = self.nc + self.reg_max * 4  # self.nc could be changed when inference with different texts
    y = self._inference(x)
    return y if self.export else (y, x)

def forward_lrpc(self, x, return_mask=False):
    """Process features with fused text embeddings to generate detections for prompt-free model."""
    masks = []
    assert self.is_fused, "Prompt-free inference requires model to be fused!"
    for i in range(self.nl):
        cls_feat = self.cv3[i](x[i])
        loc_feat = self.cv2[i](x[i])
        assert isinstance(self.lrpc[i], LRPCHead)
        x[i], mask = self.lrpc[i](
            cls_feat, loc_feat, 0 if self.export and not self.dynamic else getattr(self, "conf", 0.001)
        )
        masks.append(mask)
    shape = x[0][0].shape
    if self.dynamic or self.shape != shape:
        self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors([b[0] for b in x], self.stride, 0.5))
        self.shape = shape
    box = torch.cat([xi[0].view(shape[0], self.reg_max * 4, -1) for xi in x], 2)
    cls = torch.cat([xi[1] for xi in x], 2)

    if self.export and self.format in {"tflite", "edgetpu"}:
        # Precompute normalization factor to increase numerical stability
        # See https://github.com/ultralytics/ultralytics/issues/7371
        grid_h = shape[2]
        grid_w = shape[3]
        grid_size = torch.tensor([grid_w, grid_h, grid_w, grid_h], device=box.device).reshape(1, 4, 1)
        norm = self.strides / (self.stride[0] * grid_size)
        dbox = self.decode_bboxes(self.dfl(box) * norm, self.anchors.unsqueeze(0) * norm[:, :2])
    else:
        dbox = self.decode_bboxes(self.dfl(box), self.anchors.unsqueeze(0)) * self.strides

    mask = torch.cat(masks)
    y = torch.cat((dbox if self.export and not self.dynamic else dbox[..., mask], cls.sigmoid()), 1)

    if return_mask:
        return (y, mask) if self.export else ((y, x), mask)
    else:
        return y if self.export else (y, x)

@smart_inference_mode()
def fuse(self, txt_feats):
    """Fuse text features with model weights for efficient inference."""
    if self.is_fused:
        return

    assert not self.training
    txt_feats = txt_feats.to(torch.float32).squeeze(0)
    for cls_head, bn_head in zip(self.cv3, self.cv4):
        assert isinstance(cls_head, nn.Sequential)
        assert isinstance(bn_head, BNContrastiveHead)
        conv = cls_head[-1]
        assert isinstance(conv, nn.Conv2d)
        logit_scale = bn_head.logit_scale
        bias = bn_head.bias
        norm = bn_head.norm

        t = txt_feats * logit_scale.exp()
        conv: nn.Conv2d = fuse_conv_and_bn(conv, norm)

        w = conv.weight.data.squeeze(-1).squeeze(-1)
        b = conv.bias.data

        w = t @ w
        b1 = (t @ b.reshape(-1).unsqueeze(-1)).squeeze(-1)
        b2 = torch.ones_like(b1) * bias

        conv = (
            nn.Conv2d(
                conv.in_channels,
                w.shape[0],
                kernel_size=1,
            )
            .requires_grad_(False)
            .to(conv.weight.device)
        )

        conv.weight.data.copy_(w.unsqueeze(-1).unsqueeze(-1))
        conv.bias.data.copy_(b1 + b2)
        cls_head[-1] = conv

        bn_head.fuse()

    del self.reprta
    self.reprta = nn.Identity()
    self.is_fused = True

def get_tpe(self, tpe):
    """Get text prompt embeddings with normalization."""
    return None if tpe is None else F.normalize(self.reprta(tpe), dim=-1, p=2)

def get_vpe(self, x, vpe):
    """Get visual prompt embeddings with spatial awareness."""
    if vpe.shape[1] == 0:  # no visual prompt embeddings
        return torch.zeros(x[0].shape[0], 0, self.embed, device=x[0].device)
    if vpe.ndim == 4:  # (B, N, H, W)
        vpe = self.savpe(x, vpe)
    assert vpe.ndim == 3  # (B, N, D)
    return vpe

def forward(self, x, text):
    """Return model outputs and mask coefficients if training, otherwise return outputs and mask coefficients."""
    p = self.proto(x[0])  # mask protos
    bs = p.shape[0]  # batch size

    mc = torch.cat([self.cv5[i](x[i]).view(bs, self.nm, -1) for i in range(self.nl)], 2)  # mask coefficients
    has_lrpc = hasattr(self, "lrpc")

    if not has_lrpc:
        x = YOLOEDetect.forward(self, x, text)
    else:
        x, mask = YOLOEDetect.forward(self, x, text, return_mask=True)

    if self.training:
        return x, mc, p

    if has_lrpc:
        mc = (mc * mask.int()) if self.export and not self.dynamic else mc[..., mask]

    return (torch.cat([x, mc], 1), p) if self.export else (torch.cat([x[0], mc], 1), (x[1], mc, p))