init project

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

.gitignore

@@ -0,0 +1,28 @@
+.claude*
+
+# Python
+__pycache__/
+*.py[cod]
+*.pyo
+.venv/
+*.egg-info/
+
+# PyTorch / model weights
+*.pt
+*.pth
+*.onnx
+*.torchscript
+
+# Ultralytics runs output
+runs/
+
+# Downloaded / temp images
+*.mp4
+*.avi
+
+# uv / pip
+.python-version
+
+# OS
+.DS_Store
+Thumbs.db

README.md

@@ -10,26 +10,12 @@ tags:
 - real-time
 - pytorch
 library_name: ultralytics
-pipeline_tag: object-detection
+pipeline_tag: image-segmentation
 datasets:
 - coco
 model-index:
 - name: CSN
   results:
-  - task:
-      type: object-detection
-      name: Object Detection
-    dataset:
-      name: COCO 2017
-      type: coco
-      split: val2017
-    metrics:
-    - type: mAP50-95
-      value: [FILL_ME]
-      name: Box mAP50-95
-    - type: mAP50
-      value: [FILL_ME]
-      name: Box mAP50
   - task:
       type: instance-segmentation
       name: Instance Segmentation
@@ -39,63 +25,83 @@ model-index:
       split: val2017
     metrics:
     - type: mAP50-95
-      value: [FILL_ME]
+      value: 0.34080
       name: Mask mAP50-95
     - type: mAP50
-      value: [FILL_ME]
+      value: 0.53316
       name: Mask mAP50
 ---
 
 # Symbolic Capsule Network
 
-**Symbolic Capsule Network (SCN)** is a capsule-enhanced real-time vision model for **COCO object detection** and **instance segmentation**.  
-This release provides two model variants:
+This repository contains the HuggingFace model page assets for the **Symbolic Capsule Network (SCN)** segmentation checkpoint trained on **COCO 2017**.
+
+The repo now includes the custom capsule modules required to load the checkpoint, plus a bundled default weight copied from:
+
+`D:\Codes\CSNet\runs\CapsNeckHeadv6_SegV2_AuxDet_Mild\weights\best.pt`
+
+The bundled checkpoint is stored at:
 
-- **Symbolic Capsule Network Detection**
-- **Symbolic Capsule Network Segmentation**
+`weights/symbolic_capsule_network_segmentation.pt`
 
-SCN introduces capsule-style structured feature representation into a YOLO-style dense prediction pipeline, aiming to improve semantic compositionality while preserving practical real-time performance.
+## Model
 
-## Available Models
+| Model | Task | Dataset | Input Size | Mask mAP50 | Mask mAP50-95 |
+|---|---|---|---:|---:|---:|
+| Symbolic Capsule Network Segmentation | Instance Segmentation | COCO 2017 | 640 | 0.53316 | 0.34080 |
 
-| Model | Task | Dataset | Input Size | Best Validation Performance |
-|---|---|---|---:|---|
-| **Symbolic Capsule Network Detection** | Object Detection | COCO 2017 | 640 | Box mAP50 **0.55776**, Box mAP50-95 **0.40319** |
-| **Symbolic Capsule Network Segmentation** | Instance Segmentation | COCO 2017 | 640 | Mask mAP50 **0.53316**, Mask mAP50-95 **0.34080** |
+## Files Included
 
-## Model Description
+- `configs/data/coco-seg.yaml`: COCO segmentation dataset yaml used by the bundled best checkpoint.
+- `configs/data/coco.yaml`: COCO detection dataset yaml retained for the paired detection setup.
+- `configs/seg_model/Capsneck_v6_seg_v2.yaml`: segmentation model yaml used by the bundled best checkpoint.
+- `configs/det_model/yolo26_capsneckhead_v6_gated.yaml`: paired detection model yaml for the v6 gated setup.
+- `modules/`: custom capsule backbone, neck, and head implementations.
+- `models/custom_yolo.py`: Ultralytics registration hook for custom layers.
+- `predict.py`: minimal inference entrypoint with the bundled checkpoint as default.
+- `weights/symbolic_capsule_network_segmentation.pt`: default segmentation checkpoint.
 
-Symbolic Capsule Network is designed for real-time dense visual prediction.  
-The model replaces part of the standard feature interaction pipeline with capsule-inspired transformations and structured routing, enabling richer intermediate representations for localization and mask prediction.
+## Environment
 
-This page focuses on two public COCO releases:
+```bash
+uv sync
+```
 
-1. **Detection model**
-2. **Segmentation model**
+This repository now ships the same `pyproject.toml`, `uv.lock`, and `.python-version` setup as the original `D:\Codes\CSNet` project.
 
-## Performance on COCO 2017
+The local `.venv` in this directory is linked to:
 
-### Detection
+`D:\Codes\CSNet\.venv`
 
-Best result from the current training record:
+So you can run `uv sync` here without creating a new virtual environment, and both directories will use the same existing environment.
 
-- **Box mAP50:** `0.55776`
-- **Box mAP50-95:** `0.40319`
+## Run Inference
 
-### Segmentation
+```bash
+python predict.py path/to/image.jpg
+```
 
-Best result from the current training record:
+By default, `predict.py` loads `weights/symbolic_capsule_network_segmentation.pt`.
 
-- **Mask mAP50:** `0.53316`
-- **Mask mAP50-95:** `0.34080`
+You can also override the checkpoint explicitly:
 
-## Usage
+```bash
+python predict.py path/to/image.jpg --weights path/to/other.pt
+```
 
-### Detection
+## Python Usage
 
 ```python
 from ultralytics import YOLO
 
-model = YOLO("symbolic_capsule_network_detection.pt")
-results = model("image.jpg")
-results[0].show()
+from models import register_ultralytics_modules
+
+register_ultralytics_modules()
+model = YOLO("weights/symbolic_capsule_network_segmentation.pt")
+results = model.predict("image.jpg", imgsz=640, conf=0.25)
+```
+
+## Notes
+
+- This checkpoint is a segmentation model, so masks are available in `results[0].masks`.
+- The custom code is required before loading the checkpoint with Ultralytics.

configs/data/coco-seg.yaml

@@ -0,0 +1,88 @@
+# Ultralytics AGPL-3.0 License - https://ultralytics.com/license
+# COCO 2017 instance segmentation dataset config
+
+path: C:\Datasets\coco-seg
+train: train2017.txt
+val: val2017.txt
+
+names:
+  0: person
+  1: bicycle
+  2: car
+  3: motorcycle
+  4: airplane
+  5: bus
+  6: train
+  7: truck
+  8: boat
+  9: traffic light
+  10: fire hydrant
+  11: stop sign
+  12: parking meter
+  13: bench
+  14: bird
+  15: cat
+  16: dog
+  17: horse
+  18: sheep
+  19: cow
+  20: elephant
+  21: bear
+  22: zebra
+  23: giraffe
+  24: backpack
+  25: umbrella
+  26: handbag
+  27: tie
+  28: suitcase
+  29: frisbee
+  30: skis
+  31: snowboard
+  32: sports ball
+  33: kite
+  34: baseball bat
+  35: baseball glove
+  36: skateboard
+  37: surfboard
+  38: tennis racket
+  39: bottle
+  40: wine glass
+  41: cup
+  42: fork
+  43: knife
+  44: spoon
+  45: bowl
+  46: banana
+  47: apple
+  48: sandwich
+  49: orange
+  50: broccoli
+  51: carrot
+  52: hot dog
+  53: pizza
+  54: donut
+  55: cake
+  56: chair
+  57: couch
+  58: potted plant
+  59: bed
+  60: dining table
+  61: toilet
+  62: tv
+  63: laptop
+  64: mouse
+  65: remote
+  66: keyboard
+  67: cell phone
+  68: microwave
+  69: oven
+  70: toaster
+  71: sink
+  72: refrigerator
+  73: book
+  74: clock
+  75: vase
+  76: scissors
+  77: teddy bear
+  78: hair drier
+  79: toothbrush

configs/data/coco.yaml

@@ -0,0 +1,117 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+# COCO 2017 dataset https://cocodataset.org by Microsoft
+# Documentation: https://docs.ultralytics.com/datasets/detect/coco/
+# Example usage: yolo train data=coco.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── coco ← downloads here (20.1 GB)
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: C:\Datasets\coco-yolo # dataset root dir
+train: train2017.txt # train images (relative to 'path') 118287 images
+val: val2017.txt # val images (relative to 'path') 5000 images
+
+# Classes
+names:
+  0: person
+  1: bicycle
+  2: car
+  3: motorcycle
+  4: airplane
+  5: bus
+  6: train
+  7: truck
+  8: boat
+  9: traffic light
+  10: fire hydrant
+  11: stop sign
+  12: parking meter
+  13: bench
+  14: bird
+  15: cat
+  16: dog
+  17: horse
+  18: sheep
+  19: cow
+  20: elephant
+  21: bear
+  22: zebra
+  23: giraffe
+  24: backpack
+  25: umbrella
+  26: handbag
+  27: tie
+  28: suitcase
+  29: frisbee
+  30: skis
+  31: snowboard
+  32: sports ball
+  33: kite
+  34: baseball bat
+  35: baseball glove
+  36: skateboard
+  37: surfboard
+  38: tennis racket
+  39: bottle
+  40: wine glass
+  41: cup
+  42: fork
+  43: knife
+  44: spoon
+  45: bowl
+  46: banana
+  47: apple
+  48: sandwich
+  49: orange
+  50: broccoli
+  51: carrot
+  52: hot dog
+  53: pizza
+  54: donut
+  55: cake
+  56: chair
+  57: couch
+  58: potted plant
+  59: bed
+  60: dining table
+  61: toilet
+  62: tv
+  63: laptop
+  64: mouse
+  65: remote
+  66: keyboard
+  67: cell phone
+  68: microwave
+  69: oven
+  70: toaster
+  71: sink
+  72: refrigerator
+  73: book
+  74: clock
+  75: vase
+  76: scissors
+  77: teddy bear
+  78: hair drier
+  79: toothbrush
+
+# Download script/URL (optional)
+download: |
+  from pathlib import Path
+
+  from ultralytics.utils import ASSETS_URL
+  from ultralytics.utils.downloads import download
+
+  # Download labels
+  segments = True  # segment or box labels
+  dir = Path(yaml["path"])  # dataset root dir
+  urls = [ASSETS_URL + ("/coco2017labels-segments.zip" if segments else "/coco2017labels.zip")]  # labels
+  download(urls, dir=dir.parent)
+  # Download data
+  urls = [
+      "http://images.cocodataset.org/zips/train2017.zip",  # 19G, 118k images
+      "http://images.cocodataset.org/zips/val2017.zip",  # 1G, 5k images
+      "http://images.cocodataset.org/zips/test2017.zip",  # 7G, 41k images (optional)
+  ]
+  download(urls, dir=dir / "images", threads=3)

configs/det_model/yolo26_capsneckhead_v6_gated.yaml

@@ -0,0 +1,47 @@
+# Ultralytics AGPL-3.0 License - https://ultralytics.com/license
+#
+# YOLO26 + capsule-style neck(v2 route) + CapsuleDetectv6 head
+# This variant keeps the v6 neck/backbone and switches the detect head to the
+# gated act modulation used by CapsuleSegmentv1 for seg -> det transfer.
+
+nc: 80
+end2end: True
+reg_max: 1
+scales:
+  n: [0.50, 0.25, 1024]
+  s: [0.50, 0.50, 1024]
+  m: [0.50, 1.00, 512]
+  l: [1.00, 1.00, 512]
+  x: [1.00, 1.50, 512]
+
+backbone:
+  - [-1, 1, Conv, [64, 3, 2]]
+  - [-1, 1, Conv, [128, 3, 2]]
+  - [-1, 2, C3k2, [256, False, 0.25]]
+  - [-1, 1, Conv, [256, 3, 2]]
+  - [-1, 2, C3k2, [512, False, 0.25]]
+  - [-1, 1, Conv, [512, 3, 2]]
+  - [-1, 2, C3k2, [512, True]]
+  - [-1, 1, Conv, [1024, 3, 2]]
+  - [-1, 2, C3k2, [1024, True]]
+  - [-1, 1, SPPF, [1024, 5, 3, True]]
+  - [-1, 2, C2PSA, [1024]]
+
+head:
+  - [4, 1, CapsProj, [16, 16]]
+  - [6, 1, CapsProj, [32, 16]]
+  - [10, 1, CapsProj, [64, 16]]
+
+  - [13, 1, CapsAlign, [5, 4]]
+  - [[14, 12], 1, CapsRoutev2, [[64, 32], [16, 16], 32, 16]]
+
+  - [15, 1, CapsAlign, [4, 3]]
+  - [[16, 11], 1, CapsRoutev2, [[32, 16], [16, 16], 16, 16]]
+
+  - [17, 1, CapsAlign, [3, 4, 16]]
+  - [[18, 15], 1, CapsRoutev2, [[16, 32], [16, 16], 32, 16]]
+
+  - [19, 1, CapsAlign, [4, 5, 32]]
+  - [[20, 13], 1, CapsRoutev2, [[32, 64], [16, 16], 64, 16]]
+
+  - [[17, 19, 21], 1, CapsuleDetectv6, [nc, [16, 32, 64], [16, 16, 16]]]

configs/seg_model/Capsneck_v6_seg_v2.yaml

@@ -0,0 +1,47 @@
+# Ultralytics AGPL-3.0 License - https://ultralytics.com/license
+#
+# YOLO26 + CapsNeck(v2 route) + CapsuleSegmentv2 head
+# Classification branch uses raw pose features plus cls_priors, without act gating.
+
+nc: 80
+end2end: True
+reg_max: 1
+scales:
+  n: [0.50, 0.25, 1024]
+  s: [0.50, 0.50, 1024]
+  m: [0.50, 1.00, 512]
+  l: [1.00, 1.00, 512]
+  x: [1.00, 1.50, 512]
+
+backbone:
+  - [-1, 1, Conv, [64, 3, 2]]
+  - [-1, 1, Conv, [128, 3, 2]]
+  - [-1, 2, C3k2, [256, False, 0.25]]
+  - [-1, 1, Conv, [256, 3, 2]]
+  - [-1, 2, C3k2, [512, False, 0.25]]
+  - [-1, 1, Conv, [512, 3, 2]]
+  - [-1, 2, C3k2, [512, True]]
+  - [-1, 1, Conv, [1024, 3, 2]]
+  - [-1, 2, C3k2, [1024, True]]
+  - [-1, 1, SPPF, [1024, 5, 3, True]]
+  - [-1, 2, C2PSA, [1024]]
+
+head:
+  - [4, 1, CapsProj, [16, 16]]
+  - [6, 1, CapsProj, [32, 16]]
+  - [10, 1, CapsProj, [64, 16]]
+
+  - [13, 1, CapsAlign, [5, 4]]
+  - [[14, 12], 1, CapsRoutev2, [[64, 32], [16, 16], 32, 16]]
+
+  - [15, 1, CapsAlign, [4, 3]]
+  - [[16, 11], 1, CapsRoutev2, [[32, 16], [16, 16], 16, 16]]
+
+  - [17, 1, CapsAlign, [3, 4, 16]]
+  - [[18, 15], 1, CapsRoutev2, [[16, 32], [16, 16], 32, 16]]
+
+  - [19, 1, CapsAlign, [4, 5, 32]]
+  - [[20, 13], 1, CapsRoutev2, [[32, 64], [16, 16], 64, 16]]
+
+  # [nc, k_list, d_list, nm, npr]
+  - [[17, 19, 21], 1, CapsuleSegmentv2, [nc, [16, 32, 64], [16, 16, 16], 32, 256]]

models/__init__.py

@@ -0,0 +1,3 @@
+from .custom_yolo import register_ultralytics_modules
+
+__all__ = ["register_ultralytics_modules"]

models/custom_yolo.py

