export_onnx.py

"""
ONNX export scripts for WeDetect.

Two separate exports:
  1. Image Encoder + Neck + Head  →  wedetect_image_head.onnx
  2. Text Encoder                  →  wedetect_text_encoder.onnx

Usage:
  python export_onnx.py --config config/wedetect_base.py --checkpoint checkpoints/wedetect_base.pth

The text encoder is exported separately so text embeddings can be pre-computed
offline for any category set.  At inference time the image+head ONNX model takes
the pre-computed text features together with the image and produces detections.
"""

import argparse
import os.path as osp
import warnings
from io import BytesIO

import onnx
import torch
import torch.nn as nn

from mmdet.utils import register_all_modules
register_all_modules()
from mmengine.config import Config
from mmdet.apis import init_detector

warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning)
warnings.filterwarnings(action='ignore', category=torch.jit.ScriptWarning)
warnings.filterwarnings(action='ignore', category=UserWarning)
warnings.filterwarnings(action='ignore', category=FutureWarning)


# ---------------------------------------------------------------------------
# Wrapper: Image Encoder + Neck + Head
# ---------------------------------------------------------------------------

class ImageHeadWrapper(nn.Module):
    """Pack ConvNeXt → Neck → Head into a single traceable module.

    Inputs:
        image          [B, 3, 640, 640]  float32
        text_features  [B, num_classes, 768]  float32

    Outputs (three scales, strides 8 / 16 / 32):
        cls_scores_s8   [B, num_classes, 80, 80]
        cls_scores_s16  [B, num_classes, 40, 40]
        cls_scores_s32  [B, num_classes, 20, 20]
        bbox_preds_s8   [B, 4, 80, 80]
        bbox_preds_s16  [B, 4, 40, 40]
        bbox_preds_s32  [B, 4, 20, 20]
    """

    def __init__(self, model):
        super().__init__()
        self.backbone = model.backbone
        self.neck = model.neck
        self.head_module = model.bbox_head.head_module
    def forward(self, image, text_features):
        # ConvNeXt  (forward_image skips the text encoder)
        img_feats = self.backbone.forward_image(image)

        # BiFPN neck
        if self.neck is not None:
            img_feats = self.neck(img_feats)

        # Head: returns ((cls_L0, cls_L1, cls_L2), (bbox_L0, bbox_L1, bbox_L2))
        cls_scores, bbox_preds = self.head_module(img_feats, text_features)

        # Flatten to a single tuple so ONNX output names align correctly
        return (
            cls_scores[0], bbox_preds[0],    # stride  8
            cls_scores[1], bbox_preds[1],    # stride 16
            cls_scores[2], bbox_preds[2],    # stride 32
        )


# ---------------------------------------------------------------------------
# Wrapper: Text Encoder  (XLM-RoBERTa without tokenizer)
# ---------------------------------------------------------------------------

class TextEncoderWrapper(nn.Module):
    """Wrap the XLM-RoBERTa backbone so it can be exported without a tokenizer.

    Inputs:
        input_ids       [num_texts, max_seq_len]  int64   (fixed seq_len)
        attention_mask  [num_texts, max_seq_len]  int64

    Output:
        text_features   [1, num_texts, 768]  float32  (L2-normalised)

    """

    def __init__(self, text_model):
        super().__init__()
        self.model = text_model.model      # XLMRobertaModel
        self.head = text_model.head        # Linear(768 → 768)

    def forward(self, input_ids, attention_mask):
        # Compute position_ids in float32 to avoid INT32 CumSum/Mul/Add ops
        # that the NPU backend does not support.
        # XLMRoberta uses padding_idx=1, so mask = (input_ids != 1).
        mask = input_ids.ne(1).float()                         # [N, L] float32
        position_ids = torch.cumsum(mask, dim=1) * mask        # cumsum+mul in float32
        position_ids = (position_ids + 1.0).long()             # +padding_idx in float, then cast

        out = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
        )
        cls_embed = out["last_hidden_state"][:, 0, :]   # [N, 768]
        txt_feats = self.head(cls_embed)
        # L2-normalise + add batch dim → (1, N, 768)
        txt_feats = nn.functional.normalize(txt_feats, dim=-1)
        return txt_feats.unsqueeze(0)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _onnx_simplify(onnx_model, output_path):
    """Save ONNX, then try to simplify and overwrite.

    Saving first guarantees a usable file even if onnxsim crashes.
    """
    # Save unsimplified first — always produces a valid file on disk
    onnx.save(onnx_model, output_path)
    print(f"Saved (unsimplified): {output_path}")

    # Then attempt simplify
    try:
        import onnxsim
        simplified, check = onnxsim.simplify(onnx_model)
        if check:
            onnx.save(simplified, output_path)
            print("ONNX simplify: passed, overwritten")
        else:
            print("ONNX simplify: check failed, kept unsimplified version")
    except Exception as e:
        print(f"ONNX simplify skipped: {e}, kept unsimplified version")


