models/custom_yolo.py

"""Register custom modules and provide a parse_model with direct custom-layer support."""

from __future__ import annotations

import ast
import math
import contextlib
from typing import Any

import torch
import torch.nn as nn

from ultralytics.nn.modules import (
    AIFI,
    C1,
    C2,
    C2PSA,
    C3,
    C3TR,
    ELAN1,
    OBB,
    OBB26,
    PSA,
    SPP,
    SPPELAN,
    SPPF,
    A2C2f,
    AConv,
    ADown,
    Bottleneck,
    BottleneckCSP,
    C2f,
    C2fAttn,
    C2fCIB,
    C2fPSA,
    C3Ghost,
    C3k2,
    C3x,
    CBFuse,
    CBLinear,
    Classify,
    Concat,
    Conv,
    ConvTranspose,
    Detect,
    DWConv,
    DWConvTranspose2d,
    Focus,
    GhostBottleneck,
    GhostConv,
    HGBlock,
    HGStem,
    ImagePoolingAttn,
    Index,
    Pose,
    Pose26,
    RepC3,
    RepNCSPELAN4,
    ResNetLayer,
    RTDETRDecoder,
    SCDown,
    Segment,
    Segment26,
    TorchVision,
    WorldDetect,
    YOLOEDetect,
    YOLOESegment,
    YOLOESegment26,
    v10Detect,
)
from ultralytics.utils import LOGGER, colorstr
from ultralytics.utils.ops import make_divisible

from modules import (
    CapsAlign,
    CapsDecode,
    CapsProj,
    CapsRoute,
    CapsRoutev2,
    CapsRoutev3,
    CapsRoutev4,
    CapsuleDetect,
    CapsuleDetectv1,
    CapsuleDetectv2,
    CapsuleDetectv4,
    CapsuleDetectv5,
    CapsuleDetectv6,
    CapsuleDetectv7,
    CapsuleDetectv8,
    CapsuleOpenVocabDetect,
    CapsuleSegmentv1,
    CapsuleSegmentv2,
    CapsuleSegmentv3,
    CapsuleTap,
    DeformableCapsBlock,
)


CUSTOM_MODULES = {
    "DeformableCapsBlock": DeformableCapsBlock,
    "CapsuleDetect": CapsuleDetect,
    "CapsuleDetectv1": CapsuleDetectv1,
    "CapsuleDetectv2": CapsuleDetectv2,
    "CapsuleDetectv4": CapsuleDetectv4,
    "CapsuleDetectv5": CapsuleDetectv5,
    "CapsuleDetectv6": CapsuleDetectv6,
    "CapsuleDetectv7": CapsuleDetectv7,
    "CapsuleDetectv8": CapsuleDetectv8,
    "CapsuleOpenVocabDetect": CapsuleOpenVocabDetect,
    "CapsuleSegmentv1": CapsuleSegmentv1,
    "CapsuleSegmentv2": CapsuleSegmentv2,
    "CapsuleSegmentv3": CapsuleSegmentv3,
    "CapsProj": CapsProj,
    "CapsAlign": CapsAlign,
    "CapsRoute": CapsRoute,
    "CapsRoutev2": CapsRoutev2,
    "CapsRoutev3": CapsRoutev3,
    "CapsRoutev4": CapsRoutev4,
    "CapsDecode": CapsDecode,
    "CapsuleTap": CapsuleTap,
}