@@ -0,0 +1,501 @@
+"""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

modules/__init__.py

@@ -0,0 +1,56 @@
+from .backbone import (
+    CapsuleBackbone,
+    ConvBNAct,
+    DeformableCaps2d,
+    DeformableCapsBlock,
+    PrimaryCaps2d,
+    RoutingCaps,
+    squash,
+)
+from .head import (
+    CapsuleDetect,
+    CapsuleDetectv1,
+    CapsuleDetectv2,
+    CapsuleDetectv4,
+    CapsuleDetectv5,
+    CapsuleDetectv6,
+    CapsuleDetectv7,
+    CapsuleDetectv8,
+    CapsuleOpenVocabDetect,
+    CapsuleDualHead,
+    CapsuleSegmentv1,
+    CapsuleSegmentv2,
+    CapsuleSegmentv3,
+)
+from .neck import CapsAlign, CapsDecode, CapsProj, CapsRoute, CapsRoutev2, CapsRoutev3, CapsRoutev4, CapsuleTap
+
+__all__ = [
+    "CapsuleBackbone",
+    "CapsuleTap",
+    "CapsRoute",
+    "CapsRoutev2",
+    "CapsRoutev3",
+    "CapsRoutev4",
+    "CapsProj",
+    "CapsDecode",
+    "CapsAlign",
+    "CapsuleDetect",
+    "CapsuleDetectv1",
+    "CapsuleDetectv2",
+    "CapsuleDetectv4",
+    "CapsuleDetectv5",
+    "CapsuleDetectv6",
+    "CapsuleDetectv7",
+    "CapsuleDetectv8",
+    "CapsuleOpenVocabDetect",
+    "CapsuleSegmentv1",
+    "CapsuleSegmentv2",
+    "CapsuleSegmentv3",
+    "CapsuleDualHead",
+    "ConvBNAct",
+    "DeformableCaps2d",
+    "DeformableCapsBlock",
+    "PrimaryCaps2d",
+    "RoutingCaps",
+    "squash",
+]

modules/backbone.py

@@ -0,0 +1,276 @@
+from __future__ import annotations
+
+from typing import Iterable, Sequence
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def squash(x: torch.Tensor, dim: int = -1, eps: float = 1e-7) -> torch.Tensor:
+    """Squash nonlinearity used by capsule networks."""
+    squared_norm = (x * x).sum(dim=dim, keepdim=True)
+    scale = squared_norm / (1.0 + squared_norm)
+    return scale * x / torch.sqrt(squared_norm + eps)
+
+
+class ConvBNAct(nn.Module):
+    """Convolution + BatchNorm + SiLU."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        padding: int | None = None,
+    ) -> None:
+        super().__init__()
+        if padding is None:
+            padding = kernel_size // 2
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False)
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.act = nn.SiLU(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.act(self.bn(self.conv(x)))
+
+
+class PrimaryCaps2d(nn.Module):
+    """Primary capsule layer for 2D feature maps."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        num_caps: int,
+        dim_caps: int,
+        kernel_size: int = 1,
+        stride: int = 1,
+        padding: int | None = None,
+    ) -> None:
+        super().__init__()
+        if padding is None:
+            padding = kernel_size // 2
+        out_channels = num_caps * dim_caps
+        self.num_caps = num_caps
+        self.dim_caps = dim_caps
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False)
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.act = nn.SiLU(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.act(self.bn(self.conv(x)))
+        bsz, _, h, w = x.shape
+        x = x.view(bsz, self.num_caps, self.dim_caps, h, w)
+        return squash(x, dim=2)
+
+
+class RoutingCaps(nn.Module):
+    """Dynamic routing between capsules."""
+
+    def __init__(
+        self,
+        num_in_caps: int,
+        dim_in: int,
+        num_out_caps: int,
+        dim_out: int,
+        routing_iters: int = 3,
+    ) -> None:
+        super().__init__()
+        self.num_in_caps = num_in_caps
+        self.dim_in = dim_in
+        self.num_out_caps = num_out_caps
+        self.dim_out = dim_out
+        self.routing_iters = routing_iters
+
+        weight = torch.randn(1, num_in_caps, num_out_caps, dim_out, dim_in) * 0.01
+        self.W = nn.Parameter(weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.ndim != 3:
+            raise ValueError(f"RoutingCaps expects [B, N, D], got {tuple(x.shape)}")
+        bsz = x.shape[0]
+        x = x.unsqueeze(2).unsqueeze(-1)  # [B, N, 1, D, 1]
+        u_hat = torch.matmul(self.W, x).squeeze(-1)  # [B, N, M, Dout]
+        b = x.new_zeros(bsz, self.num_in_caps, self.num_out_caps)
+        for idx in range(self.routing_iters):
+            c = F.softmax(b, dim=-1)
+            s = (c.unsqueeze(-1) * u_hat).sum(dim=1)
+            v = squash(s, dim=-1)
+            if idx < self.routing_iters - 1:
+                b = b + (u_hat * v.unsqueeze(1)).sum(dim=-1)
+        return v
+
+
+class DeformableCaps2d(nn.Module):
+    """Deformable capsule layer with learned sampling offsets."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        num_child_caps: int = 8,
+        dim_child: int = 8,
+        num_parent_caps: int = 8,
+        dim_parent: int = 8,
+        num_samples: int = 4,
+        routing_iters: int = 3,
+        offset_scale: float = 1.0,
+        out_channels: int | None = None,
+    ) -> None:
+        super().__init__()
+        self.num_child_caps = num_child_caps
+        self.dim_child = dim_child
+        self.num_parent_caps = num_parent_caps
+        self.dim_parent = dim_parent
+        self.num_samples = num_samples
+        self.routing_iters = routing_iters
+        self.offset_scale = offset_scale
+
+        self.primary = PrimaryCaps2d(in_channels, num_child_caps, dim_child, kernel_size=1, stride=1, padding=0)
+        self.offset = nn.Conv2d(in_channels, 2 * num_samples, kernel_size=3, stride=1, padding=1)
+        nn.init.zeros_(self.offset.weight)
+        nn.init.zeros_(self.offset.bias)
+
+        self.routing = RoutingCaps(
+            num_in_caps=num_samples * num_child_caps,
+            dim_in=dim_child,
+            num_out_caps=num_parent_caps,
+            dim_out=dim_parent,
+            routing_iters=routing_iters,
+        )
+
+        caps_channels = num_parent_caps * dim_parent
+        self.out_channels = out_channels or caps_channels
+        self.project = None
+        if self.out_channels != caps_channels:
+            self.project = ConvBNAct(caps_channels, self.out_channels, kernel_size=1, stride=1, padding=0)
+
+    @staticmethod
+    def _base_grid(h: int, w: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        ys = torch.linspace(-1.0, 1.0, h, device=device, dtype=dtype)
+        xs = torch.linspace(-1.0, 1.0, w, device=device, dtype=dtype)
+        yy, xx = torch.meshgrid(ys, xs, indexing="ij")
+        return torch.stack((xx, yy), dim=-1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        child = self.primary(x)  # [B, Nc, Dc, H, W]
+        bsz, _, _, h, w = child.shape
+
+        child_flat = child.view(bsz, self.num_child_caps * self.dim_child, h, w)
+        offsets = self.offset(x).view(bsz, self.num_samples, 2, h, w)
+        offsets = torch.tanh(offsets) * self.offset_scale
+
+        scale_x = max(w - 1, 1) / 2.0
+        scale_y = max(h - 1, 1) / 2.0
+        scale = offsets.new_tensor([scale_x, scale_y]).view(1, 1, 2, 1, 1)
+        offsets = offsets / scale
+
+        base = self._base_grid(h, w, x.device, x.dtype).view(1, 1, h, w, 2)
+        grids = base + offsets.permute(0, 1, 3, 4, 2)
+
+        sampled = []
+        for idx in range(self.num_samples):
+            grid = grids[:, idx]
+            feat = F.grid_sample(
+                child_flat,
+                grid,
+                mode="bilinear",
+                padding_mode="zeros",
+                align_corners=True,
+            )
+            sampled.append(feat)
+
+        sampled = torch.stack(sampled, dim=1)
+        sampled = sampled.view(bsz, self.num_samples, self.num_child_caps, self.dim_child, h, w)
+        sampled = sampled.permute(0, 4, 5, 1, 2, 3).contiguous()
+        sampled = sampled.view(bsz * h * w, self.num_samples * self.num_child_caps, self.dim_child)
+
+        routed = self.routing(sampled)
+        routed = routed.view(bsz, h, w, self.num_parent_caps, self.dim_parent)
+        routed = routed.permute(0, 3, 4, 1, 2).contiguous()
+        out = routed.view(bsz, self.num_parent_caps * self.dim_parent, h, w)
+
+        if self.project is not None:
+            out = self.project(out)
+        return out
+
+
+class DeformableCapsBlock(nn.Module):
+    """Backbone block: Conv downsample + deformable capsule routing."""
+
+    def __init__(
+        self,
+        c1: int,
+        c2: int,
+        num_child_caps: int = 8,
+        dim_child: int = 8,
+        num_parent_caps: int = 8,
+        dim_parent: int = 8,
+        num_samples: int = 4,
+        routing_iters: int = 3,
+        stride: int = 1,
+    ) -> None:
+        super().__init__()
+        self.down = ConvBNAct(c1, c2, kernel_size=3, stride=stride)
+        self.caps = DeformableCaps2d(
+            in_channels=c2,
+            num_child_caps=num_child_caps,
+            dim_child=dim_child,
+            num_parent_caps=num_parent_caps,
+            dim_parent=dim_parent,
+            num_samples=num_samples,
+            routing_iters=routing_iters,
+            out_channels=c2,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.down(x)
+        return self.caps(x)
+
+
+class CapsuleBackbone(nn.Module):
+    """Simple capsule-based backbone that returns multi-scale features."""
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        stem_channels: int = 64,
+        stages: Sequence[int] = (128, 256, 512),
+        capsule_cfgs: Iterable[dict] | None = None,
+    ) -> None:
+        super().__init__()
+        self.stem = ConvBNAct(in_channels, stem_channels, kernel_size=3, stride=2)
+        stage_cfgs = list(capsule_cfgs) if capsule_cfgs is not None else [{}] * len(stages)
+        if len(stage_cfgs) != len(stages):
+            raise ValueError("capsule_cfgs must match stages length")
+
+        blocks = []
+        in_ch = stem_channels
+        for out_ch, cfg in zip(stages, stage_cfgs):
+            blocks.append(
+                nn.Sequential(
+                    ConvBNAct(in_ch, out_ch, kernel_size=3, stride=2),
+                    DeformableCaps2d(out_ch, out_channels=out_ch, **cfg),
+                )
+            )
+            in_ch = out_ch
+        self.stages = nn.ModuleList(blocks)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, ...]:
+        x = self.stem(x)
+        outputs = []
+        for stage in self.stages:
+            x = stage(x)
+            outputs.append(x)
+        return tuple(outputs)
+
+
+__all__ = [
+    "CapsuleBackbone",
+    "ConvBNAct",
+    "DeformableCaps2d",
+    "DeformableCapsBlock",
+    "PrimaryCaps2d",
+    "RoutingCaps",
+    "squash",
+]

modules/head.py

@@ -0,0 +1,1801 @@
+from __future__ import annotations
+
+import copy
+import math
+import time
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ultralytics.nn.modules import Conv, DWConv, Detect, Segment
+from ultralytics.nn.modules.block import Proto26
+
+
+class PrimaryCaps(nn.Module):
+    r"""Primary convolutional capsules.
+
+    Outputs pose and activation, plus a concatenated NHWC capsule tensor.
+
+    Args:
+        A: Input feature channels.
+        B: Number of capsule types.
+        K: Convolution kernel size.
+        P: Pose matrix side length (pose size is ``P*P``).
+        stride: Convolution stride.
+
+    Input shape:
+        x: ``(N, A, H, W)``
+
+    Output shape:
+        a: ``(N, B, H_out, W_out)``
+        p: ``(N, B*P*P, H_out, W_out)``
+        out: ``(N, H_out, W_out, B*(P*P+1))``
+
+    Parameter size:
+        pose conv + act conv
+        ``(K*K*A*B*P*P + B*P*P) + (K*K*A*B + B)``
+    """
+
+    def __init__(self, A: int = 32, B: int = 32, K: int = 1, P: int = 4, stride: int = 1):
+        super().__init__()
+        self.B = B
+        self.P = P
+        self.psize = P * P
+
+        self.pose = nn.Conv2d(in_channels=A, out_channels=B * self.psize, kernel_size=K, stride=stride, bias=True)
+        self.a = nn.Conv2d(in_channels=A, out_channels=B, kernel_size=K, stride=stride, bias=True)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # p: (B, B*psize, H, W), a: (B, B, H, W)
+        p = self.pose(x)
+        a = self.sigmoid(self.a(x))
+        out = torch.cat([p, a], dim=1).permute(0, 2, 3, 1).contiguous()  # (B, H, W, B*(psize+1))
+        return a, p, out
+
+
+class ConvCaps(nn.Module):
+    r"""Convolutional capsules with EM routing.
+
+    Args:
+        B: Input capsule types.
+        C: Output capsule types.
+        K: Patch kernel size.
+        P: Pose matrix side length (pose size is ``P*P``).
+        stride: Spatial stride for patch extraction.
+        iters: Number of EM routing iterations.
+        coor_add: Add coordinate offsets (class-caps style option).
+        w_shared: Share transform matrices across spatial positions.
+
+    Input shape:
+        x: ``(N, H, W, B*(P*P+1))``
+
+    Output shape:
+        p_out: ``(N, H_out, W_out, C*P*P)``
+        a_out: ``(N, H_out, W_out, C)``
+        out: ``(N, H_out, W_out, C*(P*P+1))``
+
+    Parameter size:
+        If ``w_shared=False``:
+        ``weights: (K*K*B*C*P*P*P*P)``, ``beta_u: C``, ``beta_a: C``
+
+        If ``w_shared=True``:
+        ``weights: (B*C*P*P*P*P)``, ``beta_u: C``, ``beta_a: C``
+
+        Total = ``weights + 2*C`` (excluding non-trainable buffers).
+    """
+
+    def __init__(
+        self,
+        B: int = 32,
+        C: int = 32,
+        K: int = 3,
+        P: int = 4,
+        stride: int = 1,
+        iters: int = 3,
+        coor_add: bool = False,
+        w_shared: bool = False,
+    ):
+        super().__init__()
+        self.B = B
+        self.C = C
+        self.K = K
+        self.P = P
+        self.psize = P * P
+        self.stride = stride
+        self.iters = iters
+        self.coor_add = coor_add
+        self.w_shared = w_shared
+
+        self.eps = 1e-6
+        self._lambda = 1e-3
+        self.register_buffer("ln_2pi", torch.tensor(math.log(2 * math.pi), dtype=torch.float32), persistent=False)
+
+        # Matrix-caps paper uses per-capsule beta scalars.
+        self.beta_u = nn.Parameter(torch.zeros(C))
+        self.beta_a = nn.Parameter(torch.zeros(C))
+
+        # For non-shared conv-caps, input vote count is K*K*B. For shared mode it is B then repeated by HW.
+        weight_in = B if w_shared else (K * K * B)
+        self.weights = nn.Parameter(torch.randn(1, weight_in, C, self.psize, self.psize) * 0.02)
+
+        self.sigmoid = nn.Sigmoid()
+        self.softmax = nn.Softmax(dim=2)
+
+    def m_step(
+        self,
+        a_in: torch.Tensor,
+        r: torch.Tensor,
+        v: torch.Tensor,
+        eps: float,
+        b: int,
+        B: int,
+        C: int,
+        psize: int,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # a_in: (b, B, 1) or (b, B, 1, 1), r: (b, B, C, 1), v: (b, B, C, psize)
+        if a_in.ndim == 3:
+            a_in = a_in.unsqueeze(2)
+        r = r * a_in
+        r = r / (r.sum(dim=2, keepdim=True) + eps)
+        r_sum = r.sum(dim=1, keepdim=True)
+        coeff = r / (r_sum + eps)
+
+        mu = torch.sum(coeff * v, dim=1, keepdim=True)  # (b, 1, C, psize)
+        sigma_sq = torch.sum(coeff * (v - mu).pow(2), dim=1, keepdim=True) + eps
+        sigma_sq = sigma_sq.clamp_min(1e-4)
+
+        r_sum_flat = r_sum.view(b, C, 1)
+        sigma_sq_flat = sigma_sq.view(b, C, psize).clamp_min(1e-4)
+        cost_h = (self.beta_u.view(1, C, 1) + torch.log(torch.sqrt(sigma_sq_flat))) * r_sum_flat
+        a_out = self.sigmoid(self._lambda * (self.beta_a.view(1, C) - cost_h.sum(dim=2))).clamp(1e-4, 1.0 - 1e-4)
+
+        mu = torch.nan_to_num(mu, nan=0.0, posinf=1e4, neginf=-1e4)
+        sigma_sq = torch.nan_to_num(sigma_sq, nan=1e-4, posinf=1e4, neginf=1e-4)
+        a_out = torch.nan_to_num(a_out, nan=0.5, posinf=1.0 - 1e-4, neginf=1e-4)
+        return a_out, mu, sigma_sq
+
+    def e_step(
+        self,
+        mu: torch.Tensor,
+        sigma_sq: torch.Tensor,
+        a_out: torch.Tensor,
+        v: torch.Tensor,
+        eps: float,
+        b: int,
+        C: int,
+    ) -> torch.Tensor:
+        # mu: (b,1,C,psize), sigma_sq: (b,1,C,psize), a_out: (b,C), v: (b,B,C,psize)
+        sigma_sq = sigma_sq.clamp_min(1e-4)
+        a_out = a_out.clamp(1e-4, 1.0 - 1e-4)
+        ln_p_j_h = -1.0 * (v - mu).pow(2) / (2.0 * sigma_sq) - torch.log(torch.sqrt(sigma_sq)) - 0.5 * self.ln_2pi
+        ln_ap = ln_p_j_h.sum(dim=3) + torch.log(a_out.view(b, 1, C) + eps)
+        ln_ap = torch.nan_to_num(ln_ap, nan=0.0, posinf=50.0, neginf=-50.0)
+        r = self.softmax(ln_ap).unsqueeze(-1)  # (b,B,C,1)
+        r = torch.nan_to_num(r, nan=(1.0 / max(C, 1)), posinf=1.0, neginf=0.0)
+        return r
+
+    def caps_em_routing(self, v: torch.Tensor, a_in: torch.Tensor, C: int, eps: float) -> tuple[torch.Tensor, torch.Tensor]:
+        b, B, _, psize = v.shape
+        r = v.new_full((b, B, C, 1), 1.0 / C)
+
+        for t in range(self.iters):
+            a_out, mu, sigma_sq = self.m_step(a_in, r, v, eps, b, B, C, psize)
+            if t < self.iters - 1:
+                r = self.e_step(mu, sigma_sq, a_out, v, eps, b, C)
+
+        # p_out: (b, C, psize), a_out: (b, C)
+        p_out = torch.nan_to_num(mu.squeeze(1), nan=0.0, posinf=1e4, neginf=-1e4)
+        a_out = torch.nan_to_num(a_out, nan=0.5, posinf=1.0 - 1e-4, neginf=1e-4)
+        return p_out, a_out
+
+    def add_pathes(self, x: torch.Tensor, B: int, K: int, psize: int, stride: int) -> tuple[torch.Tensor, int, int]:
+        # x: (b, h, w, B*(psize+1)) -> patches: (b, oh, ow, K*K, B*(psize+1))
+        b, h, w, c = x.shape
+        x_chw = x.permute(0, 3, 1, 2).contiguous()
+        pad = K // 2
+        patches = F.unfold(x_chw, kernel_size=K, padding=pad, stride=stride)
+
+        oh = (h + 2 * pad - K) // stride + 1
+        ow = (w + 2 * pad - K) // stride + 1
+        patches = patches.transpose(1, 2).contiguous().view(b, oh, ow, K * K, c)
+        return patches, oh, ow
+
+    def transform_view(self, x: torch.Tensor, w: torch.Tensor, C: int, P: int, w_shared: bool = False) -> torch.Tensor:
+        # x: (b, in_votes, psize), w: (1, in_votes_base, C, psize, psize)
+        b, in_votes, psize = x.shape
+        assert psize == P * P
+
+        w0 = w[0]
+        if w_shared:
+            base = w0.size(0)
+            reps = in_votes // base
+            w0 = w0.repeat(reps, 1, 1, 1)
+
+        # (b, in_votes, C, psize)
+        v = torch.einsum("bip,icpq->bicq", x, w0)
+        return v
+
+    def add_coord(self, v: torch.Tensor, b: int, h: int, w: int, B: int, C: int, psize: int) -> torch.Tensor:
+        # v: (b, h*w*B, C, psize)
+        # Supports rectangular feature maps (h != w).
+        v = v.view(b, h, w, B, C, psize)
+
+        device = v.device
+        dtype = v.dtype
+        coor_h_vals = torch.arange(h, dtype=dtype, device=device) / float(max(h, 1))
+        coor_w_vals = torch.arange(w, dtype=dtype, device=device) / float(max(w, 1))
+
+        coor_h = torch.zeros(1, h, 1, 1, 1, psize, dtype=dtype, device=device)
+        coor_w = torch.zeros(1, 1, w, 1, 1, psize, dtype=dtype, device=device)
+        coor_h[0, :, 0, 0, 0, 0] = coor_h_vals
+        coor_w[0, 0, :, 0, 0, 1] = coor_w_vals
+
+        v = (v + coor_h + coor_w).view(b, h * w * B, C, psize)
+        return v
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # x shape: (b, h, w, B*(psize+1))
+        b, h, w, c = x.shape
+
+        if not self.w_shared:
+            patches, oh, ow = self.add_pathes(x, self.B, self.K, self.psize, self.stride)
+
+            p_in = patches[..., : self.B * self.psize].contiguous().view(b * oh * ow, self.K * self.K * self.B, self.psize)
+            a_in = patches[..., self.B * self.psize :].contiguous().view(b * oh * ow, self.K * self.K * self.B, 1)
+
+            v = self.transform_view(p_in, self.weights, self.C, self.P, w_shared=False)
+            p_out, a_out = self.caps_em_routing(v, a_in, self.C, self.eps)
+
+            p_out = p_out.view(b, oh, ow, self.C * self.psize)
+            a_out = a_out.view(b, oh, ow, self.C)
+            out = torch.cat([p_out, a_out], dim=3)
+        else:
+            assert c == self.B * (self.psize + 1)
+            assert self.K == 1 
+            assert self.stride == 1
+
+            p_in = x[..., : self.B * self.psize].contiguous().view(b, h * w * self.B, self.psize)
+            a_in = x[..., self.B * self.psize :].contiguous().view(b, h * w * self.B, 1)
+
+            v = self.transform_view(p_in, self.weights, self.C, self.P, w_shared=True)
+            if self.coor_add:
+                v = self.add_coord(v, b, h, w, self.B, self.C, self.psize)
+
+            p_cls, a_cls = self.caps_em_routing(v, a_in, self.C, self.eps)
+
+            # Broadcast class capsules back to spatial map for Detect-style dense outputs.
+            p_out = p_cls.reshape(b, 1, 1, self.C * self.psize).expand(b, h, w, self.C * self.psize)
+            a_out = a_cls.unsqueeze(1).unsqueeze(1).expand(b, h, w, self.C)
+            out = torch.cat([p_out, a_out], dim=3)
+
+        return p_out, a_out, out
+
+
+
+class DynamicConvCaps(nn.Module):
+    r"""Convolutional capsules with Sabour-style dynamic routing.
+
+    This layer keeps the same tensor interface as ``ConvCaps``:
+        input:  (N, H, W, B*(P*P+1))
+        output: p_out (N, H_out, W_out, C*P*P), a_out (N, H_out, W_out, C), out concat
+
+    Args:
+        B: Input capsule types.
+        C: Output capsule types.
+        K: Patch kernel size.
+        P: Pose matrix side length.
+        stride: Patch stride.
+        iters: Routing iterations.
+        coor_add: Add coordinates in shared mode.
+        w_shared: Share transforms across spatial positions (requires K=1, stride=1).
+    """
+
+    def __init__(
+        self,
+        B: int = 32,
+        C: int = 32,
+        K: int = 3,
+        P: int = 4,
+        stride: int = 1,
+        iters: int = 3,
+        coor_add: bool = False,
+        w_shared: bool = False,
+    ):
+        super().__init__()
+        self.B = B
+        self.C = C
+        self.K = K
+        self.P = P
+        self.psize = P * P
+        self.stride = stride
+        self.iters = iters
+        self.coor_add = coor_add
+        self.w_shared = w_shared
+        self.eps = 1e-6
+
+        weight_in = B if w_shared else (K * K * B)
+        self.weights = nn.Parameter(torch.randn(1, weight_in, C, self.psize, self.psize) * 0.02)
+
+    @staticmethod
+    def squash(s: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor:
+        s2 = (s * s).sum(dim=dim, keepdim=True)
+        scale = s2 / (1.0 + s2)
+        return scale * s / torch.sqrt(s2 + eps)
+
+    def add_pathes(self, x: torch.Tensor, K: int, stride: int) -> tuple[torch.Tensor, int, int]:
+        b, h, w, c = x.shape
+        x_chw = x.permute(0, 3, 1, 2).contiguous()
+        pad = K // 2
+        patches = F.unfold(x_chw, kernel_size=K, padding=pad, stride=stride)
+        oh = (h + 2 * pad - K) // stride + 1
+        ow = (w + 2 * pad - K) // stride + 1
+        patches = patches.transpose(1, 2).contiguous().view(b, oh, ow, K * K, c)
+        return patches, oh, ow
+
+    def transform_view(self, x: torch.Tensor, w_shared: bool) -> torch.Tensor:
+        # x: (b, in_votes, psize) -> votes: (b, in_votes, C, psize)
+        b, in_votes, psize = x.shape
+        if psize != self.psize:
+            raise ValueError('Invalid pose size for DynamicConvCaps')
+
+        w0 = self.weights[0]
+        if w_shared:
+            base = w0.size(0)
+            reps = in_votes // base
+            w0 = w0.repeat(reps, 1, 1, 1)
+
+        return torch.einsum('bip,icpq->bicq', x, w0)
+
+    def add_coord(self, v: torch.Tensor, b: int, h: int, w: int, B: int, C: int, psize: int) -> torch.Tensor:
+        # v: (b, h*w*B, C, psize)
+        v = v.view(b, h, w, B, C, psize)
+        device, dtype = v.device, v.dtype
+        coor_h_vals = torch.arange(h, dtype=dtype, device=device) / float(max(h, 1))
+        coor_w_vals = torch.arange(w, dtype=dtype, device=device) / float(max(w, 1))
+
+        coor_h = torch.zeros(1, h, 1, 1, 1, psize, dtype=dtype, device=device)
+        coor_w = torch.zeros(1, 1, w, 1, 1, psize, dtype=dtype, device=device)
+        coor_h[0, :, 0, 0, 0, 0] = coor_h_vals
+        coor_w[0, 0, :, 0, 0, 1] = coor_w_vals
+
+        return (v + coor_h + coor_w).view(b, h * w * B, C, psize)
+
+    def dynamic_routing(self, v: torch.Tensor, a_in: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # v: (n, in_votes, C, psize), a_in: (n, in_votes, 1)
+        n, in_votes, C, psize = v.shape
+        b_ij = v.new_zeros(n, in_votes, C)
+
+        a_in = a_in.clamp(1e-4, 1.0)
+        for t in range(self.iters):
+            c_ij = F.softmax(b_ij, dim=2)
+            c_ij = c_ij * a_in
+            c_ij = c_ij / (c_ij.sum(dim=2, keepdim=True) + self.eps)
+
+            s_j = (c_ij.unsqueeze(-1) * v).sum(dim=1)
+            v_j = self.squash(s_j, dim=-1, eps=self.eps)
+            if t < self.iters - 1:
+                agreement = (v * v_j.unsqueeze(1)).sum(dim=-1)
+                b_ij = b_ij + agreement
+
+        # activation from vector length in (0,1)
+        a_out = torch.sqrt((v_j * v_j).sum(dim=-1) + self.eps).clamp(1e-4, 1.0 - 1e-4)
+        return v_j, a_out
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        b, h, w, c = x.shape
+
+        if not self.w_shared:
+            patches, oh, ow = self.add_pathes(x, self.K, self.stride)
+            p_in = patches[..., : self.B * self.psize].contiguous().view(b * oh * ow, self.K * self.K * self.B, self.psize)
+            a_in = patches[..., self.B * self.psize :].contiguous().view(b * oh * ow, self.K * self.K * self.B, 1)
+
+            votes = self.transform_view(p_in, w_shared=False)
+            p_vec, a_vec = self.dynamic_routing(votes, a_in)
+
+            p_out = p_vec.view(b, oh, ow, self.C * self.psize)
+            a_out = a_vec.view(b, oh, ow, self.C)
+            out = torch.cat([p_out, a_out], dim=3)
+        else:
+            if c != self.B * (self.psize + 1) or self.K != 1 or self.stride != 1:
+                raise ValueError('DynamicConvCaps shared mode requires K=1, stride=1 and matching capsule channels')
+
+            p_in = x[..., : self.B * self.psize].contiguous().view(b, h * w * self.B, self.psize)
+            a_in = x[..., self.B * self.psize :].contiguous().view(b, h * w * self.B, 1)
+            votes = self.transform_view(p_in, w_shared=True)
+            if self.coor_add:
+                votes = self.add_coord(votes, b, h, w, self.B, self.C, self.psize)
+
+            p_vec, a_vec = self.dynamic_routing(votes, a_in)
+            p_out = p_vec.reshape(b, 1, 1, self.C * self.psize).expand(b, h, w, self.C * self.psize)
+            a_out = a_vec.unsqueeze(1).unsqueeze(1).expand(b, h, w, self.C)
+            out = torch.cat([p_out, a_out], dim=3)
+
+        p_out = torch.nan_to_num(p_out, nan=0.0, posinf=1e4, neginf=-1e4)
+        a_out = torch.nan_to_num(a_out, nan=0.5, posinf=1.0 - 1e-4, neginf=1e-4)
+        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)
+        return p_out, a_out, out
+
+
+class SelfRoutingConvCaps(nn.Module):
+    r"""Convolutional self-routing capsules.
+
+    Keeps the same output contract as ``ConvCaps``/``DynamicConvCaps``:
+        input:  (N, H, W, B*(P*P+1))
+        output: p_out (N, H_out, W_out, C*P*P), a_out (N, H_out, W_out, C), out concat
+    """
+
+    def __init__(
+        self,
+        B: int = 32,
+        C: int = 32,
+        K: int = 3,
+        P: int = 4,
+        stride: int = 1,
+        iters: int = 1,
+        coor_add: bool = False,
+        w_shared: bool = False,
+    ):
+        super().__init__()
+        _ = (iters, w_shared)  # kept for API compatibility with other capsule layers.
+
+        self.B = B
+        self.C = C
+        self.K = K
+        self.P = P
+        self.psize = P * P
+        self.stride = stride
+        self.coor_add = coor_add
+        self.eps = 1e-6
+
+        self.kk = K * K
+        self.kkB = self.kk * B
+
+        # Pose transform for each input capsule vote -> output capsule pose.
+        self.W1 = nn.Parameter(torch.empty(self.kkB, C, self.psize, self.psize))
+        nn.init.kaiming_uniform_(self.W1, a=math.sqrt(5))
+
+        # Routing logits from local pose vectors.
+        self.W2 = nn.Parameter(torch.zeros(self.kkB, C, self.psize))
+        self.b2 = nn.Parameter(torch.zeros(1, 1, self.kkB, C))
+
+    def _output_hw(self, h: int, w: int) -> tuple[int, int]:
+        pad = self.K // 2
+        oh = (h + 2 * pad - self.K) // self.stride + 1
+        ow = (w + 2 * pad - self.K) // self.stride + 1
+        return oh, ow
+
+    def _add_coord(self, pose_unf: torch.Tensor, oh: int, ow: int) -> torch.Tensor:
+        # pose_unf: (b, L, kkB, psize)
+        if self.psize < 2:
+            return pose_unf
+
+        b, L, kkB, _ = pose_unf.shape
+        device, dtype = pose_unf.device, pose_unf.dtype
+        gy = torch.arange(oh, device=device, dtype=dtype) / float(max(oh, 1))
+        gx = torch.arange(ow, device=device, dtype=dtype) / float(max(ow, 1))
+        yy, xx = torch.meshgrid(gy, gx, indexing='ij')
+        coords = torch.stack((yy, xx), dim=-1).view(1, L, 1, 2)
+
+        pose_unf = pose_unf.clone()
+        pose_unf[..., :2] = pose_unf[..., :2] + coords
+        return pose_unf
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # x: (b, h, w, B*(psize+1))
+        b, h, w, c = x.shape
+        expected = self.B * (self.psize + 1)
+        if c != expected:
+            raise ValueError(f'SelfRoutingConvCaps expected {expected} channels, got {c}')
+
+        pose = x[..., : self.B * self.psize]
+        act = x[..., self.B * self.psize :]
+
+        pose_chw = pose.permute(0, 3, 1, 2).contiguous()
+        act_chw = act.permute(0, 3, 1, 2).contiguous()
+
+        pad = self.K // 2
+        pose_unf = F.unfold(pose_chw, kernel_size=self.K, stride=self.stride, padding=pad)
+        act_unf = F.unfold(act_chw, kernel_size=self.K, stride=self.stride, padding=pad)
+
+        oh, ow = self._output_hw(h, w)
+        l = pose_unf.shape[-1]
+
+        pose_unf = pose_unf.view(b, self.B, self.psize, self.kk, l).permute(0, 4, 3, 1, 2).contiguous()
+        pose_unf = pose_unf.view(b, l, self.kkB, self.psize)
+
+        act_unf = act_unf.view(b, self.B, self.kk, l).permute(0, 3, 2, 1).contiguous()
+        act_unf = act_unf.view(b, l, self.kkB)
+
+        if self.coor_add:
+            pose_unf = self._add_coord(pose_unf, oh, ow)
+
+        # Routing logits and couplings.
+        logit = torch.einsum('blip,icp->blic', pose_unf, self.W2) + self.b2
+        r = F.softmax(logit, dim=3)
+
+        ar = act_unf.unsqueeze(-1) * r
+        ar_sum = ar.sum(dim=2, keepdim=True) + self.eps
+        coeff = ar / ar_sum
+
+        a_norm = act_unf.sum(dim=2, keepdim=True) + self.eps
+        a_out = (ar_sum.squeeze(2) / a_norm).clamp(1e-4, 1.0 - 1e-4)
+
+        pose_votes = torch.einsum('blip,icpq->blicq', pose_unf, self.W1)
+        pose_out = (coeff.unsqueeze(-1) * pose_votes).sum(dim=2)
+
+        p_out = pose_out.view(b, oh, ow, self.C * self.psize)
+        a_out = a_out.view(b, oh, ow, self.C)
+        out = torch.cat([p_out, a_out], dim=3)
+
+        p_out = torch.nan_to_num(p_out, nan=0.0, posinf=1e4, neginf=-1e4)
+        a_out = torch.nan_to_num(a_out, nan=0.5, posinf=1.0 - 1e-4, neginf=1e-4)
+        out = torch.nan_to_num(out, nan=0.0, posinf=1e4, neginf=-1e4)
+        return p_out, a_out, out
+
+
+
+class CapsuleDualHead(nn.Module):
+    """Capsule detection head for one feature level.
+
+    Args:
+        c_in: Input channels of this feature scale (from parser-provided ``ch``).
+        nc: Number of classes (final activation capsule count in ``ConvCaps2``).
+        reg_max: Detect DFL bins, box channels are ``4 * reg_max``.
+        k: Number of capsule types in ``PrimaryCaps``.
+        d: Requested pose descriptor size; internally mapped to square ``P*P``.
+
+    Input shape:
+        x: ``(N, c_in, H, W)``
+
+    Output shape:
+        boxes: ``(N, 4*reg_max, H, W)``
+        scores: ``(N, nc, H, W)``
+        aux: dict with final capsule activations when ``return_aux=True`` else ``None``
+
+    Parameter size:
+        ``PrimaryCaps(c_in,k) + ConvCaps(k,nc,w_shared=True) + box_bias(4*reg_max)``
+
+    Structure:
+        PrimaryCaps -> ConvCaps(class caps only, shared)
+    """
+
+    def __init__(self, c_in: int, nc: int, reg_max: int, k: int, d: int):
+        super().__init__()
+        # Matrix-caps pose is square; choose smallest square >= requested d.
+        p = max(1, int(math.ceil(math.sqrt(d))))
+
+        self.nc = nc
+        self.reg_max = reg_max
+        self.P = p
+        self.psize = self.P * self.P
+
+        # A=c_in, B=k, P controls pose channels as B*(P*P).
+        self.primary = PrimaryCaps(A=c_in, B=k, K=1, P=self.P, stride=1)
+        # Single class-caps layer with shared transforms for parameter reduction.
+        self.conv_caps2 = ConvCaps(B=k, C=nc, K=1, P=self.P, stride=1, iters=1, coor_add=True, w_shared=True)
+
+        # Detect-style localization prior set in CapsuleDetect.bias_init().
+        self.box_bias = nn.Parameter(torch.zeros(4 * reg_max))
+
+    def _pose_to_box(self, p_out: torch.Tensor, a_out: torch.Tensor) -> torch.Tensor:
+        # p_out: (b,h,w,nc*psize), a_out is intentionally unused here.
+        # Simple rule requested: use first 4*reg_max pose values as box channels.
+        _ = a_out
+        box_ch = 4 * self.reg_max
+
+        if p_out.shape[-1] >= box_ch:
+            box = p_out[..., :box_ch]
+        else:
+            # If pose channels are fewer than required box channels, repeat and trim.
+            reps = math.ceil(box_ch / p_out.shape[-1])
+            box = p_out.repeat(1, 1, 1, reps)[..., :box_ch]
+
+        return box + self.box_bias.view(1, 1, 1, box_ch)
+
+    def forward(self, x: torch.Tensor, return_aux: bool = False) -> tuple[torch.Tensor, torch.Tensor, dict | None]:
+        _, _, caps0 = self.primary(x)
+        p2, a2, _ = self.conv_caps2(caps0)
+
+        boxes = self._pose_to_box(p2, a2).permute(0, 3, 1, 2).contiguous()  # (b,4*reg_max,h,w)
+        a2_logits = torch.logit(a2.clamp(1e-4, 1.0 - 1e-4))
+        scores = a2_logits.permute(0, 3, 1, 2).contiguous()  # (b,nc,h,w) logits
+
+        aux = None
+        if return_aux:
+            aux = {
+                "caps2_a": a2.permute(0, 3, 1, 2).contiguous(),
+            }
+        return boxes, scores, aux
+
+
+class CapsuleClsHead(nn.Module):
+    """Capsule classification branch used as a drop-in replacement for Detect.cv3."""
+
+    def __init__(self, c_in: int, nc: int, k: int = 4, d: int = 16, iters: int = 1):
+        super().__init__()
+        p = max(1, int(math.ceil(math.sqrt(d))))
+        self.primary = PrimaryCaps(A=c_in, B=k, K=1, P=p, stride=1)
+        # Internal capsule refinement layer.
+        self.mid_caps = SelfRoutingConvCaps(B=k, C=int((k+nc)/2), K=1, P=p, stride=1, iters=iters, coor_add=False, w_shared=True)
+        self.class_caps = SelfRoutingConvCaps(B=int((k+nc)/2), C=nc, K=1, P=p, stride=1, iters=iters, coor_add=False, w_shared=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Output Detect-compatible class logits in BCHW.
+        _, _, caps = self.primary(x)
+        _, _, caps_mid = self.mid_caps(caps)
+        _, a_out, _ = self.class_caps(caps_mid)
+        logits = torch.logit(a_out.clamp(1e-4, 1.0 - 1e-4)).permute(0, 3, 1, 2).contiguous()
+        return torch.nan_to_num(logits, nan=0.0, posinf=20.0, neginf=-20.0).float()
+
+
+class CapsuleDetect(Detect):
+    """Detect head with capsule vote aggregation for both box and cls branches.
+
+    Input feature of level i is packed as interleaved channels:
+        [pose(d_i), act(1)] repeated k_i times -> C_i = k_i * (d_i + 1)
+
+    In forward_head:
+    - split pose/act per capsule type
+    - run Detect box/cls heads on each type-specific pose tensor
+    - aggregate type predictions with act-driven vote weights
+
+    Detect decode/postprocess/end2end flow is reused unchanged.
+    """
+
+    def __init__(
+        self,
+        nc: int = 80,
+        *args,
+        reg_max: int = 16,
+        end2end: bool = False,
+        k: list[int] | tuple[int, ...] = (4, 8, 16),
+        d: list[int] | tuple[int, ...] = (16, 16, 16),
+        ch: tuple = (),
+    ):
+        parsed = list(args)
+        if parsed and isinstance(parsed[-1], (list, tuple)):
+            ch = tuple(parsed.pop(-1))
+
+        # Parser layout: [k_list, d_list, reg_max, end2end, ch]
+        if len(parsed) not in (2, 4):
+            raise ValueError('CapsuleDetect expects [k_list, d_list, reg_max, end2end, ch].')
+
+        k, d = parsed[0], parsed[1]
+        if len(parsed) == 4:
+            reg_max = int(parsed[2])
+            end2end = bool(parsed[3])
+
+        if not isinstance(k, (list, tuple)) or not isinstance(d, (list, tuple)):
+            raise TypeError('CapsuleDetect requires list/tuple k and d (per-level settings).')
+
+        ch = tuple(int(c) for c in ch)
+        nl = len(ch)
+        if len(k) != nl or len(d) != nl:
+            raise ValueError(f'CapsuleDetect k/d length must equal number of levels ({nl}).')
+
+        self.k_list = tuple(int(v) for v in k)
+        self.d_list = tuple(int(v) for v in d)
+
+        for i, c in enumerate(ch):
+            expected = self.k_list[i] * (self.d_list[i] + 1)
+            if c != expected:
+                raise ValueError(
+                    f'CapsuleDetect level-{i} channel mismatch: got {c}, expected {expected} from k={self.k_list[i]}, d={self.d_list[i]}.'
+                )
+
+        # Detect heads operate on per-type pose tensors (d_i channels).
+        super().__init__(nc=nc, reg_max=reg_max, end2end=end2end, ch=self.d_list)
+
+        # Vote weights from activation channels (K_i channels), separate for cls/box.
+        self.box_vote = nn.ModuleList(
+            nn.Sequential(Conv(k_i, k_i, 3), nn.Conv2d(k_i, k_i, 1, bias=True)) for k_i in self.k_list
+        )
+        self.cls_vote = nn.ModuleList(
+            nn.Sequential(Conv(k_i, k_i, 3), nn.Conv2d(k_i, k_i, 1, bias=True)) for k_i in self.k_list
+        )
+
+    def _split_caps(self, x: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+        """Split packed feature into pose and activation tensors per level.
+
+        Returns:
+            pose_caps: list of tensors, each (B, K, D, H, W)
+            act_map:   list of tensors, each (B, K, H, W)
+        """
+        pose_caps, act_map = [], []
+        for i, xi in enumerate(x):
+            k_i = self.k_list[i]
+            d_i = self.d_list[i]
+            c = int(xi.shape[1])
+            expected = k_i * (d_i + 1)
+            if c != expected:
+                raise ValueError(f'CapsuleDetect level-{i} channel mismatch: got {c}, expected {expected}.')
+
+            b, _, h, w = xi.shape
+            caps = xi.view(b, k_i, d_i + 1, h, w)
+            pose_caps.append(caps[:, :, :d_i].contiguous())
+            act_map.append(caps[:, :, d_i].contiguous())
+        return pose_caps, act_map
+
+    @staticmethod
+    def _normalized_votes(raw: torch.Tensor, eps: float = 1e-4) -> torch.Tensor:
+        # No softmax/sigmoid: use softplus + sum-normalization.
+        w = F.softplus(raw) + eps
+        return w / (w.sum(dim=1, keepdim=True) + eps)
+
+    def _run_voted_head(
+        self,
+        pose: torch.Tensor,
+        act: torch.Tensor,
+        head: torch.nn.Module,
+        vote_head: torch.nn.Module,
+        out_ch: int,
+    ) -> torch.Tensor:
+        """Apply one Detect head per type and aggregate by vote weights.
+
+        Args:
+            pose: (B, K, D, H, W)
+            act: (B, K, H, W)
+            head: Detect box or cls head module for this level
+            vote_head: vote logits module for this level
+            out_ch: output channels of target prediction
+
+        Returns:
+            (B, out_ch, H, W)
+        """
+        b, k, d, h, w = pose.shape
+
+        # No voting needed when there is only one capsule type.
+        if k == 1:
+            return head(pose[:, 0])
+
+        pose_bt = pose.reshape(b * k, d, h, w)
+        pred_bt = head(pose_bt).reshape(b, k, out_ch, h, w)
+
+        vote_raw = vote_head(act)  # (B, K, H, W)
+        vote = self._normalized_votes(vote_raw).unsqueeze(2)  # (B, K, 1, H, W)
+        pred = (pred_bt * vote).sum(dim=1)
+        return pred
+
+    def forward_head(
+        self, x: list[torch.Tensor], box_head: torch.nn.Module = None, cls_head: torch.nn.Module = None
+    ) -> dict[str, torch.Tensor]:
+        if box_head is None or cls_head is None:
+            return dict()
+
+        pose_caps, act_map = self._split_caps(x)
+        bs = x[0].shape[0]
+
+        box_list = []
+        cls_list = []
+        for i in range(self.nl):
+            box_i = self._run_voted_head(
+                pose_caps[i],
+                act_map[i],
+                box_head[i],
+                self.box_vote[i],
+                out_ch=4 * self.reg_max,
+            )
+            cls_i = self._run_voted_head(
+                pose_caps[i],
+                act_map[i],
+                cls_head[i],
+                self.cls_vote[i],
+                out_ch=self.nc,
+            )
+            box_list.append(box_i.view(bs, 4 * self.reg_max, -1))
+            cls_list.append(cls_i.view(bs, self.nc, -1))
+
+        boxes = torch.cat(box_list, dim=-1)
+        scores = torch.cat(cls_list, dim=-1)
+        return dict(boxes=boxes, scores=scores, feats=x)
+    
+
+class CapsuleDetectv1(Detect):
+    """Capsule Detect variant with activation-gated pose fusion.
+
+    Per level:
+    1) Split packed capsule channels into pose/activation (interleaved by type).
+    2) Use a 2-layer 1x1 gate net on activation channels.
+    3) Gate pose channels with residual scaling.
+    4) Flatten to K*D channels and run original Detect cv2/cv3 heads.
+    """
+
+    def __init__(
+        self,
+        nc: int = 80,
+        *args,
+        reg_max: int = 16,
+        end2end: bool = False,
+        k: list[int] | tuple[int, ...] = (4, 8, 16),
+        d: list[int] | tuple[int, ...] = (16, 16, 16),
+        ch: tuple = (),
+    ):
+        parsed = list(args)
+        if parsed and isinstance(parsed[-1], (list, tuple)):
+            ch = tuple(parsed.pop(-1))
+
+        # Parser layout: [k_list, d_list, reg_max, end2end, ch]
+        if len(parsed) not in (2, 4):
+            raise ValueError("CapsuleDetectv1 expects [k_list, d_list, reg_max, end2end, ch].")
+
+        k, d = parsed[0], parsed[1]
+        if len(parsed) == 4:
+            reg_max = int(parsed[2])
+            end2end = bool(parsed[3])
+
+        if not isinstance(k, (list, tuple)) or not isinstance(d, (list, tuple)):
+            raise TypeError("CapsuleDetectv1 requires list/tuple k and d (per-level settings).")
+
+        ch = tuple(int(c) for c in ch)
+        nl = len(ch)
+        if len(k) != nl or len(d) != nl:
+            raise ValueError(f"CapsuleDetectv1 k/d length must equal number of levels ({nl}).")
+
+        self.k_list = tuple(int(v) for v in k)
+        self.d_list = tuple(int(v) for v in d)
+
+        # Input from neck is packed as K*(D+1): [pose(D), act(1)] repeated K types.
+        for i, c in enumerate(ch):
+            expected = self.k_list[i] * (self.d_list[i] + 1)
+            if c != expected:
+                raise ValueError(
+                    f"CapsuleDetectv1 level-{i} channel mismatch: got {c}, "
+                    f"expected {expected} from k={self.k_list[i]}, d={self.d_list[i]}."
+                )
+
+        # Detect heads consume merged pose channels: K*D.
+        merged_ch = tuple(k_i * d_i for k_i, d_i in zip(self.k_list, self.d_list))
+        super().__init__(nc=nc, reg_max=reg_max, end2end=end2end, ch=merged_ch)
+
+        self.pose_gates = nn.ModuleList()
+        self.gate_alpha = nn.ParameterList()
+        for k_i, d_i in zip(self.k_list, self.d_list):
+            out_ch = k_i * d_i
+            hidden = max(8, k_i * 2)
+            self.pose_gates.append(
+                nn.Sequential(
+                    nn.Conv2d(k_i, hidden, 1, bias=True),
+                    nn.SiLU(inplace=True),
+                    nn.Conv2d(hidden, out_ch, 1, bias=True),
+                )
+            )
+            self.gate_alpha.append(nn.Parameter(torch.tensor(0.5)))
+
+    def _split_pose_act(self, x: torch.Tensor, i: int) -> tuple[torch.Tensor, torch.Tensor]:
+        """Split one level packed tensor into pose and activation maps."""
+        k_i = self.k_list[i]
+        d_i = self.d_list[i]
+        b, c, h, w = x.shape
+        expected = k_i * (d_i + 1)
+        if c != expected:
+            raise ValueError(f"CapsuleDetectv1 level-{i} channel mismatch: got {c}, expected {expected}.")
+
+        caps = x.view(b, k_i, d_i + 1, h, w)
+        pose = caps[:, :, :d_i].reshape(b, k_i * d_i, h, w).contiguous()
+        act = caps[:, :, d_i].contiguous()
+        return pose, act
+
+    def _merge_pose(self, x: list[torch.Tensor]) -> list[torch.Tensor]:
+        merged = []
+        for i, xi in enumerate(x):
+            pose, act = self._split_pose_act(xi, i)
+            gate = torch.sigmoid(self.pose_gates[i](act))
+            # Residual gating keeps base pose information and improves stability.
+            pose = pose * (1.0 + self.gate_alpha[i] * gate)
+            merged.append(pose)
+        return merged
+
+    def forward_head(
+        self, x: list[torch.Tensor], box_head: torch.nn.Module = None, cls_head: torch.nn.Module = None
+    ) -> dict[str, torch.Tensor]:
+        if box_head is None or cls_head is None:
+            return dict()
+
+        pose_feats = self._merge_pose(x)
+        bs = pose_feats[0].shape[0]
+
+        box_list = [box_head[i](pose_feats[i]).view(bs, 4 * self.reg_max, -1) for i in range(self.nl)]
+        cls_list = [cls_head[i](pose_feats[i]).view(bs, self.nc, -1) for i in range(self.nl)]
+        boxes = torch.cat(box_list, dim=-1)
+        scores = torch.cat(cls_list, dim=-1)
+        return dict(boxes=boxes, scores=scores, feats=x)
+
+
+class CapsuleDetectv2(Detect):
+    """Capsule Detect v2: activation-gated pose + activation bypass for classification."""
+
+    def __init__(
+        self,
+        nc: int = 80,
+        *args,
+        reg_max: int = 16,
+        end2end: bool = False,
+        k: list[int] | tuple[int, ...] = (4, 8, 16),
+        d: list[int] | tuple[int, ...] = (16, 16, 16),
+        ch: tuple = (),
+    ):
+        parsed = list(args)
+        if parsed and isinstance(parsed[-1], (list, tuple)):
+            ch = tuple(parsed.pop(-1))
+
+        # Parser layout: [k_list, d_list, reg_max, end2end, ch]
+        if len(parsed) not in (2, 4):
+            raise ValueError("CapsuleDetectv2 expects [k_list, d_list, reg_max, end2end, ch].")
+
+        k, d = parsed[0], parsed[1]
+        if len(parsed) == 4:
+            reg_max = int(parsed[2])
+            end2end = bool(parsed[3])
+
+        if not isinstance(k, (list, tuple)) or not isinstance(d, (list, tuple)):
+            raise TypeError("CapsuleDetectv2 requires list/tuple k and d (per-level settings).")
+
+        ch = tuple(int(c) for c in ch)
+        nl = len(ch)
+        if len(k) != nl or len(d) != nl:
+            raise ValueError(f"CapsuleDetectv2 k/d length must equal number of levels ({nl}).")
+
+        self.k_list = tuple(int(v) for v in k)
+        self.d_list = tuple(int(v) for v in d)
+
+        # Input from neck is packed as K*(D+1): [pose(D), act(1)] repeated K types.
+        for i, c in enumerate(ch):
+            expected = self.k_list[i] * (self.d_list[i] + 1)
+            if c != expected:
+                raise ValueError(
+                    f"CapsuleDetectv2 level-{i} channel mismatch: got {c}, "
+                    f"expected {expected} from k={self.k_list[i]}, d={self.d_list[i]}."
+                )
+
+        # Detect heads consume merged pose channels: K*D.
+        merged_ch = tuple(k_i * d_i for k_i, d_i in zip(self.k_list, self.d_list))
+        super().__init__(nc=nc, reg_max=reg_max, end2end=end2end, ch=merged_ch)
+
+        self.pose_gates = nn.ModuleList()
+        self.gate_alpha = nn.ParameterList()
+        self.cls_bypass = nn.ModuleList()
+        self.cls_beta = nn.ParameterList()
+
+        for k_i, d_i in zip(self.k_list, self.d_list):
+            pose_ch = k_i * d_i
+            gate_hidden = max(8, k_i * 2)
+            self.pose_gates.append(
+                nn.Sequential(
+                    nn.Conv2d(k_i, gate_hidden, 1, bias=True),
+                    nn.SiLU(inplace=True),
+                    nn.Conv2d(gate_hidden, pose_ch, 1, bias=True),
+                )
+            )
+            self.gate_alpha.append(nn.Parameter(torch.tensor(0.5)))
+
+            cls_hidden = max(16, k_i * 2)
+            self.cls_bypass.append(
+                nn.Sequential(
+                    nn.Conv2d(k_i, cls_hidden, 1, bias=True),
+                    nn.SiLU(inplace=True),
+                    nn.Conv2d(cls_hidden, pose_ch, 1, bias=True),
+                )
+            )
+            self.cls_beta.append(nn.Parameter(torch.tensor(0.1)))
+
+    def _split_pose_act(self, x: torch.Tensor, i: int) -> tuple[torch.Tensor, torch.Tensor]:
+        """Split one level packed tensor into pose and activation maps."""
+        k_i = self.k_list[i]
+        d_i = self.d_list[i]
+        b, c, h, w = x.shape
+        expected = k_i * (d_i + 1)
+        if c != expected:
+            raise ValueError(f"CapsuleDetectv2 level-{i} channel mismatch: got {c}, expected {expected}.")
+
+        caps = x.view(b, k_i, d_i + 1, h, w)
+        pose = caps[:, :, :d_i].reshape(b, k_i * d_i, h, w).contiguous()
+        act = caps[:, :, d_i].contiguous()
+        return pose, act
+
+    def _fuse_pose(self, x: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+        box_feats, cls_feats = [], []
+        for i, xi in enumerate(x):
+            pose, act = self._split_pose_act(xi, i)
+            gate = torch.sigmoid(self.pose_gates[i](act))
+            pose_g = pose * (1.0 + self.gate_alpha[i] * gate)
+
+            # Classification bypass from activation channels.
+            act_skip = self.cls_bypass[i](act)
+            cls_in = pose_g + self.cls_beta[i] * act_skip
+
+            box_feats.append(pose_g)
+            cls_feats.append(cls_in)
+        return box_feats, cls_feats
+
+    def forward_head(
+        self, x: list[torch.Tensor], box_head: torch.nn.Module = None, cls_head: torch.nn.Module = None
+    ) -> dict[str, torch.Tensor]:
+        if box_head is None or cls_head is None:
+            return dict()
+
+        box_feats, cls_feats = self._fuse_pose(x)
+        bs = x[0].shape[0]
+
+        box_list = [box_head[i](box_feats[i]).view(bs, 4 * self.reg_max, -1) for i in range(self.nl)]
+        cls_list = [cls_head[i](cls_feats[i]).view(bs, self.nc, -1) for i in range(self.nl)]
+        boxes = torch.cat(box_list, dim=-1)
+        scores = torch.cat(cls_list, dim=-1)
+        return dict(boxes=boxes, scores=scores, feats=x)
+
+
+class CapsuleDetectv4(Detect):
+    """Capsule Detect v4: box uses raw pose, cls uses act bypass + symbolic type prior."""
+
+    def __init__(
+        self,
+        nc: int = 80,
+        *args,
+        reg_max: int = 16,
+        end2end: bool = False,
+        k: list[int] | tuple[int, ...] = (4, 8, 16),
+        d: list[int] | tuple[int, ...] = (16, 16, 16),
+        ch: tuple = (),
+    ):
+        parsed = list(args)
+        if parsed and isinstance(parsed[-1], (list, tuple)):
+            ch = tuple(parsed.pop(-1))
+
+        if len(parsed) not in (2, 4):
+            raise ValueError("CapsuleDetectv4 expects [k_list, d_list, reg_max, end2end, ch].")
+
+        k, d = parsed[0], parsed[1]
+        if len(parsed) == 4:
+            reg_max = int(parsed[2])
+            end2end = bool(parsed[3])
+
+        if not isinstance(k, (list, tuple)) or not isinstance(d, (list, tuple)):
+            raise TypeError("CapsuleDetectv4 requires list/tuple k and d (per-level settings).")
+
+        ch = tuple(int(c) for c in ch)
+        nl = len(ch)
+        if len(k) != nl or len(d) != nl:
+            raise ValueError(f"CapsuleDetectv4 k/d length must equal number of levels ({nl}).")
+
+        self.k_list = tuple(int(v) for v in k)
+        self.d_list = tuple(int(v) for v in d)
+
+        for i, c in enumerate(ch):
+            expected = self.k_list[i] * (self.d_list[i] + 1)
+            if c != expected:
+                raise ValueError(
+                    f"CapsuleDetectv4 level-{i} channel mismatch: got {c}, "
+                    f"expected {expected} from k={self.k_list[i]}, d={self.d_list[i]}."
+                )
+
+        merged_ch = tuple(k_i * d_i for k_i, d_i in zip(self.k_list, self.d_list))
+        super().__init__(nc=nc, reg_max=reg_max, end2end=end2end, ch=merged_ch)
+
+        self.cls_bypass = nn.ModuleList()
+        self.cls_beta = nn.ParameterList()
+        self.sym_prior = nn.ModuleList()
+        self.sym_beta = nn.ParameterList()
+
+        for k_i, d_i in zip(self.k_list, self.d_list):
+            pose_ch = k_i * d_i
+            cls_hidden = max(16, k_i * 2)
+            self.cls_bypass.append(
+                nn.Sequential(
+                    nn.Conv2d(k_i, cls_hidden, 1, bias=True),
+                    nn.SiLU(inplace=True),
+                    nn.Conv2d(cls_hidden, pose_ch, 1, bias=True),
+                )
+            )
+            self.cls_beta.append(nn.Parameter(torch.tensor(0.1)))
+            self.sym_prior.append(nn.Conv2d(k_i, self.nc, 1, bias=False))
+            self.sym_beta.append(nn.Parameter(torch.tensor(0.1)))
+
+    def _split_pose_act(self, x: torch.Tensor, i: int) -> tuple[torch.Tensor, torch.Tensor]:
+        k_i = self.k_list[i]
+        d_i = self.d_list[i]
+        b, c, h, w = x.shape
+        expected = k_i * (d_i + 1)
+        if c != expected:
+            raise ValueError(f"CapsuleDetectv4 level-{i} channel mismatch: got {c}, expected {expected}.")
+
+        caps = x.view(b, k_i, d_i + 1, h, w)
+        pose = caps[:, :, :d_i].reshape(b, k_i * d_i, h, w).contiguous()
+        act = caps[:, :, d_i].contiguous()
+        return pose, act
+
+    def _build_feats(self, x: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+        box_feats, cls_feats, cls_priors = [], [], []
+        for i, xi in enumerate(x):
+            pose, act = self._split_pose_act(xi, i)
+            cls_in = pose + self.cls_beta[i] * self.cls_bypass[i](act)
+            cls_prior = self.sym_beta[i] * self.sym_prior[i](act)
+            box_feats.append(pose)
+            cls_feats.append(cls_in)
+            cls_priors.append(cls_prior)
+        return box_feats, cls_feats, cls_priors
+
+    def forward_head(
+        self, x: list[torch.Tensor], box_head: torch.nn.Module = None, cls_head: torch.nn.Module = None
+    ) -> dict[str, torch.Tensor]:
+        if box_head is None or cls_head is None:
+            return dict()
+
+        box_feats, cls_feats, cls_priors = self._build_feats(x)
+        bs = x[0].shape[0]
+
+        box_list = [box_head[i](box_feats[i]).view(bs, 4 * self.reg_max, -1) for i in range(self.nl)]
+        cls_list = [
+            (cls_head[i](cls_feats[i]) + cls_priors[i]).view(bs, self.nc, -1)
+            for i in range(self.nl)
+        ]
+        boxes = torch.cat(box_list, dim=-1)
+        scores = torch.cat(cls_list, dim=-1)
+        return dict(boxes=boxes, scores=scores, feats=x)
+
+
+def _setup_capsule_layout(
+    k: list[int] | tuple[int, ...],
+    d: list[int] | tuple[int, ...],
+    ch: tuple,
+    cls_name: str,
+) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+    if not isinstance(k, (list, tuple)) or not isinstance(d, (list, tuple)):
+        raise TypeError(f"{cls_name} requires list/tuple k and d (per-level settings).")
+
+    ch = tuple(int(c) for c in ch)
+    nl = len(ch)
+    if len(k) != nl or len(d) != nl:
+        raise ValueError(f"{cls_name} k/d length must equal number of levels ({nl}).")
+
+    k_list = tuple(int(v) for v in k)
+    d_list = tuple(int(v) for v in d)
+    for i, c in enumerate(ch):
+        expected = k_list[i] * (d_list[i] + 1)
+        if c != expected:
+            raise ValueError(
+                f"{cls_name} level-{i} channel mismatch: got {c}, "
+                f"expected {expected} from k={k_list[i]}, d={d_list[i]}."
+            )
+    merged_ch = tuple(k_i * d_i for k_i, d_i in zip(k_list, d_list))
+    return k_list, d_list, merged_ch
+
+
+def _init_capsule_semantic_heads(obj: nn.Module) -> None:
+    obj.cls_bypass = nn.ModuleList()
+    obj.cls_beta = nn.ParameterList()
+    obj.sym_prior = nn.ModuleList()
+    obj.sym_norm = nn.ModuleList()
+    obj.sym_dropout = nn.ModuleList()
+    obj.sym_beta = nn.ParameterList()
+
+    for k_i, d_i in zip(obj.k_list, obj.d_list):
+        pose_ch = k_i * d_i
+        cls_hidden = max(16, k_i * 2)
+        obj.cls_bypass.append(
+            nn.Sequential(
+                nn.Conv2d(k_i, cls_hidden, 1, bias=True),
+                nn.SiLU(inplace=True),
+                nn.Conv2d(cls_hidden, pose_ch, 1, bias=True),
+            )
+        )
+        obj.cls_beta.append(nn.Parameter(torch.tensor(0.1)))
+        obj.sym_dropout.append(nn.Dropout2d(p=0.1))
+        obj.sym_prior.append(nn.Conv2d(k_i, obj.nc, 1, bias=False))
+        obj.sym_norm.append(nn.GroupNorm(1, obj.nc))
+        obj.sym_beta.append(nn.Parameter(torch.tensor(0.1)))
+
+
+def _capsule_split_pose_act(
+    x: torch.Tensor,
+    k_i: int,
+    d_i: int,
+    cls_name: str,
+    level_i: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    b, c, h, w = x.shape
+    expected = k_i * (d_i + 1)
+    if c != expected:
+        raise ValueError(f"{cls_name} level-{level_i} channel mismatch: got {c}, expected {expected}.")
+    caps = x.view(b, k_i, d_i + 1, h, w)
+    pose = caps[:, :, :d_i].reshape(b, k_i * d_i, h, w).contiguous()
+    act = caps[:, :, d_i].contiguous()
+    return pose, act
+
+
+def _capsule_build_feats(obj: nn.Module, x: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+    box_feats, cls_feats, cls_priors = [], [], []
+    cls_name = obj.__class__.__name__
+    for i, xi in enumerate(x):
+        pose, act = _capsule_split_pose_act(xi, obj.k_list[i], obj.d_list[i], cls_name, i)
+        cls_scale = torch.tanh(obj.cls_beta[i])
+        cls_in = pose + cls_scale * obj.cls_bypass[i](act)
+        act_s = obj.sym_dropout[i](act)
+        prior = obj.sym_prior[i](act_s)
+        prior = obj.sym_norm[i](prior)
+        prior = prior - prior.mean(dim=1, keepdim=True)
+        sym_scale = torch.tanh(obj.sym_beta[i])
+        cls_prior = sym_scale * prior
+        box_feats.append(pose)
+        cls_feats.append(cls_in)
+        cls_priors.append(cls_prior)
+    return box_feats, cls_feats, cls_priors
+
+
+def _capsule_build_feats_gated(
+    obj: nn.Module, x: list[torch.Tensor]
+) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+    box_feats, cls_feats, cls_priors = [], [], []
+    cls_name = obj.__class__.__name__
+    for i, xi in enumerate(x):
+        pose, act = _capsule_split_pose_act(xi, obj.k_list[i], obj.d_list[i], cls_name, i)
+        cls_scale = torch.tanh(obj.cls_beta[i])
+        gate = torch.sigmoid(obj.cls_bypass[i](act))
+        cls_in = pose * (1.0 + cls_scale * gate)
+
+        act_s = obj.sym_dropout[i](act)
+        prior = obj.sym_prior[i](act_s)
+        prior = obj.sym_norm[i](prior)
+        prior = prior - prior.mean(dim=1, keepdim=True)
+        sym_scale = torch.tanh(obj.sym_beta[i])
+        cls_prior = sym_scale * prior
+
+        box_feats.append(pose)
+        cls_feats.append(cls_in)
+        cls_priors.append(cls_prior)
+    return box_feats, cls_feats, cls_priors
+
+
+def _capsule_build_feats_boxcls(
+    obj: nn.Module, x: list[torch.Tensor]
+) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+    box_feats, cls_feats, cls_priors = [], [], []
+    cls_name = obj.__class__.__name__
+    for i, xi in enumerate(x):
+        pose, act = _capsule_split_pose_act(xi, obj.k_list[i], obj.d_list[i], cls_name, i)
+        act_s = obj.sym_dropout[i](act)
+        prior = obj.sym_prior[i](act_s)
+        prior = obj.sym_norm[i](prior)
+        prior = prior - prior.mean(dim=1, keepdim=True)
+        sym_scale = torch.tanh(obj.sym_beta[i])
+        cls_prior = sym_scale * prior
+
+        box_feats.append(pose)
+        cls_feats.append(pose)
+        cls_priors.append(cls_prior)
+    return box_feats, cls_feats, cls_priors
+
+
+def _capsule_build_feats_boxcls_simpleprior(
+    obj: nn.Module, x: list[torch.Tensor]
+) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+    box_feats, cls_feats, cls_priors = [], [], []
+    cls_name = obj.__class__.__name__
+    for i, xi in enumerate(x):
+        pose, act = _capsule_split_pose_act(xi, obj.k_list[i], obj.d_list[i], cls_name, i)
+        act_s = obj.sym_dropout[i](act)
+        prior = obj.sym_prior[i](act_s)
+        sym_scale = torch.tanh(obj.sym_beta[i])
+        cls_prior = sym_scale * prior
+
+        box_feats.append(pose)
+        cls_feats.append(pose)
+        cls_priors.append(cls_prior)
+    return box_feats, cls_feats, cls_priors
+
+
+def _capsule_build_feats_open_vocab(
+    obj: nn.Module, x: list[torch.Tensor]
+) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+    box_feats, cls_feats, acts = [], [], []
+    cls_name = obj.__class__.__name__
+    for i, xi in enumerate(x):
+        pose, act = _capsule_split_pose_act(xi, obj.k_list[i], obj.d_list[i], cls_name, i)
+        cls_in = pose
+        if getattr(obj, "with_act_gate", False):
+            cls_scale = torch.tanh(obj.ov_beta[i])
+            gate = torch.sigmoid(obj.ov_gate[i](act))
+            cls_in = pose * (1.0 + cls_scale * gate)
+        box_feats.append(pose)
+        cls_feats.append(cls_in)
+        acts.append(act)
+    return box_feats, cls_feats, acts
+
+
+class CapsuleDetectv5(Detect):
+    """Capsule Detect v5: box uses raw pose, cls uses stabilized symbolic prior."""
+
+    def __init__(
+        self,
+        nc: int = 80,
+        *args,
+        reg_max: int = 16,
+        end2end: bool = False,
+        k: list[int] | tuple[int, ...] = (4, 8, 16),
+        d: list[int] | tuple[int, ...] = (16, 16, 16),
+        ch: tuple = (),
+    ):
+        parsed = list(args)
+        if parsed and isinstance(parsed[-1], (list, tuple)):
+            ch = tuple(parsed.pop(-1))
+
+        if len(parsed) not in (2, 4):
+            raise ValueError("CapsuleDetectv5 expects [k_list, d_list, reg_max, end2end, ch].")
+
+        k, d = parsed[0], parsed[1]
+        if len(parsed) == 4:
+            reg_max = int(parsed[2])
+            end2end = bool(parsed[3])
+
+        self.k_list, self.d_list, merged_ch = _setup_capsule_layout(k, d, ch, "CapsuleDetectv5")
+        super().__init__(nc=nc, reg_max=reg_max, end2end=end2end, ch=merged_ch)
+        _init_capsule_semantic_heads(self)
+
+    def _split_pose_act(self, x: torch.Tensor, i: int) -> tuple[torch.Tensor, torch.Tensor]:
+        return _capsule_split_pose_act(x, self.k_list[i], self.d_list[i], "CapsuleDetectv5", i)
+
+    def _build_feats(self, x: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+        return _capsule_build_feats(self, x)
+
+    def forward_head(
+        self, x: list[torch.Tensor], box_head: torch.nn.Module = None, cls_head: torch.nn.Module = None
+    ) -> dict[str, torch.Tensor]:
+        if box_head is None or cls_head is None:
+            return dict()
+
+        box_feats, cls_feats, cls_priors = self._build_feats(x)
+        bs = x[0].shape[0]
+
+        box_list = [box_head[i](box_feats[i]).view(bs, 4 * self.reg_max, -1) for i in range(self.nl)]
+        cls_list = [
+            (cls_head[i](cls_feats[i]) + cls_priors[i]).view(bs, self.nc, -1)
+            for i in range(self.nl)
+        ]
+        boxes = torch.cat(box_list, dim=-1)
+        scores = torch.cat(cls_list, dim=-1)
+        return dict(boxes=boxes, scores=scores, feats=x)
+
+
+class CapsuleDetectv6(CapsuleDetectv5):
+    """Capsule Detect v6: replace additive cls correction with multiplicative act gate."""
+
+    def _build_feats(self, x: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+        return _capsule_build_feats_gated(self, x)
+
+
+class CapsuleDetectv7(CapsuleDetectv5):
+    """Capsule Detect v7: cls head consumes raw pose features plus symbolic priors only."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.profile_head = False
+        self._head_profile: dict[str, float] = {}
+        self._head_profile_calls = 0
+
+    def _ensure_profile_attrs(self) -> None:
+        if not hasattr(self, "profile_head"):
+            self.profile_head = False
+        if not hasattr(self, "_head_profile"):
+            self._head_profile = {}
+        if not hasattr(self, "_head_profile_calls"):
+            self._head_profile_calls = 0
+
+    def reset_head_profile(self) -> None:
+        self._ensure_profile_attrs()
+        self._head_profile = {
+            "split_pose_act_ms": 0.0,
+            "cls_prior_ms": 0.0,
+            "box_head_ms": 0.0,
+            "cls_head_ms": 0.0,
+            "cat_ms": 0.0,
+        }
+        self._head_profile_calls = 0
+
+    def get_head_profile(self) -> dict[str, float]:
+        self._ensure_profile_attrs()
+        if not self._head_profile:
+            return {}
+        out = dict(self._head_profile)
+        calls = max(self._head_profile_calls, 1)
+        out["calls"] = float(self._head_profile_calls)
+        out["total_ms"] = sum(v for k, v in out.items() if k.endswith("_ms"))
+        for key in list(self._head_profile):
+            out[key.replace("_ms", "_avg_ms")] = self._head_profile[key] / calls
+        return out
+
+    def _sync_profile(self) -> None:
+        self._ensure_profile_attrs()
+        if self.profile_head and torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+    def _build_feats(self, x: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+        self._ensure_profile_attrs()
+        if not self.profile_head:
+            return _capsule_build_feats_boxcls(self, x)
+
+        if not self._head_profile:
+            self.reset_head_profile()
+
+        box_feats, cls_feats, cls_priors = [], [], []
+        cls_name = self.__class__.__name__
+        for i, xi in enumerate(x):
+            self._sync_profile()
+            t0 = time.perf_counter()
+            pose, act = _capsule_split_pose_act(xi, self.k_list[i], self.d_list[i], cls_name, i)
+            self._sync_profile()
+            self._head_profile["split_pose_act_ms"] += (time.perf_counter() - t0) * 1000.0
+
+            self._sync_profile()
+            t0 = time.perf_counter()
+            act_s = self.sym_dropout[i](act)
+            prior = self.sym_prior[i](act_s)
+            prior = self.sym_norm[i](prior)
+            prior = prior - prior.mean(dim=1, keepdim=True)
+            sym_scale = torch.tanh(self.sym_beta[i])
+            cls_prior = sym_scale * prior
+            self._sync_profile()
+            self._head_profile["cls_prior_ms"] += (time.perf_counter() - t0) * 1000.0
+
+            box_feats.append(pose)
+            cls_feats.append(pose)
+            cls_priors.append(cls_prior)
+        return box_feats, cls_feats, cls_priors
+
+    def forward_head(
+        self, x: list[torch.Tensor], box_head: torch.nn.Module = None, cls_head: torch.nn.Module = None
+    ) -> dict[str, torch.Tensor]:
+        self._ensure_profile_attrs()
+        if box_head is None or cls_head is None:
+            return dict()
+
+        box_feats, cls_feats, cls_priors = self._build_feats(x)
+        bs = x[0].shape[0]
+
+        if not self.profile_head:
+            box_list = [box_head[i](box_feats[i]).view(bs, 4 * self.reg_max, -1) for i in range(self.nl)]
+            cls_list = [
+                (cls_head[i](cls_feats[i]) + cls_priors[i]).view(bs, self.nc, -1)
+                for i in range(self.nl)
+            ]
+            boxes = torch.cat(box_list, dim=-1)
+            scores = torch.cat(cls_list, dim=-1)
+            return dict(boxes=boxes, scores=scores, feats=x)
+
+        if not self._head_profile:
+            self.reset_head_profile()
+        self._head_profile_calls += 1
+
+        box_list, cls_list = [], []
+        for i in range(self.nl):
+            self._sync_profile()
+            t0 = time.perf_counter()
+            box_i = box_head[i](box_feats[i]).view(bs, 4 * self.reg_max, -1)
+            self._sync_profile()
+            self._head_profile["box_head_ms"] += (time.perf_counter() - t0) * 1000.0
+            box_list.append(box_i)
+
+            self._sync_profile()
+            t0 = time.perf_counter()
+            cls_i = (cls_head[i](cls_feats[i]) + cls_priors[i]).view(bs, self.nc, -1)
+            self._sync_profile()
+            self._head_profile["cls_head_ms"] += (time.perf_counter() - t0) * 1000.0
+            cls_list.append(cls_i)
+
+        self._sync_profile()
+        t0 = time.perf_counter()
+        boxes = torch.cat(box_list, dim=-1)
+        scores = torch.cat(cls_list, dim=-1)
+        self._sync_profile()
+        self._head_profile["cat_ms"] += (time.perf_counter() - t0) * 1000.0
+        return dict(boxes=boxes, scores=scores, feats=x)
+
+
+class CapsuleDetectv8(CapsuleDetectv5):
+    """Capsule Detect v8: raw pose cls path with simplified cls_prior (no norm/centering)."""
+
+    def _build_feats(self, x: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+        return _capsule_build_feats_boxcls_simpleprior(self, x)
+
+
+class CapsuleOpenVocabDetect(Detect):
+    """Capsule detection head with open-vocabulary classification via text embedding matching."""
+
+    def __init__(
+        self,
+        nc: int = 80,
+        *args,
+        reg_max: int = 16,
+        end2end: bool = False,
+        embed: int = 256,
+        with_act_gate: bool = False,
+        with_objectness_prior: bool = True,
+        k: list[int] | tuple[int, ...] = (4, 8, 16),
+        d: list[int] | tuple[int, ...] = (16, 16, 16),
+        ch: tuple = (),
+    ):
+        parsed = list(args)
+        if parsed and isinstance(parsed[-1], (list, tuple)):
+            ch = tuple(parsed.pop(-1))
+
+        if len(parsed) not in (2, 4, 7):
+            raise ValueError(
+                "CapsuleOpenVocabDetect expects [k_list, d_list, (reg_max, end2end, embed, with_act_gate, "
+                "with_objectness_prior), ch]."
+            )
+
+        k, d = parsed[0], parsed[1]
+        if len(parsed) == 4:
+            reg_max = int(parsed[2])
+            end2end = bool(parsed[3])
+        elif len(parsed) == 7:
+            # Support both direct args order:
+            #   [k_list, d_list, reg_max, end2end, embed, with_act_gate, with_objectness_prior]
+            # and parser-appended order:
+            #   [k_list, d_list, embed, with_act_gate, with_objectness_prior, reg_max, end2end]
+            if type(parsed[3]) is bool and type(parsed[4]) is bool and type(parsed[6]) is bool:
+                embed = int(parsed[2])
+                with_act_gate = bool(parsed[3])
+                with_objectness_prior = bool(parsed[4])
+                reg_max = int(parsed[5])
+                end2end = bool(parsed[6])
+            else:
+                reg_max = int(parsed[2])
+                end2end = bool(parsed[3])
+                embed = int(parsed[4])
+                with_act_gate = bool(parsed[5])
+                with_objectness_prior = bool(parsed[6])
+
+        self.k_list, self.d_list, merged_ch = _setup_capsule_layout(k, d, ch, "CapsuleOpenVocabDetect")
+        super().__init__(nc=nc, reg_max=reg_max, end2end=end2end, ch=merged_ch)
+
+        self.embed = int(embed)
+        self.with_act_gate = bool(with_act_gate)
+        self.with_objectness_prior = bool(with_objectness_prior)
+
+        self.emb_head = nn.ModuleList()
+        self.ov_gate = nn.ModuleList()
+        self.ov_beta = nn.ParameterList()
+        self.obj_prior = nn.ModuleList()
+
+        for k_i, d_i in zip(self.k_list, self.d_list):
+            pose_ch = k_i * d_i
+            self.emb_head.append(
+                nn.Sequential(
+                    Conv(pose_ch, pose_ch, 3),
+                    DWConv(pose_ch, pose_ch, 3),
+                    nn.Conv2d(pose_ch, self.embed, 1, bias=True),
+                )
+            )
+            if self.with_act_gate:
+                hidden = max(16, k_i * 2)
+                self.ov_gate.append(
+                    nn.Sequential(
+                        nn.Conv2d(k_i, hidden, 1, bias=True),
+                        nn.SiLU(inplace=True),
+                        nn.Conv2d(hidden, pose_ch, 1, bias=True),
+                    )
+                )
+                self.ov_beta.append(nn.Parameter(torch.tensor(0.1)))
+            else:
+                self.ov_gate.append(nn.Identity())
+                self.ov_beta.append(nn.Parameter(torch.tensor(0.0), requires_grad=False))
+
+            if self.with_objectness_prior:
+                self.obj_prior.append(nn.Conv2d(k_i, 1, 1, bias=True))
+            else:
+                self.obj_prior.append(nn.Identity())
+
+        self.logit_scale = nn.Parameter(torch.tensor(math.log(1 / 0.07), dtype=torch.float32))
+        self.register_buffer("cached_text_embeddings", torch.empty(0), persistent=False)
+
+    def set_text_embeddings(self, text_embs: torch.Tensor | None) -> None:
+        """Cache normalized text embeddings for inference."""
+        if text_embs is None:
+            self.cached_text_embeddings = torch.empty(0, device=self.logit_scale.device)
+            return
+        if text_embs.ndim != 2:
+            raise ValueError(f"text_embs must be 2D [num_classes, embed_dim], got shape {tuple(text_embs.shape)}.")
+        self.cached_text_embeddings = F.normalize(text_embs.detach().to(self.logit_scale.device), dim=-1)
+
+    def _split_pose_act(self, x: torch.Tensor, i: int) -> tuple[torch.Tensor, torch.Tensor]:
+        return _capsule_split_pose_act(x, self.k_list[i], self.d_list[i], "CapsuleOpenVocabDetect", i)
+
+    def _build_feats(self, x: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+        return _capsule_build_feats_open_vocab(self, x)
+
+    def _prepare_text_embeddings(self, text_embs: torch.Tensor | None, bs: int, device: torch.device) -> torch.Tensor | None:
+        if text_embs is None:
+            if self.cached_text_embeddings.numel() == 0:
+                return None
+            text = self.cached_text_embeddings
+        else:
+            text = text_embs
+
+        if text.ndim == 2:
+            text = text.unsqueeze(0).expand(bs, -1, -1)
+        elif text.ndim != 3:
+            raise ValueError(f"text_embs must be 2D or 3D, got shape {tuple(text.shape)}.")
+
+        if text.shape[-1] != self.embed:
+            raise ValueError(f"text_embs last dim must equal embed={self.embed}, got {text.shape[-1]}.")
+        return F.normalize(text.to(device=device, dtype=self.logit_scale.dtype), dim=-1)
+
+    def _compute_ov_scores(
+        self, cls_feats: list[torch.Tensor], acts: list[torch.Tensor], text_embs: torch.Tensor | None
+    ) -> tuple[torch.Tensor | None, list[torch.Tensor], torch.Tensor | None]:
+        bs = cls_feats[0].shape[0]
+        level_embeddings = []
+        for i in range(self.nl):
+            emb = self.emb_head[i](cls_feats[i])
+            if self.with_objectness_prior:
+                emb = emb * (1.0 + torch.sigmoid(self.obj_prior[i](acts[i])))
+            level_embeddings.append(emb)
+
+        text = self._prepare_text_embeddings(text_embs, bs, cls_feats[0].device)
+        if text is None:
+            return None, level_embeddings, None
+
+        visual_tokens = torch.cat(
+            [F.normalize(emb.flatten(2).transpose(1, 2), dim=-1) for emb in level_embeddings],
+            dim=1,
+        )
+        scale = self.logit_scale.exp().clamp(max=100.0)
+        scores = torch.einsum("bnd,bcd->bcn", visual_tokens, text) * scale
+        return scores, level_embeddings, text
+
+    def forward_head(
+        self,
+        x: list[torch.Tensor],
+        text_embs: torch.Tensor | None = None,
+        box_head: torch.nn.Module = None,
+        cls_head: torch.nn.Module = None,
+    ) -> dict[str, torch.Tensor]:
+        del cls_head  # fixed-class cls head is unused in open-vocabulary mode
+        if box_head is None:
+            return dict()
+
+        box_feats, cls_feats, acts = self._build_feats(x)
+        bs = x[0].shape[0]
+        boxes = torch.cat([box_head[i](box_feats[i]).view(bs, 4 * self.reg_max, -1) for i in range(self.nl)], dim=-1)
+        scores, level_embeddings, text = self._compute_ov_scores(cls_feats, acts, text_embs)
+
+        preds = {
+            "boxes": boxes,
+            "embeddings": level_embeddings,
+            "cls_feats": cls_feats,
+            "acts": acts,
+            "feats": x,
+        }
+        if scores is not None:
+            preds["scores"] = scores
+            preds["text_embeddings"] = text
+        return preds
+
+    def forward(
+        self, x: list[torch.Tensor], text_embs: torch.Tensor | None = None
+    ) -> dict[str, torch.Tensor] | torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        preds = self.forward_head(x, text_embs=text_embs, **self.one2many)
+        if self.end2end:
+            x_detach = [xi.detach() for xi in x]
+            one2one = self.forward_head(x_detach, text_embs=text_embs, **self.one2one)
+            preds = {"one2many": preds, "one2one": one2one}
+        if self.training:
+            return preds
+
+        infer_preds = preds["one2one"] if self.end2end else preds
+        if "scores" not in infer_preds:
+            raise ValueError("CapsuleOpenVocabDetect inference requires text_embs or cached text embeddings.")
+
+        original_nc = self.nc
+        self.nc = int(infer_preds["scores"].shape[1])
+        try:
+            y = self._inference(infer_preds)
+            if self.end2end:
+                y = self.postprocess(y.permute(0, 2, 1))
+        finally:
+            self.nc = original_nc
+        return y if self.export else (y, preds)
+
+
+class CapsuleSegmentv1(Segment):
+    """Capsule-style Segment head aligned with CapsuleDetectv6 semantics."""
+
+    def __init__(
+        self,
+        nc: int = 80,
+        *args,
+        nm: int = 32,
+        npr: int = 256,
+        reg_max: int = 16,
+        end2end: bool = False,
+        k: list[int] | tuple[int, ...] = (4, 8, 16),
+        d: list[int] | tuple[int, ...] = (16, 16, 16),
+        ch: tuple = (),
+    ):
+        parsed = list(args)
+        if parsed and isinstance(parsed[-1], (list, tuple)):
+            ch = tuple(parsed.pop(-1))
+
+        if len(parsed) not in (2, 4, 6):
+            raise ValueError("CapsuleSegmentv1 expects [k_list, d_list, (nm, npr), reg_max, end2end, ch].")
+
+        k, d = parsed[0], parsed[1]
+        if len(parsed) == 4:
+            if isinstance(parsed[3], bool):
+                reg_max = int(parsed[2])
+                end2end = bool(parsed[3])
+            else:
+                nm = int(parsed[2])
+                npr = int(parsed[3])
+        elif len(parsed) == 6:
+            nm = int(parsed[2])
+            npr = int(parsed[3])
+            reg_max = int(parsed[4])
+            end2end = bool(parsed[5])
+
+        self.k_list, self.d_list, merged_ch = _setup_capsule_layout(k, d, ch, "CapsuleSegmentv1")
+        super().__init__(nc=nc, nm=nm, npr=npr, reg_max=reg_max, end2end=end2end, ch=merged_ch)
+        _init_capsule_semantic_heads(self)
+        self.proto = Proto26(merged_ch, self.npr, self.nm, nc)
+
+    def _split_pose_act(self, x: torch.Tensor, i: int) -> tuple[torch.Tensor, torch.Tensor]:
+        return _capsule_split_pose_act(x, self.k_list[i], self.d_list[i], "CapsuleSegmentv1", i)
+
+    def _build_feats(self, x: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+        return _capsule_build_feats_gated(self, x)
+
+    def forward_head(
+        self,
+        x: list[torch.Tensor],
+        box_head: torch.nn.Module = None,
+        cls_head: torch.nn.Module = None,
+        mask_head: torch.nn.Module = None,
+    ) -> dict[str, torch.Tensor]:
+        if box_head is None or cls_head is None:
+            return dict()
+
+        box_feats, cls_feats, cls_priors = self._build_feats(x)
+        bs = x[0].shape[0]
+        boxes = torch.cat([box_head[i](box_feats[i]).view(bs, 4 * self.reg_max, -1) for i in range(self.nl)], dim=-1)
+        scores = torch.cat(
+            [(cls_head[i](cls_feats[i]) + cls_priors[i]).view(bs, self.nc, -1) for i in range(self.nl)],
+            dim=-1,
+        )
+        preds = dict(boxes=boxes, scores=scores, feats=cls_feats)
+        if mask_head is not None:
+            preds["mask_coefficient"] = torch.cat(
+                [mask_head[i](cls_feats[i]).view(bs, self.nm, -1) for i in range(self.nl)],
+                dim=-1,
+            )
+        return preds
+
+    def forward(self, x: list[torch.Tensor]) -> tuple | list[torch.Tensor] | dict[str, torch.Tensor]:
+        _, cls_feats, _ = self._build_feats(x)
+        outputs = Detect.forward(self, x)
+        preds = outputs[1] if isinstance(outputs, tuple) else outputs
+        proto_in = cls_feats
+        proto = self.proto(proto_in)  # multi-level Proto26 over merged capsule features
+        if isinstance(preds, dict):
+            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)
+
+
+class CapsuleSegmentv2(CapsuleSegmentv1):
+    """Capsule Segment v2: cls head consumes raw pose features and symbolic priors only."""
+
+    def _build_feats(self, x: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+        return _capsule_build_feats_boxcls(self, x)
+
+
+class CapsuleSegmentv3(CapsuleSegmentv1):
+    """Capsule Segment v3: raw pose cls path with simplified cls_prior (no norm/centering)."""
+
+    def _build_feats(self, x: list[torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]:
+        return _capsule_build_feats_boxcls_simpleprior(self, x)
+

modules/neck.py

@@ -0,0 +1,769 @@
+"""
+CapsNeck: efficient capsule-style neck blocks for Ultralytics YAML models.
+
+Design intent:
+- Keep capsule semantics (type/channel grouping + routing-style fusion).
+- Stay lightweight and export-friendly for detection training/inference.
+- Avoid expensive iterative EM/dynamic routing inside the neck path.
+
+This neck is "capsule-style" rather than a full matrix-capsule network:
+1) CapsProj  : CNN feature -> packed capsules (K types * D dims)
+2) CapsAlign : scale alignment between pyramid levels (no global context)
+3) CapsRoute : efficient self-routing proxy across sources (softmax source gating)
+4) CapsDecode: packed capsules -> standard feature map for Detect
+5) CapsuleTap: optional pass-through cache hook for analysis/aux losses
+
+Note:
+- Routing here is source-level and single-step by default (iters=1), chosen for speed.
+- If stronger capsule routing is needed, it should be added in the head where cost is lower.
+"""
+
+
+from __future__ import annotations
+
+from typing import List, Optional, Tuple, Union
+
+import math
+import time
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ultralytics.nn.modules import C3k2, Conv, DWConv
+
+
+# -------------------------
+# 1) CapsProj
+# -------------------------
+
+class CapsProj(nn.Module):
+    """
+    Project a standard feature map into packed capsule channels using one C3k2 block.
+
+    Input:  x  [B, C, H, W]
+    Output: u  [B, K*(D+1), H, W]
+
+    Args:
+      K: number of capsule types
+      D: capsule pose dimension per type
+      mix/mix_kernel: kept for backward YAML compatibility (unused)
+    """
+
+    def __init__(self, c1: int, K: int = 4, D: int = 16):
+        super().__init__()
+        self.K = int(K)
+        self.D = int(D)
+        self.c_out = self.K * (self.D + 1)
+
+        # Use a single C3k2 block as the capsule projection operator.
+        self.map = C3k2(c1, self.c_out, n=1, c3k=False, e=0.5, g=1, shortcut=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.map(x)
+
+
+# -------------------------
+# 2) CapsAlign (no context)
+# -------------------------
+
+class CapsAlign(nn.Module):
+    """
+    Align packed capsules across pyramid levels with YOLO-style ops.
+
+    - Upsampling uses ``nn.Upsample(scale_factor=2, mode='nearest')``.
+    - Downsampling uses stride-2 ``Conv`` blocks.
+
+    Args:
+      c1: input/output channel count.
+      src_level: source pyramid level in {3,4,5}.
+      tgt_level: target pyramid level in {3,4,5}.
+      down_groups: groups for downsample Conv.
+        Use capsule-type count K to keep each capsule block isolated.
+    """
+
+    def __init__(self, c1: int, src_level: int, tgt_level: int, down_groups: int = 1):
+        super().__init__()
+        self.c1 = int(c1)
+        self.src_level = int(src_level)
+        self.tgt_level = int(tgt_level)
+        self.down_groups = int(down_groups)
+
+        if self.src_level not in (3, 4, 5) or self.tgt_level not in (3, 4, 5):
+            raise ValueError("CapsAlign levels must be in {3,4,5}.")
+
+        if self.down_groups < 1 or self.c1 % self.down_groups != 0:
+            raise ValueError(f"CapsAlign down_groups={self.down_groups} must divide c1={self.c1}.")
+
+        steps = abs(self.src_level - self.tgt_level)
+        if self.src_level == self.tgt_level:
+            self.mode = 'identity'
+            self.ops = nn.ModuleList()
+        elif self.src_level > self.tgt_level:
+            self.mode = 'up'
+            # YOLO-style top-down path: nearest-neighbor upsample x2 per level.
+            self.ops = nn.ModuleList(nn.Upsample(scale_factor=2, mode='nearest') for _ in range(steps))
+        else:
+            self.mode = 'down'
+            # YOLO-style bottom-up path: stride-2 grouped Conv per level.
+            self.ops = nn.ModuleList(Conv(self.c1, self.c1, 3, 2, g=self.down_groups) for _ in range(steps))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.mode == 'identity':
+            return x
+
+        for op in self.ops:
+            x = op(x)
+        return x
+
+
+# -------------------------
+# 3) CapsRoute (light, parser-friendly)
+# -------------------------
+
+class ConvSelfRouting(nn.Module):
+    """Grouped-conv self-routing over stacked capsule sources.
+
+    Args:
+      K_in: input capsule type count.
+      P_in: input pose dimension.
+      K_out: output capsule type count.
+      P_out: output pose dimension.
+      kernel_size: grouped conv kernel for local capsule mixing.
+    """
+
+    def __init__(self, K_in: int, P_in: int, K_out: int, P_out: int, kernel_size: int = 3):
+        super().__init__()
+        self.K_in = int(K_in)
+        self.P_in = int(P_in)
+        self.K_out = int(K_out)
+        self.P_out = int(P_out)
+
+        if min(self.K_in, self.P_in, self.K_out, self.P_out) <= 0:
+            raise ValueError('ConvSelfRouting expects positive K/P values.')
+
+        self.c_in = self.K_in * (self.P_in + 1)
+        self.c_out = self.K_out * (self.P_out + 1)
+
+        k = int(kernel_size)
+        padding = k//2
+        self.mix = nn.Conv2d(self.c_in, self.c_in, kernel_size=k, stride=1, padding=padding, groups=self.K_in, bias=False)
+        self.gate = nn.Conv2d(self.c_in, self.K_in, kernel_size=1, stride=1, padding=0, groups=self.K_in, bias=True)
+
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B,C,H,W], C = K_in*(P_in+1)
+
+        b, c, h, w = x.shape
+        if c != self.c_in:
+            raise ValueError(f'ConvSelfRouting expected C={self.c_in}, got C={c}')
+
+        mixed = self.mix(x)
+        logits = self.gate(mixed).reshape(b, self.K_in, h, w)
+        weights = logits.softmax(dim=1)
+
+        caps = mixed.reshape(b, self.K_in, self.P_in + 1, h, w)
+        routed = weights.unsqueeze(2) * caps
+        routed = routed.reshape(b, self.c_in, h, w)
+
+        return routed
+
+
+class SelfRouting(nn.Module):
+    """Pose-transform self-routing on packed capsule tensor.
+
+    Args:
+      K_in: input capsule type count.
+      P_in: input pose dimension.
+      K_out: output capsule type count.
+      P_out: output pose dimension.
+
+    Input:
+      x: [B, K_in*(P_in+1), H, W]
+
+    Output:
+      y: [B, K_out*(P_out+1), H, W]
+    """
+
+    def __init__(self, K_in: int, P_in: int, K_out: int, P_out: int):
+        super().__init__()
+        self.K_in = int(K_in)
+        self.P_in = int(P_in)
+        self.K_out = int(K_out)
+        self.P_out = int(P_out)
+        if min(self.K_in, self.P_in, self.K_out, self.P_out) <= 0:
+            raise ValueError('SelfRouting expects positive K/P values.')
+
+        self.c_in = self.K_in * (self.P_in + 1)
+        self.c_out = self.K_out * (self.P_out + 1)
+        self.eps = 1e-6
+
+        self.W_pose = nn.Parameter(torch.empty(self.K_in, self.K_out, self.P_in, self.P_out))
+        nn.init.kaiming_uniform_(self.W_pose, a=math.sqrt(5))
+        self.W_gate = nn.Parameter(torch.zeros(self.K_in, self.K_out, self.P_in))
+        self.b_gate = nn.Parameter(torch.zeros(1, self.K_in, self.K_out, 1, 1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [B, C, H, W], C = K_in*(P_in+1)
+        if x.ndim != 4:
+            raise TypeError(f'SelfRouting expects [B,C,H,W], got {tuple(x.shape)}')
+
+        b, c, h, w = x.shape
+        if c != self.c_in:
+            raise ValueError(f'SelfRouting expected C={self.c_in}, got C={c}')
+
+        # Packed capsule layout is interleaved per type: [pose(P), act(1)].
+        # x_caps: [B, K_in, P_in+1, H, W]
+        x_caps = x.reshape(b, self.K_in, self.P_in + 1, h, w)
+        pose = x_caps[:, :, :self.P_in]  # [B, K_in, P_in, H, W]
+        act = x_caps[:, :, self.P_in : self.P_in + 1].sigmoid()  # [B, K_in, 1, H, W]
+
+        # votes: [B, K_in, K_out, H, W, P_out]
+        votes = torch.einsum('bkphw,kopq->bkohwq', pose, self.W_pose)
+        # logits/weights: [B, K_in, K_out, H, W]
+        logits = torch.einsum('bkphw,kop->bkohw', pose, self.W_gate) + self.b_gate
+        weights = logits.softmax(dim=2)
+
+        ar = weights * act  # [B, K_in, K_out, H, W]
+        ar_sum = ar.sum(dim=1, keepdim=True) + self.eps
+        coeff = ar / ar_sum
+
+        pose_out = (coeff.unsqueeze(-1) * votes).sum(dim=1)  # [B, K_out, H, W, P_out]
+        pose_out = pose_out.permute(0, 1, 4, 2, 3)  # [B, K_out, P_out, H, W]
+        act_out = ar_sum.squeeze(1).unsqueeze(2)  # [B, K_out, 1, H, W]
+
+        # Keep interleaved packed output: [pose(P_out), act(1)] per capsule type.
+        out = torch.cat([pose_out, act_out], dim=2).reshape(b, self.c_out, h, w)
+        return out
+
+
+class HybridRoute1(nn.Module):
+    """Conv-heavy replacement for SelfRouting with lightweight capsule-aware gating."""
+
+    def __init__(self, K_in: int, P_in: int, K_out: int, P_out: int):
+        super().__init__()
+        self.K_in = int(K_in)
+        self.P_in = int(P_in)
+        self.K_out = int(K_out)
+        self.P_out = int(P_out)
+        self.c_in = self.K_in * (self.P_in + 1)
+        self.c_out = self.K_out * (self.P_out + 1)
+
+        pose_in = self.K_in * self.P_in
+        pose_out = self.K_out * self.P_out
+        vote_groups = math.gcd(self.K_in, self.K_out)
+        vote_groups = max(int(vote_groups), 1)
+        self.vote_proj = Conv(pose_in, pose_out, 1, 1, g=vote_groups)
+        self.gate_proj = nn.Conv2d(self.c_in, self.K_out, kernel_size=1, stride=1, padding=0, bias=True)
+        self.act_proj = nn.Conv2d(self.K_in, self.K_out, kernel_size=1, stride=1, padding=0, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.ndim != 4:
+            raise TypeError(f'HybridRoute1 expects [B,C,H,W], got {tuple(x.shape)}')
+
+        b, c, h, w = x.shape
+        if c != self.c_in:
+            raise ValueError(f'HybridRoute1 expected C={self.c_in}, got C={c}')
+
+        x_caps = x.reshape(b, self.K_in, self.P_in + 1, h, w)
+        pose = x_caps[:, :, :self.P_in].reshape(b, self.K_in * self.P_in, h, w)
+        act = x_caps[:, :, self.P_in].contiguous()
+
+        pose_votes = self.vote_proj(pose).reshape(b, self.K_out, self.P_out, h, w)
+        gate = self.gate_proj(x).sigmoid().unsqueeze(2)
+        pose_out = pose_votes * gate
+
+        act_out = self.act_proj(act).sigmoid().unsqueeze(2)
+        out = torch.cat([pose_out, act_out], dim=2).reshape(b, self.c_out, h, w)
+        return out
+
+
+class CapsRoute(nn.Module):
+    """Capsule routing fusion by direct capsule concatenation.
+
+    Args:
+      K_in: list of input capsule type counts per source.
+      P_in: list of input pose dimensions per source.
+      K_out: target output capsule type count.
+      P_out: target output pose dimension.
+      kernel_size: grouped-conv kernel for ``ConvSelfRouting``.
+
+    Notes:
+      Inputs are concatenated directly (no pre-projection).
+      For direct packed concat, all ``P_in`` must be identical.
+    """
+
+    def __init__(
+        self,
+        K_in: Union[List[int], Tuple[int, ...]],
+        P_in: Union[List[int], Tuple[int, ...]],
+        K_out: int,
+        P_out: int,
+        kernel_size: int = 3,
+        pre_k: int = 3,
+        post_k: int = 3,
+        pre_groups: Optional[int] = None,
+        post_groups: Optional[int] = None,
+    ):
+        super().__init__()
+        self.K_in_list = [int(v) for v in K_in]
+        self.P_in_list = [int(v) for v in P_in]
+        if len(self.K_in_list) < 2 or len(self.K_in_list) != len(self.P_in_list):
+            raise ValueError('CapsRoute expects K_in/P_in lists with same length >= 2.')
+        if min(*self.K_in_list, *self.P_in_list) <= 0:
+            raise ValueError('CapsRoute expects positive K_in/P_in values.')
+
+        # Direct capsule concat requires a shared pose dimension.
+        if len(set(self.P_in_list)) != 1:
+            raise ValueError('CapsRoute direct concat requires all P_in to be identical.')
+
+        self.num_sources = len(self.K_in_list)
+        self.P_cat = int(self.P_in_list[0])
+        self.K_cat = int(sum(self.K_in_list))
+        self.c_cat = self.K_cat * (self.P_cat + 1)
+
+        self.K_out = int(K_out)
+        self.P_out = int(P_out)
+        if min(self.K_out, self.P_out) <= 0:
+            raise ValueError('CapsRoute expects positive K_out/P_out values.')
+        self.c_out = self.K_out * (self.P_out + 1)
+
+        # self.conv_route = ConvSelfRouting(
+        #     K_in=self.K_cat,
+        #     P_in=self.P_cat,
+        #     K_out=self.K_cat,
+        #     P_out=self.P_cat,
+        #     kernel_size=kernel_size,
+        # )
+        # Grouped Conv before routing: C = K_cat * (P_cat + 1), groups = K_cat.
+        self.conv_route = Conv(self.c_cat, self.c_cat, 3, 1, g=self.K_cat)
+        self.route1 = SelfRouting(K_in=self.K_cat, P_in=self.P_cat, K_out=self.K_out, P_out=self.P_out)
+        # Grouped Conv after routing: C = K_out * (P_out + 1), groups = K_out.
+        self.spagg = Conv(self.c_out, self.c_out, 3, 1, g=self.K_out)
+        # self.route2 = SelfRouting(K_in=self.K_out, P_in=self.P_out, K_out=self.K_out, P_out=self.P_out)
+
+    def forward(self, xs: Union[List[torch.Tensor], Tuple[torch.Tensor, ...]]) -> torch.Tensor:
+        if not isinstance(xs, (list, tuple)):
+            raise TypeError(f'CapsRoute expects list/tuple inputs, got {type(xs)}')
+        if len(xs) != self.num_sources:
+            raise ValueError(f'CapsRoute expected {self.num_sources} sources, got {len(xs)}')
+
+        h, w = int(xs[0].shape[-2]), int(xs[0].shape[-1])
+        cat_parts = []
+        for i, x in enumerate(xs):
+            expected_c = self.K_in_list[i] * (self.P_in_list[i] + 1)
+            if int(x.shape[1]) != expected_c:
+                raise ValueError(f'CapsRoute source-{i} expected C={expected_c} from K_in/P_in, got C={int(x.shape[1])}')
+            if int(x.shape[-2]) != h or int(x.shape[-1]) != w:
+                raise ValueError('CapsRoute inputs must share H,W. Use CapsAlign before routing.')
+            cat_parts.append(x)
+
+        x_cat = torch.cat(cat_parts, dim=1)  # [B, K_cat*(P+1), H, W]
+        routed = self.route1(self.conv_route(x_cat))
+        routed = self.spagg(routed)
+        return routed
+
+
+class CapsRoutev2(CapsRoute):
+    """CapsRoute with per-capsule pose refinement and act residual update."""
+
+    def __init__(
+        self,
+        K_in: Union[List[int], Tuple[int, ...]],
+        P_in: Union[List[int], Tuple[int, ...]],
+        K_out: int,
+        P_out: int,
+        kernel_size: int = 3,
+        pre_k: int = 3,
+        post_k: int = 3,
+        pre_groups: Optional[int] = None,
+        post_groups: Optional[int] = None,
+    ):
+        super().__init__(K_in, P_in, K_out, P_out, kernel_size, pre_k, post_k, pre_groups, post_groups)
+        _ = (post_k, post_groups, pre_k, pre_groups)  # kept for YAML/API compatibility
+        self.profile_route = False
+        self._route_profile = {
+            'cat_ms': 0.0,
+            'conv_route_ms': 0.0,
+            'route1_ms': 0.0,
+            'pose_refine_ms': 0.0,
+            'act_from_pose_ms': 0.0,
+            'pack_ms': 0.0,
+            'calls': 0.0,
+        }
+
+        deep_stage = self.K_out >= 64
+        pose_ch = self.K_out * self.P_out
+        # Match YOLO26 neck style:
+        # - shallow/mid stages: C3k2(n=2, c3k=True, attn=False)
+        # - deep stage:       C3k2(n=1, c3k=True, attn=True)
+        pose_e = 0.5 if (self.P_out % 2 == 0) else 1.0
+        self.pose_refine = C3k2(
+            pose_ch,
+            pose_ch,
+            n=1 if deep_stage else 2,
+            c3k=True,
+            e=pose_e,
+            attn=deep_stage,
+            g=self.K_out,
+            shortcut=True,
+        )
+        self.act_from_pose = Conv(pose_ch, self.K_out, 1, 1, g=self.K_out)
+        self.act_alpha = nn.Parameter(torch.tensor(0.1))
+
+    @staticmethod
+    def _sync_profile() -> None:
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+    def _ensure_route_profile_state(self) -> None:
+        if not hasattr(self, "profile_route"):
+            self.profile_route = False
+        if not hasattr(self, "_route_profile"):
+            self._route_profile = {
+                'cat_ms': 0.0,
+                'conv_route_ms': 0.0,
+                'route1_ms': 0.0,
+                'pose_refine_ms': 0.0,
+                'act_from_pose_ms': 0.0,
+                'pack_ms': 0.0,
+                'calls': 0.0,
+            }
+
+    def reset_route_profile(self) -> None:
+        self._ensure_route_profile_state()
+        for k in self._route_profile:
+            self._route_profile[k] = 0.0
+
+    def get_route_profile(self) -> dict:
+        self._ensure_route_profile_state()
+        calls = max(float(self._route_profile.get('calls', 0.0)), 1.0)
+        total = (
+            self._route_profile['cat_ms']
+            + self._route_profile['conv_route_ms']
+            + self._route_profile['route1_ms']
+            + self._route_profile['pose_refine_ms']
+            + self._route_profile['act_from_pose_ms']
+            + self._route_profile['pack_ms']
+        )
+        out = dict(self._route_profile)
+        out['total_ms'] = total
+        out['cat_avg_ms'] = self._route_profile['cat_ms'] / calls
+        out['conv_route_avg_ms'] = self._route_profile['conv_route_ms'] / calls
+        out['route1_avg_ms'] = self._route_profile['route1_ms'] / calls
+        out['pose_refine_avg_ms'] = self._route_profile['pose_refine_ms'] / calls
+        out['act_from_pose_avg_ms'] = self._route_profile['act_from_pose_ms'] / calls
+        out['pack_avg_ms'] = self._route_profile['pack_ms'] / calls
+        out['total_avg_ms'] = total / calls
+        return out
+
+    def forward(self, xs: Union[List[torch.Tensor], Tuple[torch.Tensor, ...]]) -> torch.Tensor:
+        if not isinstance(xs, (list, tuple)):
+            raise TypeError(f'CapsRoutev2 expects list/tuple inputs, got {type(xs)}')
+        if len(xs) != self.num_sources:
+            raise ValueError(f'CapsRoutev2 expected {self.num_sources} sources, got {len(xs)}')
+
+        h, w = int(xs[0].shape[-2]), int(xs[0].shape[-1])
+        cat_parts = []
+        for i, x in enumerate(xs):
+            expected_c = self.K_in_list[i] * (self.P_in_list[i] + 1)
+            if int(x.shape[1]) != expected_c:
+                raise ValueError(f'CapsRoutev2 source-{i} expected C={expected_c}, got C={int(x.shape[1])}')
+            if int(x.shape[-2]) != h or int(x.shape[-1]) != w:
+                raise ValueError('CapsRoutev2 inputs must share H,W. Use CapsAlign before routing.')
+            cat_parts.append(x)
+
+        self._ensure_route_profile_state()
+        if getattr(self, "profile_route", False):
+            self._route_profile['calls'] += 1.0
+            self._sync_profile()
+            t0 = time.perf_counter()
+            x_cat = torch.cat(cat_parts, dim=1)  # [B, K_cat*(P+1), H, W]
+            self._sync_profile()
+            self._route_profile['cat_ms'] += (time.perf_counter() - t0) * 1000.0
+
+            t0 = time.perf_counter()
+            conv_out = self.conv_route(x_cat)
+            self._sync_profile()
+            self._route_profile['conv_route_ms'] += (time.perf_counter() - t0) * 1000.0
+
+            t0 = time.perf_counter()
+            routed = self.route1(conv_out)  # [B, K_out*(P_out+1), H, W]
+            self._sync_profile()
+            self._route_profile['route1_ms'] += (time.perf_counter() - t0) * 1000.0
+        else:
+            x_cat = torch.cat(cat_parts, dim=1)  # [B, K_cat*(P+1), H, W]
+            routed = self.route1(self.conv_route(x_cat))  # [B, K_out*(P_out+1), H, W]
+
+        b, _, _, _ = routed.shape
+        # Packed layout by type: [pose(P), act(1)] repeated K times.
+        caps = routed.reshape(b, self.K_out, self.P_out + 1, h, w)
+        pose = caps[:, :, :self.P_out].contiguous()  # [B, K_out, P_out, H, W]
+        act = caps[:, :, self.P_out].contiguous()  # [B, K_out, H, W]
+
+        # Grouped pose refinement across type blocks (equivalent to per-type grouped processing).
+        pose_flat = pose.reshape(b, self.K_out * self.P_out, h, w)
+        if getattr(self, "profile_route", False):
+            t0 = time.perf_counter()
+            pose_flat = self.pose_refine(pose_flat)
+            self._sync_profile()
+            self._route_profile['pose_refine_ms'] += (time.perf_counter() - t0) * 1000.0
+
+            t0 = time.perf_counter()
+            act_delta = self.act_from_pose(pose_flat)
+            act_final = act + act_delta
+            self._sync_profile()
+            self._route_profile['act_from_pose_ms'] += (time.perf_counter() - t0) * 1000.0
+        else:
+            pose_flat = self.pose_refine(pose_flat)
+            act_delta = self.act_from_pose(pose_flat)
+            act_final = act + act_delta
+
+        if getattr(self, "profile_route", False):
+            t0 = time.perf_counter()
+            pose_pack = pose_flat.reshape(b, self.K_out, self.P_out, h, w)
+            out = torch.cat([pose_pack, act_final.unsqueeze(2)], dim=2).reshape(b, self.c_out, h, w)
+            self._sync_profile()
+            self._route_profile['pack_ms'] += (time.perf_counter() - t0) * 1000.0
+        else:
+            pose_pack = pose_flat.reshape(b, self.K_out, self.P_out, h, w)
+            out = torch.cat([pose_pack, act_final.unsqueeze(2)], dim=2).reshape(b, self.c_out, h, w)
+        return out
+
+
+# -------------------------
+# 4) CapsDecode
+# -------------------------
+
+class CapsRoutev3(CapsRoute):
+    """CapsRoute with DS-style lightweight pose refinement and act residual update."""
+
+    def __init__(
+        self,
+        K_in: Union[List[int], Tuple[int, ...]],
+        P_in: Union[List[int], Tuple[int, ...]],
+        K_out: int,
+        P_out: int,
+        kernel_size: int = 3,
+        pre_k: int = 3,
+        post_k: int = 3,
+        pre_groups: Optional[int] = None,
+        post_groups: Optional[int] = None,
+    ):
+        super().__init__(K_in, P_in, K_out, P_out, kernel_size, pre_k, post_k, pre_groups, post_groups)
+        _ = (post_k, post_groups, pre_k, pre_groups)
+        self.profile_route = False
+        self._route_profile = {
+            'cat_ms': 0.0,
+            'conv_route_ms': 0.0,
+            'route1_ms': 0.0,
+            'pose_refine_ms': 0.0,
+            'act_from_pose_ms': 0.0,
+            'pack_ms': 0.0,
+            'calls': 0.0,
+        }
+
+        pose_ch = self.K_out * self.P_out
+        # Keep refinement fully type-grouped to preserve capsule semantics:
+        # each capsule type only mixes its own pose channels.
+        self.pose_refine = nn.Sequential(
+            Conv(pose_ch, pose_ch, 1, 1, g=self.K_out),
+            Conv(pose_ch, pose_ch, 3, 1, g=self.K_out),
+            Conv(pose_ch, pose_ch, 1, 1, g=self.K_out),
+        )
+        self.act_from_pose = Conv(pose_ch, self.K_out, 1, 1, g=self.K_out)
+        self.act_alpha = nn.Parameter(torch.tensor(0.1))
+
+    @staticmethod
+    def _sync_profile() -> None:
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+
+    def _ensure_route_profile_state(self) -> None:
+        if not hasattr(self, "profile_route"):
+            self.profile_route = False
+        if not hasattr(self, "_route_profile"):
+            self._route_profile = {
+                'cat_ms': 0.0,
+                'conv_route_ms': 0.0,
+                'route1_ms': 0.0,
+                'pose_refine_ms': 0.0,
+                'act_from_pose_ms': 0.0,
+                'pack_ms': 0.0,
+                'calls': 0.0,
+            }
+
+    def reset_route_profile(self) -> None:
+        self._ensure_route_profile_state()
+        for k in self._route_profile:
+            self._route_profile[k] = 0.0
+
+    def get_route_profile(self) -> dict:
+        self._ensure_route_profile_state()
+        calls = max(float(self._route_profile.get('calls', 0.0)), 1.0)
+        total = (
+            self._route_profile['cat_ms']
+            + self._route_profile['conv_route_ms']
+            + self._route_profile['route1_ms']
+            + self._route_profile['pose_refine_ms']
+            + self._route_profile['act_from_pose_ms']
+            + self._route_profile['pack_ms']
+        )
+        out = dict(self._route_profile)
+        out['total_ms'] = total
+        out['cat_avg_ms'] = self._route_profile['cat_ms'] / calls
+        out['conv_route_avg_ms'] = self._route_profile['conv_route_ms'] / calls
+        out['route1_avg_ms'] = self._route_profile['route1_ms'] / calls
+        out['pose_refine_avg_ms'] = self._route_profile['pose_refine_ms'] / calls
+        out['act_from_pose_avg_ms'] = self._route_profile['act_from_pose_ms'] / calls
+        out['pack_avg_ms'] = self._route_profile['pack_ms'] / calls
+        out['total_avg_ms'] = total / calls
+        return out
+
+    def forward(self, xs: Union[List[torch.Tensor], Tuple[torch.Tensor, ...]]) -> torch.Tensor:
+        if not isinstance(xs, (list, tuple)):
+            raise TypeError(f'CapsRoutev3 expects list/tuple inputs, got {type(xs)}')
+        if len(xs) != self.num_sources:
+            raise ValueError(f'CapsRoutev3 expected {self.num_sources} sources, got {len(xs)}')
+
+        h, w = int(xs[0].shape[-2]), int(xs[0].shape[-1])
+        cat_parts = []
+        for i, x in enumerate(xs):
+            expected_c = self.K_in_list[i] * (self.P_in_list[i] + 1)
+            if int(x.shape[1]) != expected_c:
+                raise ValueError(f'CapsRoutev3 source-{i} expected C={expected_c}, got C={int(x.shape[1])}')
+            if int(x.shape[-2]) != h or int(x.shape[-1]) != w:
+                raise ValueError('CapsRoutev3 inputs must share H,W. Use CapsAlign before routing.')
+            cat_parts.append(x)
+
+        self._ensure_route_profile_state()
+        if getattr(self, "profile_route", False):
+            self._route_profile['calls'] += 1.0
+            self._sync_profile()
+            t0 = time.perf_counter()
+            x_cat = torch.cat(cat_parts, dim=1)
+            self._sync_profile()
+            self._route_profile['cat_ms'] += (time.perf_counter() - t0) * 1000.0
+
+            t0 = time.perf_counter()
+            conv_out = self.conv_route(x_cat)
+            self._sync_profile()
+            self._route_profile['conv_route_ms'] += (time.perf_counter() - t0) * 1000.0
+
+            t0 = time.perf_counter()
+            routed = self.route1(conv_out)
+            self._sync_profile()
+            self._route_profile['route1_ms'] += (time.perf_counter() - t0) * 1000.0
+        else:
+            x_cat = torch.cat(cat_parts, dim=1)
+            routed = self.route1(self.conv_route(x_cat))
+
+        b, _, _, _ = routed.shape
+        caps = routed.reshape(b, self.K_out, self.P_out + 1, h, w)
+        pose = caps[:, :, :self.P_out].contiguous()
+        act = caps[:, :, self.P_out].contiguous()
+
+        pose_flat = pose.reshape(b, self.K_out * self.P_out, h, w)
+        if getattr(self, "profile_route", False):
+            t0 = time.perf_counter()
+            pose_flat = pose_flat + self.pose_refine(pose_flat)
+            self._sync_profile()
+            self._route_profile['pose_refine_ms'] += (time.perf_counter() - t0) * 1000.0
+
+            t0 = time.perf_counter()
+            act_delta = self.act_from_pose(pose_flat)
+            act_final = act + act_delta
+            self._sync_profile()
+            self._route_profile['act_from_pose_ms'] += (time.perf_counter() - t0) * 1000.0
+        else:
+            pose_flat = pose_flat + self.pose_refine(pose_flat)
+            act_delta = self.act_from_pose(pose_flat)
+            act_final = act + act_delta
+
+        if getattr(self, "profile_route", False):
+            t0 = time.perf_counter()
+            pose_pack = pose_flat.reshape(b, self.K_out, self.P_out, h, w)
+            out = torch.cat([pose_pack, act_final.unsqueeze(2)], dim=2).reshape(b, self.c_out, h, w)
+            self._sync_profile()
+            self._route_profile['pack_ms'] += (time.perf_counter() - t0) * 1000.0
+        else:
+            pose_pack = pose_flat.reshape(b, self.K_out, self.P_out, h, w)
+        out = torch.cat([pose_pack, act_final.unsqueeze(2)], dim=2).reshape(b, self.c_out, h, w)
+        return out
+
+
+class CapsRoutev4(CapsRoutev2):
+    """CapsRoutev2 with conv-heavy HybridRoute1 to reduce routing overhead."""
+
+    def __init__(
+        self,
+        K_in: Union[List[int], Tuple[int, ...]],
+        P_in: Union[List[int], Tuple[int, ...]],
+        K_out: int,
+        P_out: int,
+        kernel_size: int = 3,
+        pre_k: int = 3,
+        post_k: int = 3,
+        pre_groups: Optional[int] = None,
+        post_groups: Optional[int] = None,
+    ):
+        super().__init__(K_in, P_in, K_out, P_out, kernel_size, pre_k, post_k, pre_groups, post_groups)
+        self.route1 = HybridRoute1(K_in=self.K_cat, P_in=self.P_cat, K_out=self.K_out, P_out=self.P_out)
+
+
+class CapsDecode(nn.Module):
+    """
+    Decode routed capsule features to standard feature map for Detect.
+
+    Input:  y [B, C_in, H, W]  (often concat of weighted sources, so C_in = S*(K*D))
+    Output: f [B, C_out, H, W]
+
+    Args:
+      c2: output channels (e.g., 256/512/1024)
+    """
+
+    def __init__(self, c1: int, c2: int):
+        super().__init__()
+        self.conv = nn.Conv2d(c1, c2, kernel_size=1, stride=1, padding=0, bias=False)
+        self.bn = nn.BatchNorm2d(c2)
+        self.act = nn.SiLU(inplace=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.act(self.bn(self.conv(x)))
+
+
+# -------------------------
+# 5) CapsuleTap
+# -------------------------
+
+class CapsuleTap(nn.Module):
+    """
+    Pass-through hook to cache feature maps for explainability/aux loss.
+
+    MUST NOT change tensor shape. Returns x unchanged.
+
+    Args:
+      tag: string identifier ("F3"/"F4"/"F5")
+      K,D: capsule hyperparams (metadata only)
+      cache_enabled: if True, cache during training (disabled in tracing/scripting)
+    """
+
+    def __init__(self, tag: str = "F", K: int = 4, D: int = 16, cache_enabled: bool = True):
+        super().__init__()
+        self.tag = str(tag)
+        self.K = int(K)
+        self.D = int(D)
+        self.cache_enabled = bool(cache_enabled)
+        self.last_x: Optional[torch.Tensor] = None
+
+    def clear_cache(self) -> None:
+        self.last_x = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if (
+            self.cache_enabled
+            and self.training
+            and (not torch.jit.is_scripting())
+            and (not torch.jit.is_tracing())
+        ):
+            self.last_x = x.detach()
+        return x

predict.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from ultralytics import YOLO
+
+from models import register_ultralytics_modules
+
+
+ROOT = Path(__file__).resolve().parent
+DEFAULT_WEIGHTS = ROOT / "weights" / "symbolic_capsule_network_segmentation.pt"
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Run Symbolic Capsule Network segmentation inference.")
+    parser.add_argument("source", help="Image, directory, video, or glob pattern.")
+    parser.add_argument("--weights", default=str(DEFAULT_WEIGHTS), help="Checkpoint path.")
+    parser.add_argument("--imgsz", type=int, default=640)
+    parser.add_argument("--conf", type=float, default=0.25)
+    parser.add_argument("--device", default="")
+    parser.add_argument("--save", action="store_true", default=True)
+    parser.add_argument("--show", action="store_true")
+    return parser
+
+
+def main() -> None:
+    args = build_parser().parse_args()
+    weights = Path(args.weights).expanduser().resolve()
+    if not weights.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {weights}")
+
+    register_ultralytics_modules()
+    model = YOLO(str(weights))
+
+    predict_kwargs = {k: v for k, v in vars(args).items() if k != "weights"}
+    model.predict(**predict_kwargs)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -0,0 +1,27 @@
+[project]
+name = "csnet"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "fiftyone>=1.11.1",
+    "langchain>=1.2.15",
+    "langchain-openai>=1.1.12",
+    "mobileclip",
+    "numpy>=2.4.4",
+    "torch>=2.7.0",
+    "torchvision>=0.21.0",
+    "transformers>=5.5.0",
+    "ultralytics>=8.4.9",
+]
+
+[tool.uv.sources]
+torch = { index = "pytorch-cu128" }
+torchvision = { index = "pytorch-cu128" }
+mobileclip = { git = "https://github.com/ultralytics/mobileclip.git" }
+
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true