# ---------------------------------------------------------------------------
# Export entry-points
# ---------------------------------------------------------------------------

def export_image_encoder(model, output_path, num_classes=80, image_size=640):
    """Export Image Encoder + Neck + Head to ONNX."""
    wrapper = ImageHeadWrapper(model)
    wrapper.eval()

    device = next(model.parameters()).device

    # Dummy inputs
    image = torch.randn(1, 3, image_size, image_size, device=device)
    text_feats = torch.randn(1, num_classes, 768, device=device)

    # Dry-run to verify shapes
    with torch.no_grad():
        outputs = wrapper(image, text_feats)
    print("Image + Head output shapes:")
    for i, o in enumerate(outputs):
        print(f"  output_{i}: {list(o.shape)}")

    output_names = [
        "cls_scores_s8", "bbox_preds_s8",
        "cls_scores_s16", "bbox_preds_s16",
        "cls_scores_s32", "bbox_preds_s32",
    ]

    # Export to buffer → check → simplify → save
    with BytesIO() as f:
        torch.onnx.export(
            wrapper,
            (image, text_feats),
            f,
            input_names=["images", "text_features"],
            output_names=output_names,
            opset_version=14,
            do_constant_folding=True,
        )
        f.seek(0)
        onnx_model = onnx.load(f)
        onnx.checker.check_model(onnx_model)

    _onnx_simplify(onnx_model, output_path)
    print(f"Exported: {output_path}")


def export_text_encoder(model, output_path, num_texts=4, max_seq_len=32):
    """Export Text Encoder (XLM-RoBERTa + projection) to ONNX.

    The seq_len dimension is FIXED at export time.  At inference, tokenized
    inputs must be padded to exactly ``max_seq_len`` tokens (the tokenizer
    does this automatically when ``padding='max_length'`` is set).

    Parameters
    ----------
    max_seq_len : int
        Fixed token length.  32 is sufficient for typical Chinese class names.
    """
    text_model = model.backbone.text_model
    wrapper = TextEncoderWrapper(text_model)
    wrapper.eval()

    device = next(text_model.parameters()).device

    # Dummy inputs — fixed seq_len, dynamic num_texts
    input_ids = torch.ones(num_texts, max_seq_len, dtype=torch.int64, device=device)
    attention_mask = torch.ones(num_texts, max_seq_len, dtype=torch.int64, device=device)

    with torch.no_grad():
        out = wrapper(input_ids, attention_mask)
    print(f"Text Encoder output shape: {list(out.shape)}")

    # Export to buffer → check → simplify → save
    with BytesIO() as f:
        torch.onnx.export(
            wrapper,
            (input_ids, attention_mask),
            f,
            input_names=["input_ids", "attention_mask"],
            output_names=["text_features"],
            opset_version=17,
            do_constant_folding=True,
        )
        f.seek(0)
        onnx_model = onnx.load(f)
        onnx.checker.check_model(onnx_model)

    _onnx_simplify(onnx_model, output_path)
    print(f"Exported: {output_path}")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args():
    parser = argparse.ArgumentParser(description="Export WeDetect to ONNX")
    parser.add_argument("--config", required=True, help="Config file path")
    parser.add_argument("--checkpoint", required=True, help="Checkpoint file path")
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--num-classes", type=int, default=4,
                        help="Number of classes for dummy text features (default 4).")
    parser.add_argument("--image-size", type=int, default=640)
    #此处演示导出检测4个类别的模型
    parser.add_argument("--num-texts", type=int, default=4,
                        help="Number of dummy texts for text-encoder trace.")
    #可以根据实际检测类别prompt计算选取最大seq_len
    parser.add_argument("--max-seq-len", type=int, default=32,
                        help="Fixed token length for text-encoder ONNX. "
                             "Inference inputs must be padded to this length.")
    parser.add_argument("--output-dir", default=".",
                        help="Directory for the exported .onnx files.")
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()

    cfg = Config.fromfile(args.config)
    cfg.work_dir = osp.join("./work_dirs",
                            osp.splitext(osp.basename(args.config))[0])

    print(f"Loading model from {args.checkpoint} ...")
    model = init_detector(cfg, checkpoint=args.checkpoint, device=args.device, palette=['red'])
    model.eval()

    # ---- Image + Head ----
    export_image_encoder(
        model,
        osp.join(args.output_dir, "wedetect_image_encoder.onnx"),
        num_classes=args.num_classes,
        image_size=args.image_size,
    )

    # ---- Text Encoder ----
    export_text_encoder(
        model,
        osp.join(args.output_dir, "wedetect_text_encoder.onnx"),
        num_texts=args.num_texts,
        max_seq_len=args.max_seq_len,
    )
    print("Done.")