def parse_model(d: dict[str, Any], ch: int, verbose: bool = True):
    """Parse a model.yaml dictionary into a PyTorch model with direct custom-module support."""
    legacy = True
    max_channels = float("inf")
    nc, act, scales, end2end = (d.get(x) for x in ("nc", "activation", "scales", "end2end"))
    reg_max = d.get("reg_max", 16)
    depth, width = (d.get(x, 1.0) for x in ("depth_multiple", "width_multiple"))
    scale = d.get("scale")
    if scales:
        if not scale:
            scale = next(iter(scales.keys()))
            LOGGER.warning(f"no model scale passed. Assuming scale='{scale}'.")
        depth, width, max_channels = scales[scale]

    if act:
        Conv.default_act = eval(act)
        if verbose:
            LOGGER.info(f"{colorstr('activation:')} {act}")

    if verbose:
        LOGGER.info(f"\n{'':>3}{'from':>20}{'n':>3}{'params':>10}  {'module':<45}{'arguments':<30}")

    ch = [ch]
    layers, save = [], []

    base_modules = frozenset(
        {
            Classify,
            Conv,
            ConvTranspose,
            GhostConv,
            Bottleneck,
            GhostBottleneck,
            SPP,
            SPPF,
            C2fPSA,
            C2PSA,
            DWConv,
            Focus,
            BottleneckCSP,
            C1,
            C2,
            C2f,
            C3k2,
            RepNCSPELAN4,
            ELAN1,
            ADown,
            AConv,
            SPPELAN,
            C2fAttn,
            C3,
            C3TR,
            C3Ghost,
            torch.nn.ConvTranspose2d,
            DWConvTranspose2d,
            C3x,
            RepC3,
            PSA,
            SCDown,
            C2fCIB,
            A2C2f,
        }
    )
    repeat_modules = frozenset(
        {
            BottleneckCSP,
            C1,
            C2,
            C2f,
            C3k2,
            C2fAttn,
            C3,
            C3TR,
            C3Ghost,
            C3x,
            RepC3,
            C2fPSA,
            C2fCIB,
            C2PSA,
            A2C2f,
        }
    )

    detect_modules = frozenset(
        {
            Detect,
            CapsuleDetect,
            CapsuleDetectv1,
            CapsuleDetectv2,
            CapsuleDetectv4,
            CapsuleDetectv5,
            CapsuleDetectv6,
            CapsuleDetectv7,
            CapsuleDetectv8,
    CapsuleDetectv8,
            CapsuleOpenVocabDetect,
            CapsuleSegmentv1,
            CapsuleSegmentv2,
            CapsuleSegmentv3,
    CapsuleSegmentv3,
            WorldDetect,
            YOLOEDetect,
            Segment,
            Segment26,
            YOLOESegment,
            YOLOESegment26,
            Pose,
            Pose26,
            OBB,
            OBB26,
        }
    )

    for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]):
        m = (
            getattr(torch.nn, m[3:])
            if "nn." in m
            else getattr(__import__("torchvision").ops, m[16:])
            if "torchvision.ops." in m
            else globals()[m]
        )

        for j, a in enumerate(args):
            if isinstance(a, str):
                with contextlib.suppress(ValueError):
                    args[j] = locals()[a] if a in locals() else ast.literal_eval(a)

        n = n_ = max(round(n * depth), 1) if n > 1 else n
        c2 = None

        if m in base_modules:
            c1, c2 = ch[f], args[0]
            if c2 != nc:
                c2 = make_divisible(min(c2, max_channels) * width, 8)
            if m is C2fAttn:
                args[1] = make_divisible(min(args[1], max_channels // 2) * width, 8)
                args[2] = int(max(round(min(args[2], max_channels // 2 // 32)) * width, 1) if args[2] > 1 else args[2])

            args = [c1, c2, *args[1:]]
            if m in repeat_modules:
                args.insert(2, n)
                n = 1
            if m is C3k2:
                legacy = False
                if scale in "mlx":
                    args[3] = True
            if m is A2C2f:
                legacy = False
                if scale in "lx":
                    args.extend((True, 1.2))
            if m is C2fCIB:
                legacy = False

        elif m is CapsProj:
            c1 = ch[f]
            k_base = int(args[0]) if len(args) > 0 else 4
            # Width scaling for capsule type count keeps model-size behavior aligned with YOLO scales.
            k = max(int(round(k_base * width)), 1)
            d_caps = int(args[1]) if len(args) > 1 else 16
            args = [c1, k, d_caps]
            c2 = k * (d_caps + 1)

        elif m is CapsAlign:
            c1 = ch[f]
            if len(args) < 2:
                raise ValueError('CapsAlign args must be [src_level, tgt_level, (down_groups)].')
            src_level, tgt_level = int(args[0]), int(args[1])
            if len(args) > 2:
                dg_base = int(args[2])
                group_num = max(int(round(dg_base * width)), 1)
                # keep grouped-conv valid: down_groups must divide channels
            else:
                group_num = 1
            args = [c1, src_level, tgt_level, group_num]
            c2 = c1

        elif m in {CapsRoute, CapsRoutev2, CapsRoutev3, CapsRoutev4}:
            num_src = len(f) if isinstance(f, (list, tuple)) else 1

            # Preferred YAML args:
            #   [K_in_list, P_in_list, K_out, P_out]
            #   [K_in_list, P_in_list, K_out, P_out, kernel_size, pre_k, post_k, pre_groups, post_groups]
            # Legacy support:
            #   [P_in, K_out, P_out]
            #   [P_in]
            if len(args) >= 4:
                K_in_raw, P_in_raw, K_out, P_out = args[0], args[1], int(args[2]), int(args[3])
                kernel_size = int(args[4]) if len(args) > 4 else 3
                pre_k = int(args[5]) if len(args) > 5 else 3
                post_k = int(args[6]) if len(args) > 6 else 3
                pre_groups_raw = int(args[7]) if len(args) > 7 else 0
                post_groups_raw = int(args[8]) if len(args) > 8 else 0
            elif len(args) == 3:
                P_in_raw, K_out, P_out = int(args[0]), int(args[1]), int(args[2])
                K_in_raw = 1
                kernel_size, pre_k, post_k = 3, 3, 3
                pre_groups_raw, post_groups_raw = 0, 0
            elif len(args) == 1:
                P_in_raw = int(args[0])
                K_in_raw, K_out, P_out = 1, 1, int(P_in_raw)
                kernel_size, pre_k, post_k = 3, 3, 3
                pre_groups_raw, post_groups_raw = 0, 0
            else:
                raise ValueError('CapsRoute/CapsRoutev2 args must be [K_in,P_in,K_out,P_out,(kernel_size,pre_k,post_k,pre_groups,post_groups)] or legacy [P_in,(K_out,P_out)].')

            if isinstance(K_in_raw, (list, tuple)):
                K_in_base = [int(v) for v in K_in_raw]
            else:
                K_in_base = [int(K_in_raw)] * num_src

            if isinstance(P_in_raw, (list, tuple)):
                P_in = [int(v) for v in P_in_raw]
            else:
                P_in = [int(P_in_raw)] * num_src

            if len(K_in_base) != num_src or len(P_in) != num_src:
                raise ValueError('CapsRoute/CapsRoutev2 K_in/P_in lists must match number of sources.')

            # Width scaling follows Ultralytics scale.width behavior.
            K_in = [max(int(round(k * width)), 1) for k in K_in_base]
            K_out = max(int(round(K_out * width)), 1)

            pre_groups = None
            if pre_groups_raw > 0:
                pre_groups = max(int(round(pre_groups_raw * width)), 1)
            post_groups = None
            if post_groups_raw > 0:
                post_groups = max(int(round(post_groups_raw * width)), 1)

            args = [K_in, P_in, K_out, int(P_out), kernel_size, pre_k, post_k, pre_groups, post_groups]
            c2 = K_out * (int(P_out) + 1)

        elif m is CapsDecode:
            c1 = ch[f]
            c2 = int(args[0]) if len(args) else c1
            c2 = make_divisible(min(c2, max_channels) * width, 8)
            args = [c1, c2]

        elif m is CapsuleTap:
            c2 = ch[f]

        elif m is AIFI:
            args = [ch[f], *args]
        elif m in frozenset({HGStem, HGBlock}):
            c1, cm, c2 = ch[f], args[0], args[1]
            args = [c1, cm, c2, *args[2:]]
            if m is HGBlock:
                args.insert(4, n)
                n = 1
        elif m is ResNetLayer:
            c2 = args[1] if args[3] else args[1] * 4
        elif m is torch.nn.BatchNorm2d:
            args = [ch[f]]
            c2 = ch[f]
        elif m is Concat:
            c2 = sum(ch[x] for x in f)
        elif m in detect_modules:
            if m in {
                CapsuleDetect,
                CapsuleDetectv1,
                CapsuleDetectv2,
                CapsuleDetectv4,
                CapsuleDetectv5,
                CapsuleDetectv6,
                CapsuleDetectv7,
                CapsuleDetectv8,
            CapsuleDetectv8,
    CapsuleDetectv8,
                CapsuleOpenVocabDetect,
                CapsuleSegmentv1,
                CapsuleSegmentv2,
                CapsuleSegmentv3,
            CapsuleSegmentv3,
    CapsuleSegmentv3,
            }:
                if len(args) < 3:
                    raise ValueError('CapsuleDetect/CapsuleDetectv1/CapsuleDetectv2/CapsuleDetectv4/CapsuleDetectv5/CapsuleDetectv6/CapsuleDetectv7/CapsuleDetectv8/CapsuleOpenVocabDetect/CapsuleSegmentv1/CapsuleSegmentv2/CapsuleSegmentv3 args must include [nc, k_list, d_list].')
                if not isinstance(args[1], (list, tuple)) or not isinstance(args[2], (list, tuple)):
                    raise TypeError('CapsuleDetect/CapsuleDetectv1/CapsuleDetectv2/CapsuleDetectv4/CapsuleDetectv5/CapsuleDetectv6/CapsuleDetectv7/CapsuleDetectv8/CapsuleOpenVocabDetect/CapsuleSegmentv1/CapsuleSegmentv2/CapsuleSegmentv3 requires k_list and d_list in YAML.')
                # Width-scale capsule type counts per level; keep pose dims as provided.
                args[1] = [max(int(round(int(v) * width)), 1) for v in args[1]]
                args[2] = [int(v) for v in args[2]]

            if m is CapsuleOpenVocabDetect:
                # Keep YAML order aligned with the head API:
                # [nc, k_list, d_list, embed, with_act_gate, with_objectness_prior]
                args = [*args[:3], reg_max, end2end, *args[3:], [ch[x] for x in f]]
            else:
                args.extend([reg_max, end2end, [ch[x] for x in f]])
            if m in {Segment, Segment26, YOLOESegment, YOLOESegment26}:
                args[2] = make_divisible(min(args[2], max_channels) * width, 8)
            if m in {
                Detect,
                CapsuleDetect,
                CapsuleDetectv1,
                CapsuleDetectv2,
                CapsuleDetectv4,
                CapsuleDetectv5,
                CapsuleDetectv6,
                CapsuleDetectv7,
                CapsuleDetectv8,
            CapsuleDetectv8,
    CapsuleDetectv8,
                CapsuleOpenVocabDetect,
                CapsuleSegmentv1,
                CapsuleSegmentv2,
                CapsuleSegmentv3,
            CapsuleSegmentv3,
    CapsuleSegmentv3,
                YOLOEDetect,
                Segment,
                Segment26,
                YOLOESegment,
                YOLOESegment26,
                Pose,
                Pose26,
                OBB,
                OBB26,
            }:
                m.legacy = legacy
            c2 = ch[f[-1]] if isinstance(f, (list, tuple)) else ch[f]

        elif m is v10Detect:
            args.append([ch[x] for x in f])
            c2 = ch[f[-1]] if isinstance(f, (list, tuple)) else ch[f]
        elif m is ImagePoolingAttn:
            args.insert(1, [ch[x] for x in f])
            c2 = ch[f[-1]] if isinstance(f, (list, tuple)) else ch[f]
        elif m is RTDETRDecoder:
            args.insert(1, [ch[x] for x in f])
            c2 = ch[f[-1]] if isinstance(f, (list, tuple)) else ch[f]
        elif m is CBLinear:
            c2 = args[0]
            c1 = ch[f]
            args = [c1, c2, *args[1:]]
        elif m is CBFuse:
            c2 = ch[f[-1]]
        elif m in frozenset({TorchVision, Index}):
            c2 = args[0]
            args = [*args[1:]]
        else:
            c2 = ch[f]

        m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)

        if m in {CapsRoute, CapsRoutev2, CapsRoutev3, CapsRoutev4}:
            c2 = int(getattr(m_, "c_out", c2))

        t = str(m)[8:-2].replace("__main__.", "")
        m_.np = sum(x.numel() for x in m_.parameters())
        m_.i, m_.f, m_.type = i, f, t
        if verbose:
            LOGGER.info(f"{i:>3}{f!s:>20}{n_:>3}{m_.np:10.0f}  {t:<45}{args!s:<30}")

        save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)
        layers.append(m_)
        if i == 0:
            ch = []
        ch.append(c2)

    # Keep all intermediate outputs to avoid None entries for custom multi-source routing blocks.
    return nn.Sequential(*layers), sorted(set(save + list(range(len(layers)))))


def register_ultralytics_modules() -> None:
    """Register custom modules and replace Ultralytics parse_model with this custom parse_model."""
    import ultralytics.nn.modules as nn_modules
    import ultralytics.nn.tasks as nn_tasks

    for name, cls in CUSTOM_MODULES.items():
        setattr(nn_tasks, name, cls)
        setattr(nn_modules, name, cls)

    if getattr(nn_tasks, "_capsule_parse_patched", False):
        return

    nn_tasks.parse_model = parse_model
    nn_tasks._capsule_parse_patched = True