first commit

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.axmodel filter=lfs diff=lfs merge=lfs -text
+*.tar.gz filter=lfs diff=lfs merge=lfs -text
+*.json filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.ttc filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,85 @@
+# WeDetect demo for AX
+
+## The original repo
+
+[WeDetect](https://github.com/WeChatCV/WeDetect/tree/main)
+
+## 背景
+开放词汇检测旨在利用文本描述来检测任意的物体，Wedetect不利用跨模态交互的方案把识别任务类比成一种检索任务，即在一个统一的特征空间中匹配区域特征和文本特征。本项目用于指导开发者完成以下内容：
+
+- 导出 class num = 4 的 WeDetect ONNX 模型；
+- 生成 AXERA NPU 模型转换工具 Pulsar2 编译依赖的 text 量化校准数据集；
+- 完成ONNX模型基于Pulsa2工具链的编译及在AX650N上的部署。
+
+## 模型导出
+
+项目使用模型为wedetect_base,模型可在[huggingface](https://huggingface.co/fushh7/WeDetect/tree/main)下载。
+生成适合用于 AXera NPU 工具链 Pulsar2 模型转换的 ONNX 模型:
+
+- 下载 `wedetect_base.pth`放在checkpoints目录下
+- 使用 export_onnx.py分别导出图像编码模型和文本编码模型：
+
+```
+python export_onnx.py --config config/wedetect_base.py --checkpoint checkpoints/wedetect_base.pth
+```
+
+- 生成 Pulsar2 编译模型时所依赖的量化校准数据 `class_embedding_4cls.tar.gz`、`input_ids.tar.gz`、 `attention_mask.tar.gz`，另准备图片数据如`coco100.tar.gz`作为图像编码器量化对分数据
+
+```
+python generate_class_embedding.py --wedetect_checkpoint checkpoints/wedetect_base.pth --classname_file ./coco_zh_class_texts.json --calib-dir ./quant
+```
+
+## 模型编译
+
+- Pulsar2 安装及使用请参考相关文档
+  - [在线文档](https://pulsar2-docs.readthedocs.io/zh-cn/latest/index.html)
+
+- 编译命令
+```
+# image encoder
+pulsar2 build --config quant/image_encoder_4cls.json
+
+# text encoder
+pulsar2 build --config quant/text_encoder_4cls.json
+```
+- 模型性能
+
+| Models                                  | Platforms | latency       | CMM size(MB)  |
+| --------------------------------------- | --------- | ------------- | ------------- |
+| wedetect_image_encoder_npu3_u16.axmodel | AX650     | 100.3ms       | 152.2         |
+| wedetect_text_encoder_npu3_u16.axmodel  | AX650     | 7.2ms         | 457.8         |
+
+
+## Python demo
+
+### Onnx demo
+```
+python3 onnx_infer.py
+```
+
+### Axmodel demo
+
+需基于[PyAXEngine](https://github.com/AXERA-TECH/pyaxengine)在AX650N上进行部署
+板端执行命令：
+```
+root@ax650:~/WeDetect# python3 axmodel_infer.py
+[INFO] Available providers:  ['AxEngineExecutionProvider', 'AXCLRTExecutionProvider']
+Classes: ['鞋', '床', '人', '衣架']
+[INFO] Using provider: AxEngineExecutionProvider
+[INFO] Chip type: ChipType.MC50
+[INFO] VNPU type: VNPUType.DISABLED
+[INFO] Engine version: 2.12.0s
+[INFO] Model type: 2 (triple core)
+[INFO] Compiler version: 6.0 6965315a
+[INFO] Using provider: AxEngineExecutionProvider
+[INFO] Model type: 2 (triple core)
+[INFO] Compiler version: 6.0 62ad4ff7
+Detections: 3
+  床 0.863  (336, 427, 837, 679)
+  鞋 0.430  (316, 629, 345, 688)
+  鞋 0.397  (268, 628, 331, 679)
+Saved: axmodel_res.jpg
+
+```
+
+![](axmodel_res.jpg)

axmodel/wedetect_image_encoder_npu3_u16.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fcbce3f947d827ead3429e7b8f6b52b28ec434e9ca1482328d4e1d605a61b81b
+size 139985400

axmodel/wedetect_text_encoder_npu3_u16.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8058303286a3612e6b2e98c602bafca470236a11311904d63fd13e4175198b7b
+size 479847260

axmodel_infer.py

@@ -0,0 +1,279 @@
+"""
+AXMODEL runtime inference for WeDetect.
+
+Usage:
+  python axmodel_infer.py --image assets/demo.jpeg --text "鞋,床" --threshold 0.3
+
+Dependencies: axengine, numpy, PIL, transformers (tokenizer only)
+"""
+
+import argparse
+import os
+import sys
+
+import numpy as np
+import axengine as axe
+from PIL import Image, ImageDraw, ImageFont
+from transformers import AutoTokenizer
+
+# ---------------------------------------------------------------------------
+# Image preprocessing (mirrors WeDetectKeepRatioResize + LetterResize)
+# ---------------------------------------------------------------------------
+
+def letterbox(image: Image.Image, new_shape=(640, 640), color=(114, 114, 114)):
+    """Resize keeping aspect ratio and pad to ``new_shape``."""
+    ow, oh = image.size
+    r = min(new_shape[0] / ow, new_shape[1] / oh)
+    new_w, new_h = int(round(ow * r)), int(round(oh * r))
+    image = image.resize((new_w, new_h), Image.BILINEAR)
+
+    # Paste onto a blank canvas
+    canvas = Image.new("RGB", new_shape, color)
+    left = (new_shape[0] - new_w) // 2
+    top = (new_shape[1] - new_h) // 2
+    canvas.paste(image, (left, top))
+
+    return canvas, r, (left, top)
+
+
+def preprocess_image(image_path: str, image_size=640):
+    """Load, letterbox, and normalise an image.
+    """
+    img = Image.open(image_path).convert("RGB")
+    img, ratio, (pad_left, pad_top) = letterbox(img, (image_size, image_size))
+
+    # arr = np.array(img, dtype=np.float32) / 255.0       # [0, 1], HWC
+    arr = np.array(img, dtype=np.uint8)
+    tensor = arr[None]                # NHWC
+    return tensor, ratio, (pad_left, pad_top)
+
+
+# ---------------------------------------------------------------------------
+# Tokenization
+# ---------------------------------------------------------------------------
+
+def tokenize(texts, tokenizer, max_seq_len=32):
+    """Tokenize category names into fixed-length tensors.
+
+    Returns
+    -------
+    input_ids      : np.ndarray  shape (N, max_seq_len)  int64
+    attention_mask : np.ndarray  shape (N, max_seq_len)  int64
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+    tokens = tokenizer(texts, padding="max_length", max_length=max_seq_len,
+                       return_tensors="np")
+    return tokens["input_ids"], tokens["attention_mask"]
+
+
+# ---------------------------------------------------------------------------
+# BBox decoding
+# ---------------------------------------------------------------------------
+
+def _grid_points(h, w, stride):
+    """Generate grid anchor points for one feature-map scale."""
+    x = (np.arange(w) + 0.5) * stride
+    y = (np.arange(h) + 0.5) * stride
+    yy, xx = np.meshgrid(y, x, indexing="ij")
+    return np.stack([xx, yy], axis=-1).reshape(-1, 2)  # (H*W, 2)
+
+
+def _decode_bboxes(cls_scores, bbox_preds, h, w, stride, score_thr):
+    """Decode one scale: sigmoid → filter → distance2bbox."""
+    # cls_scores:  (1, C, H, W)   bbox_preds: (1, 4, H, W)
+    C = cls_scores.shape[1]
+    scores = 1.0 / (1.0 + np.exp(-cls_scores[0]))  # sigmoid, (C, H, W)
+    scores = scores.reshape(C, -1).transpose(1, 0)  # (H*W, C)
+
+    bbox = bbox_preds[0].reshape(4, -1).transpose(1, 0)     # (H*W, 4)
+    points = _grid_points(h, w, stride)                     # (H*W, 2)
+
+    bbox = bbox * stride  # scale distances to pixels
+    x1 = points[:, 0] - bbox[:, 0]
+    y1 = points[:, 1] - bbox[:, 1]
+    x2 = points[:, 0] + bbox[:, 2]
+    y2 = points[:, 1] + bbox[:, 3]
+    boxes = np.stack([x1, y1, x2, y2], axis=-1)             # (H*W, 4)
+
+    # Filter by score
+    max_scores = scores.max(axis=1)
+    keep = max_scores > score_thr
+    return boxes[keep], scores[keep], max_scores[keep]
+
+
+# ---------------------------------------------------------------------------
+# NMS
+# ---------------------------------------------------------------------------
+
+def nms(boxes, scores, iou_threshold=0.7, max_dets=300):
+    """Simple numpy NMS."""
+    order = scores.argsort()[::-1]
+    keep = []
+    while order.size > 0 and len(keep) < max_dets:
+        i = order[0]
+        keep.append(i)
+        if order.size == 1:
+            break
+        ious = _box_iou(boxes[i:i + 1], boxes[order[1:]])[0]
+        order = order[1:][ious < iou_threshold]
+    return np.array(keep, dtype=np.int64) if keep else np.array([], dtype=np.int64)
+
+
+def _box_iou(boxes1, boxes2):
+    """Pairwise IoU between two sets of xyxy boxes."""
+    area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
+    area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])
+
+    inter_x1 = np.maximum(boxes1[:, None, 0], boxes2[None, :, 0])
+    inter_y1 = np.maximum(boxes1[:, None, 1], boxes2[None, :, 1])
+    inter_x2 = np.minimum(boxes1[:, None, 2], boxes2[None, :, 2])
+    inter_y2 = np.minimum(boxes1[:, None, 3], boxes2[None, :, 3])
+    inter_w = np.maximum(0, inter_x2 - inter_x1)
+    inter_h = np.maximum(0, inter_y2 - inter_y1)
+    inter = inter_w * inter_h
+    return inter / (area1[:, None] + area2[None, :] - inter + 1e-7)
+
+# ---------------------------------------------------------------------------
+# Visualisation
+# ---------------------------------------------------------------------------
+
+def draw_boxes(image_path, boxes, labels, scores, output_path, font_path=None):
+    """Draw bounding boxes with Chinese-capable font and save."""
+
+    img = Image.open(image_path).convert("RGB")
+    draw = ImageDraw.Draw(img)
+
+    font = ImageFont.truetype(font_path, size=18)
+    font_small = ImageFont.truetype(font_path, size=14)
+
+    colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0),
+              (255, 0, 255), (0, 255, 255), (128, 0, 128), (255, 165, 0)]
+
+    for i, (box, label, score) in enumerate(zip(boxes, labels, scores)):
+        x1, y1, x2, y2 = box.astype(int)
+        c = colors[i % len(colors)]
+        draw.rectangle([x1, y1, x2, y2], outline=c, width=2)
+
+        # Use textbbox for accurate label background sizing (PIL >= 8.0)
+        text = f"{label} {score:.2f}"
+        if hasattr(draw, "textbbox"):
+            bbox = draw.textbbox((x1 + 2, y1 + 2), text, font=font)
+            draw.rectangle([x1, y1, bbox[2] + 4, bbox[3] + 2], fill=c)
+            draw.text((x1 + 2, y1 + 2), text, fill="white", font=font)
+        else:  # older PIL fallback
+            draw.rectangle([x1, y1, x1 + len(text) * 10, y1 + 20], fill=c)
+            draw.text((x1 + 2, y1 + 2), text, fill="white", font=font)
+
+    img.save(output_path)
+    print(f"Saved: {output_path}")
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def parse_args():
+    p = argparse.ArgumentParser(description="WeDetect AXMODEL inference")
+    p.add_argument("--image", default='assets/demo.jpeg', help="Input image path")
+    p.add_argument("--text", default="鞋,床,人,衣架",
+                   help="Chinese class names separated by comma, e.g. '鞋,床,人,衣架'")
+    p.add_argument("--image_model", default="wedetect_image_encoder_npu3_u16.axmodel")
+    p.add_argument("--text_model", default="wedetect_text_encoder_npu3_u16.axmodel")
+    p.add_argument("--tokenizer-dir", default="./xlm-roberta-base/")
+    p.add_argument("--max-seq-len", type=int, default=32)
+    p.add_argument("--threshold", type=float, default=0.3)
+    p.add_argument("--nms-threshold", type=float, default=0.7)
+    p.add_argument("--output", default="axmodel_res.jpg")
+    p.add_argument("--image-size", type=int, default=640)
+    p.add_argument("--font", default='wqy-microhei.ttc',
+                   help="Path to a .ttf/.ttc font file for Chinese label rendering. "
+                        "Auto-detected if not specified.")
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    # ---- Parse texts ----
+    texts = [t.strip() for t in args.text.split(",")]
+    print(f"Classes: {texts}")
+
+    # ---- Load AXMODEL sessions ----
+    img_sess = axe.InferenceSession(args.image_model,providers=["AxEngineExecutionProvider"])
+    txt_sess = axe.InferenceSession(args.text_model,providers=["AxEngineExecutionProvider"])
+
+    # ---- Tokenize & run text encoder ----
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_dir)
+    input_ids, attn_mask = tokenize(texts, tokenizer, args.max_seq_len)
+
+    txt_out = txt_sess.run(["text_features"], {
+        "input_ids": input_ids.astype(np.int32),
+        "attention_mask": attn_mask.astype(np.int32),
+    })
+    text_features = txt_out[0]  # (1, N, 768) — batch dim already included
+
+    # ---- Preprocess image & run image encoder ----
+    img_tensor, ratio, (pad_left, pad_top) = preprocess_image(args.image,
+                                                               args.image_size)
+
+    img_out = img_sess.run(None, {
+        "images": img_tensor,
+        "text_features": text_features.astype(np.float32),
+    })
+
+    # img_out = [cls_s8, bbox_s8, cls_s16, bbox_s16, cls_s32, bbox_s32]
+    scales = [
+        (img_out[0], img_out[1], 80, 80, 8),    # stride 8
+        (img_out[2], img_out[3], 40, 40, 16),   # stride 16
+        (img_out[4], img_out[5], 20, 20, 32),   # stride 32
+    ]
+
+    # ---- Decode all scales ----
+    all_boxes, all_scores, all_labels = [], [], []
+    for cls_scores, bbox_preds, h, w, stride in scales:
+        boxes, sc, _ = _decode_bboxes(cls_scores, bbox_preds, h, w, stride,
+                                      args.threshold)
+        if boxes.size == 0:
+            continue
+        labels = sc.argmax(axis=1)
+        scores = sc.max(axis=1)
+        all_boxes.append(boxes)
+        all_scores.append(scores)
+        all_labels.append(labels)
+
+    if not all_boxes:
+        print("No detections above threshold.")
+        return
+
+    boxes = np.concatenate(all_boxes)
+    scores = np.concatenate(all_scores)
+    labels = np.concatenate(all_labels)
+
+    # ---- NMS ----
+    keep = nms(boxes, scores, args.nms_threshold)
+    boxes, scores, labels = boxes[keep], scores[keep], labels[keep]
+
+    # ---- Rescale to original image coordinates ----
+    boxes[:, [0, 2]] -= pad_left
+    boxes[:, [1, 3]] -= pad_top
+    boxes /= ratio
+
+    # Clamp to image bounds
+    img = Image.open(args.image)
+    boxes[:, 0::2] = np.clip(boxes[:, 0::2], 0, img.size[0])
+    boxes[:, 1::2] = np.clip(boxes[:, 1::2], 0, img.size[1])
+
+    label_names = [texts[l] for l in labels]
+    print(f"Detections: {len(boxes)}")
+    for name, box, s in zip(label_names, boxes, scores):
+        print(f"  {name} {s:.3f}  ({box[0]:.0f}, {box[1]:.0f}, "
+              f"{box[2]:.0f}, {box[3]:.0f})")
+
+    # ---- Visualise ----
+    draw_boxes(args.image, boxes, label_names, scores, args.output,
+               font_path=args.font)
+
+
+if __name__ == "__main__":
+    main()

coco_zh_class_texts.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa535063b870659ec911d1cd93c5f9659465c25e86580d3e2b72d384dffe7b3f
+size 1047

config/wedetect_base.py

@@ -0,0 +1,216 @@
+_base_ = ["default_runtime.py"]
+
+# hyper-parameters
+num_classes = 1203
+num_training_classes = 80
+max_epochs = 80  # Maximum training epochs
+close_mosaic_epochs = 2
+save_epoch_intervals = 1
+text_channels = 768
+neck_embed_channels = [128, 256, 512]
+neck_num_heads = [4, 8,16]
+base_lr = 5e-4
+weight_decay = 0.05 / 2
+train_batch_size_per_gpu = 10
+
+find_unused_parameters = True
+
+model_test_cfg = dict(
+    # The config of multi-label for multi-class prediction.
+    multi_label=True,
+    # The number of boxes before NMS
+    nms_pre=30000,
+    score_thr=0.001,  # Threshold to filter out boxes.
+    nms=dict(type='nms', iou_threshold=0.7),  # NMS type and threshold
+    max_per_img=300)  # Max number of detections of each image
+
+tal_topk = 10  # Number of bbox selected in each level
+tal_alpha = 0.5  # A Hyper-parameter related to alignment_metrics
+tal_beta = 6.0  # A Hyper-parameter related to alignment_metrics
+# TODO: Automatically scale loss_weight based on number of detection layers
+loss_cls_weight = 0.5
+loss_bbox_weight = 7.5
+# Since the dfloss is implemented differently in the official
+# and mmdet, we're going to divide loss_weight by 4.
+loss_dfl_weight = 1.5 / 4
+
+custom_imports = dict(imports=["wedetect"], allow_failed_imports=False)
+# model settings
+model = dict(
+    type="YOLOWorldDetector",
+    mm_neck=False,
+    num_train_classes=num_training_classes,
+    num_test_classes=num_classes,
+    data_preprocessor=dict(
+        type="YOLOWDetDataPreprocessor",
+        mean=[0., 0., 0.],
+        std=[255., 255., 255.],
+        bgr_to_rgb=True),
+    backbone=dict(
+        type="MultiModalYOLOBackbone",
+        image_model=dict(
+            type="ConvNextVisionBackbone",
+            model_name="base",
+            frozen_modules=[],
+        ),
+        text_model=dict(
+            type="XLMRobertaLanguageBackbone",
+            model_name="./xlm-roberta-base/",
+            model_size="base",
+            frozen_modules=[],
+        ),
+    ),
+    neck=dict(
+        type="CSPRepBiFPANNeck",
+        scale_factor=1.0,
+        model_size = 'base',
+    ),
+    bbox_head=dict(
+        type="YOLOWorldHead",
+        head_module=dict(
+            type="YOLOWorldHeadModule",
+            use_bn_head=True,
+            embed_dims=text_channels,
+            num_classes=num_training_classes,
+            model_size = 'base',
+            in_channels=[256, 512, 1024],
+        ),
+        prior_generator=dict(
+            type='MlvlPointGenerator', offset=0.5, strides=[8, 16, 32]),
+        bbox_coder=dict(type='WeDetectDistancePointBBoxCoder'),
+        # scaled based on number of detection layers
+        loss_cls=dict(
+            type='CrossEntropyLoss',
+            use_sigmoid=True,
+            reduction='none',
+            loss_weight=loss_cls_weight),
+        loss_bbox=dict(
+            type='mmyoloIoULoss',
+            iou_mode='ciou',
+            bbox_format='xyxy',
+            reduction='sum',
+            loss_weight=loss_bbox_weight,
+            return_iou=False),
+        loss_dfl=dict(
+            type='DistributionFocalLoss',
+            reduction='mean',
+            loss_weight=loss_dfl_weight)),
+    train_cfg=dict(
+        assigner=dict(
+            type='BatchTaskAlignedAssigner',
+            num_classes=num_classes,
+            use_ciou=True,
+            topk=tal_topk,
+            alpha=tal_alpha,
+            beta=tal_beta,
+            eps=1e-9)),
+    test_cfg=model_test_cfg)
+
+img_scale = (640, 640)
+
+test_pipeline = [
+    dict(type='LoadImageFromFile', backend_args=None),
+    dict(type='WeDetectKeepRatioResize', scale=img_scale),
+    dict(
+        type='WeDetectLetterResize',
+        scale=img_scale,
+        allow_scale_up=False,
+        pad_val=dict(img=114)),
+    dict(type='LoadAnnotations', with_bbox=True, _scope_='mmdet'),
+    dict(type="LoadText"),
+    dict(
+        type="PackDetInputs",
+        meta_keys=(
+            "img_id",
+            "img_path",
+            "ori_shape",
+            "img_shape",
+            "scale_factor",
+            "pad_param",
+            "texts",
+        ),
+    ),
+]
+
+dist_cfg = dict(backend="nccl", timeout=10800)
+
+
+coco_val_dataset = dict(
+    type="MultiModalDataset",
+    dataset=dict(
+        type="WeCocoDataset",
+        data_root="data/coco/",
+        test_mode=True,
+        ann_file="data/coco/annotations/instances_val2017.json",
+        data_prefix=dict(img="val2017"),
+        batch_shapes_cfg=None,
+    ),
+    class_text_path="data/texts/coco_zh_class_texts.json",
+    pipeline=test_pipeline,
+)
+
+
+lvis_minival_dataset = dict(
+    type="MultiModalDataset",
+    dataset=dict(
+        type="YOLOv5LVISV1Dataset",
+        data_root=f"data/coco/",
+        test_mode=True,
+        ann_file=f"data/lvis/lvis_v1_minival_inserted_image_name.json",
+        data_prefix=dict(img=""),
+        batch_shapes_cfg=None,
+    ),
+    class_text_path=f"data/texts/lvis_v1_zh_class_texts.json",
+    pipeline=test_pipeline,
+)
+
+lvis_od_val_dataset = dict(
+    type="MultiModalDataset",
+    dataset=dict(
+        type="YOLOv5LVISV1Dataset",
+        data_root=f"data/coco/",
+        test_mode=True,
+        ann_file="data/lvis/lvis_od_val.json",
+        data_prefix=dict(img=""),
+        batch_shapes_cfg=None,
+    ),
+    class_text_path=f"data/texts/lvis_v1_zh_class_texts.json",
+    pipeline=test_pipeline,
+)
+coco_evaluator = dict(
+    type="CocoMetric",
+    ann_file=f"data/coco/annotations/instances_val2017.json",
+    metric="bbox",
+)
+
+lvis_minival_evaluator = dict(
+    type="LVISMetric",
+    ann_file=f"data/lvis/lvis_v1_minival_inserted_image_name.json",
+    metric="bbox",
+)
+lvis_od_val_evaluator = dict(
+    type="LVISMetric",
+    ann_file="data/lvis/lvis_od_val.json",
+    metric="bbox",
+)
+
+val_dataloader = dict(
+    batch_size=1,
+    num_workers=2,
+    persistent_workers=True,
+    pin_memory=True,
+    drop_last=False,
+    sampler=dict(type='DefaultSampler', shuffle=False),
+    dataset=coco_val_dataset)
+
+test_dataloader = val_dataloader
+
+# val_evaluator = _base_.lvis_minival_evaluator
+# val_evaluator = _base_.lvis_od_val_evaluator
+val_evaluator = coco_evaluator
+test_evaluator = val_evaluator
+
+
+val_cfg = dict(type='ValLoop')
+test_cfg = dict(type='TestLoop')
+

export_onnx.py

@@ -0,0 +1,289 @@
+"""
+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.")

generate_class_embedding.py

@@ -0,0 +1,161 @@
+import argparse
+import json
+import os
+import random
+from collections import OrderedDict
+from typing import List, Sequence, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from transformers import AutoTokenizer, AutoConfig, XLMRobertaModel
+except ImportError:
+    AutoTokenizer = None
+    HFBertModel = None
+
+
+class XLMRobertaLanguageBackbone(nn.Module):
+
+    def __init__(
+        self,
+        ckpt_path,
+        frozen_modules: Sequence[str] = (),
+        dropout: float = 0.0,
+        init_cfg= None,
+    ) -> None:
+
+        super().__init__()
+        if 'base' in ckpt_path:
+            self.head = nn.Linear(768, 768, bias=True) # XLarge
+            model_name = "./xlm-roberta-base/"
+        elif 'large' in ckpt_path:
+            self.head = nn.Linear(1024, 768, bias=True) # XLarge
+            model_name = "./xlm-roberta-large/"
+
+        self.frozen_modules = frozen_modules
+        cfg = AutoConfig.from_pretrained(model_name)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = XLMRobertaModel(cfg)
+        self.language_dim = cfg.hidden_size
+        
+        
+        # 加载 model 权重
+        new_state_dict = OrderedDict()
+        state_dict = torch.load(
+            ckpt_path,
+            map_location="cpu",
+            weights_only=False,
+        )['state_dict']
+        for k, v in state_dict.items():
+            if k.startswith('backbone.text_model.'):
+                name = k.split("backbone.text_model.")[-1]
+                new_state_dict[name] = v
+        msg = self.load_state_dict(new_state_dict, strict=True)
+        print(msg)
+
+        print("TEXT-ENCODER xlm-roberta-base LOADING WEIGHTS !!!!")
+
+
+
+    def forward(self, text: List[str], max_seq_len: int = 32):
+        text = self.tokenizer(text=text, return_tensors="pt",
+                              padding="max_length", max_length=max_seq_len)
+        text = text.to(device=self.model.device)
+
+        txt_feats = self.model(**text)["last_hidden_state"][:, 0]
+        txt_feats = self.head(txt_feats)
+
+        return txt_feats
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--wedetect_checkpoint', type=str, default='checkpoints/wedetect_base.pth')
+    parser.add_argument('--classname_file', type=str, default='data/texts/coco_zh_class_texts.json')
+    parser.add_argument('--max-seq-len', type=int, default=32,
+                        help='Fixed token length (must match ONNX export).')
+    parser.add_argument('--num-classes-per-group', type=int, default=4,
+                        help='Number of classes per group npy.')
+    parser.add_argument('--num-groups', type=int, default=64,
+                        help='Number of random groups to generate.')
+    parser.add_argument('--calib-dir', type=str, default='calib_data',
+                        help='Directory for text-encoder quantisation calibration data.')
+    args = parser.parse_args()
+
+    with open(args.classname_file) as f:
+        name_chinese = json.load(f)
+    name_chinese = [name[0] for name in name_chinese]
+
+    language_encoder = XLMRobertaLanguageBackbone(args.wedetect_checkpoint).cuda()
+
+    # Generate random groups: each group picks 4 random classes → shape (1, 4, 768)
+    total_classes = len(name_chinese)
+    print(f"Total classes: {total_classes}  → Generating {args.num_groups} random groups")
+    # Generate calibration data for text-encoder quantisation
+    # Directories: calib_dir/input_ids/  calib_dir/attention_mask/
+    calib_input_ids = os.path.join(args.calib_dir, "input_ids")
+    calib_attn_mask = os.path.join(args.calib_dir, "attention_mask")
+    for d in (calib_input_ids, calib_attn_mask):
+        os.makedirs(d, exist_ok=True)
+
+    tokenizer = language_encoder.tokenizer
+
+    for g in range(args.num_groups):
+        idx = random.sample(range(total_classes), args.num_classes_per_group)
+        group_texts = [name_chinese[i] for i in idx]
+        tokens = tokenizer(group_texts, padding="max_length",
+                           max_length=args.max_seq_len, return_tensors="np")
+
+        np.save(os.path.join(calib_input_ids, f"{g:03d}.npy"),
+                tokens["input_ids"].astype(np.int64))
+        np.save(os.path.join(calib_attn_mask, f"{g:03d}.npy"),
+                tokens["attention_mask"].astype(np.int64))
+        print(f"calib [{g:03d}] input_ids: {tokens['input_ids'].shape}  "
+              f"classes: {group_texts}")
+
+    # Compress each subdirectory to .tar.gz
+    import tarfile
+    for sub_name in ("input_ids", "attention_mask"):
+        sub_dir = os.path.join(args.calib_dir, sub_name)
+        tar_path = os.path.join(args.calib_dir, f"{sub_name}.tar.gz")
+        with tarfile.open(tar_path, "w:gz") as tar:
+            for fname in sorted(os.listdir(sub_dir)):
+                tar.add(os.path.join(sub_dir, fname), arcname=fname)
+        print(f"Compressed: {tar_path}")
+
+    print(f"Saved calibration data to {args.calib_dir}/")
+
+    # -------------------------------------------------------------------
+    # Generate 64 random 4-class text embeddings  →  class_embedding_4cls/
+    # Each file: (1, 4, 768) float32, L2-normalised (same as the ONNX
+    # image encoder expects via text_features input).
+    # -------------------------------------------------------------------
+    embed_dir = os.path.join(args.calib_dir, "class_embedding_4cls")
+    os.makedirs(embed_dir, exist_ok=True)
+
+    print(f"\nGenerating {args.num_groups} random {args.num_classes_per_group}-class "
+          f"text embeddings → {embed_dir}/")
+    for g in range(args.num_groups):
+        idx = random.sample(range(total_classes), args.num_classes_per_group)
+        group_texts = [name_chinese[i] for i in idx]
+        with torch.no_grad():
+            feats = language_encoder(group_texts, max_seq_len=args.max_seq_len)
+            feats = F.normalize(feats, dim=-1).unsqueeze(0)  # (1, 4, 768)
+        fpath = os.path.join(embed_dir, f"{g:03d}.npy")
+        np.save(fpath, feats.cpu().numpy().astype(np.float32))
+        if (g + 1) % 16 == 0 or g == args.num_groups - 1:
+            print(f"  [{g + 1:3d}/{args.num_groups}]  shape={feats.shape}  "
+                  f"classes: {group_texts}")
+
+    # Compress
+    tar_path = os.path.join(args.calib_dir, "class_embedding_4cls.tar.gz")
+    with tarfile.open(tar_path, "w:gz") as tar:
+        for fname in sorted(os.listdir(embed_dir)):
+            tar.add(os.path.join(embed_dir, fname), arcname=fname)
+    print(f"Compressed: {tar_path}")
+
+    print(f"Saved calibration data to {args.calib_dir}/")

onnx_infer.py

@@ -0,0 +1,278 @@
+"""
+ONNX runtime inference for WeDetect.
+
+Usage:
+  python onnx_infer.py --image assets/demo.jpeg --text "鞋,床" --threshold 0.3
+
+Dependencies: onnxruntime, numpy, PIL, transformers (tokenizer only)
+"""
+
+import argparse
+import os
+import sys
+
+import numpy as np
+import onnxruntime as ort
+from PIL import Image, ImageDraw, ImageFont
+from transformers import AutoTokenizer
+
+# ---------------------------------------------------------------------------
+# Image preprocessing (mirrors WeDetectKeepRatioResize + LetterResize)
+# ---------------------------------------------------------------------------
+
+def letterbox(image: Image.Image, new_shape=(640, 640), color=(114, 114, 114)):
+    """Resize keeping aspect ratio and pad to ``new_shape``."""
+    ow, oh = image.size
+    r = min(new_shape[0] / ow, new_shape[1] / oh)
+    new_w, new_h = int(round(ow * r)), int(round(oh * r))
+    image = image.resize((new_w, new_h), Image.BILINEAR)
+
+    # Paste onto a blank canvas
+    canvas = Image.new("RGB", new_shape, color)
+    left = (new_shape[0] - new_w) // 2
+    top = (new_shape[1] - new_h) // 2
+    canvas.paste(image, (left, top))
+
+    return canvas, r, (left, top)
+
+
+def preprocess_image(image_path: str, image_size=640):
+    """Load, letterbox, and normalise an image.
+    """
+    img = Image.open(image_path).convert("RGB")
+    img, ratio, (pad_left, pad_top) = letterbox(img, (image_size, image_size))
+
+    arr = np.array(img, dtype=np.float32) / 255.0       # [0, 1], HWC
+    tensor = arr.transpose(2, 0, 1)[None]                # NCHW
+    return tensor, ratio, (pad_left, pad_top)
+
+
+# ---------------------------------------------------------------------------
+# Tokenization
+# ---------------------------------------------------------------------------
+
+def tokenize(texts, tokenizer, max_seq_len=32):
+    """Tokenize category names into fixed-length tensors.
+
+    Returns
+    -------
+    input_ids      : np.ndarray  shape (N, max_seq_len)  int64
+    attention_mask : np.ndarray  shape (N, max_seq_len)  int64
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+    tokens = tokenizer(texts, padding="max_length", max_length=max_seq_len,
+                       return_tensors="np")
+    return tokens["input_ids"], tokens["attention_mask"]
+
+
+# ---------------------------------------------------------------------------
+# BBox decoding
+# ---------------------------------------------------------------------------
+
+def _grid_points(h, w, stride):
+    """Generate grid anchor points for one feature-map scale."""
+    x = (np.arange(w) + 0.5) * stride
+    y = (np.arange(h) + 0.5) * stride
+    yy, xx = np.meshgrid(y, x, indexing="ij")
+    return np.stack([xx, yy], axis=-1).reshape(-1, 2)  # (H*W, 2)
+
+
+def _decode_bboxes(cls_scores, bbox_preds, h, w, stride, score_thr):
+    """Decode one scale: sigmoid → filter → distance2bbox."""
+    # cls_scores:  (1, C, H, W)   bbox_preds: (1, 4, H, W)
+    C = cls_scores.shape[1]
+    scores = 1.0 / (1.0 + np.exp(-cls_scores[0]))  # sigmoid, (C, H, W)
+    scores = scores.reshape(C, -1).transpose(1, 0)  # (H*W, C)
+
+    bbox = bbox_preds[0].reshape(4, -1).transpose(1, 0)     # (H*W, 4)
+    points = _grid_points(h, w, stride)                     # (H*W, 2)
+
+    bbox = bbox * stride  # scale distances to pixels
+    x1 = points[:, 0] - bbox[:, 0]
+    y1 = points[:, 1] - bbox[:, 1]
+    x2 = points[:, 0] + bbox[:, 2]
+    y2 = points[:, 1] + bbox[:, 3]
+    boxes = np.stack([x1, y1, x2, y2], axis=-1)             # (H*W, 4)
+
+    # Filter by score
+    max_scores = scores.max(axis=1)
+    keep = max_scores > score_thr
+    return boxes[keep], scores[keep], max_scores[keep]
+
+
+# ---------------------------------------------------------------------------
+# NMS
+# ---------------------------------------------------------------------------
+
+def nms(boxes, scores, iou_threshold=0.7, max_dets=300):
+    """Simple numpy NMS."""
+    order = scores.argsort()[::-1]
+    keep = []
+    while order.size > 0 and len(keep) < max_dets:
+        i = order[0]
+        keep.append(i)
+        if order.size == 1:
+            break
+        ious = _box_iou(boxes[i:i + 1], boxes[order[1:]])[0]
+        order = order[1:][ious < iou_threshold]
+    return np.array(keep, dtype=np.int64) if keep else np.array([], dtype=np.int64)
+
+
+def _box_iou(boxes1, boxes2):
+    """Pairwise IoU between two sets of xyxy boxes."""
+    area1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
+    area2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])
+
+    inter_x1 = np.maximum(boxes1[:, None, 0], boxes2[None, :, 0])
+    inter_y1 = np.maximum(boxes1[:, None, 1], boxes2[None, :, 1])
+    inter_x2 = np.minimum(boxes1[:, None, 2], boxes2[None, :, 2])
+    inter_y2 = np.minimum(boxes1[:, None, 3], boxes2[None, :, 3])
+    inter_w = np.maximum(0, inter_x2 - inter_x1)
+    inter_h = np.maximum(0, inter_y2 - inter_y1)
+    inter = inter_w * inter_h
+    return inter / (area1[:, None] + area2[None, :] - inter + 1e-7)
+
+# ---------------------------------------------------------------------------
+# Visualisation
+# ---------------------------------------------------------------------------
+
+def draw_boxes(image_path, boxes, labels, scores, output_path, font_path=None):
+    """Draw bounding boxes with Chinese-capable font and save."""
+
+    img = Image.open(image_path).convert("RGB")
+    draw = ImageDraw.Draw(img)
+
+    font = ImageFont.truetype(font_path, size=18)
+    font_small = ImageFont.truetype(font_path, size=14)
+
+    colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0),
+              (255, 0, 255), (0, 255, 255), (128, 0, 128), (255, 165, 0)]
+
+    for i, (box, label, score) in enumerate(zip(boxes, labels, scores)):
+        x1, y1, x2, y2 = box.astype(int)
+        c = colors[i % len(colors)]
+        draw.rectangle([x1, y1, x2, y2], outline=c, width=2)
+
+        # Use textbbox for accurate label background sizing (PIL >= 8.0)
+        text = f"{label} {score:.2f}"
+        if hasattr(draw, "textbbox"):
+            bbox = draw.textbbox((x1 + 2, y1 + 2), text, font=font)
+            draw.rectangle([x1, y1, bbox[2] + 4, bbox[3] + 2], fill=c)
+            draw.text((x1 + 2, y1 + 2), text, fill="white", font=font)
+        else:  # older PIL fallback
+            draw.rectangle([x1, y1, x1 + len(text) * 10, y1 + 20], fill=c)
+            draw.text((x1 + 2, y1 + 2), text, fill="white", font=font)
+
+    img.save(output_path)
+    print(f"Saved: {output_path}")
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def parse_args():
+    p = argparse.ArgumentParser(description="WeDetect ONNX inference")
+    p.add_argument("--image", default='assets/demo.jpeg', help="Input image path")
+    p.add_argument("--text", default="鞋,床,人,衣架",
+                   help="Chinese class names separated by comma, e.g. '鞋,床,人,衣架'")
+    p.add_argument("--image_model", default="wedetect_image_encoder.onnx")
+    p.add_argument("--text_model", default="wedetect_text_encoder.onnx")
+    p.add_argument("--tokenizer-dir", default="./xlm-roberta-base/")
+    p.add_argument("--max-seq-len", type=int, default=32)
+    p.add_argument("--threshold", type=float, default=0.3)
+    p.add_argument("--nms-threshold", type=float, default=0.7)
+    p.add_argument("--output", default="onnx_res.jpg")
+    p.add_argument("--image-size", type=int, default=640)
+    p.add_argument("--font", default='wqy-microhei.ttc',
+                   help="Path to a .ttf/.ttc font file for Chinese label rendering. "
+                        "Auto-detected if not specified.")
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    # ---- Parse texts ----
+    texts = [t.strip() for t in args.text.split(",")]
+    print(f"Classes: {texts}")
+
+    # ---- Load ONNX sessions ----
+    img_sess = ort.InferenceSession(args.image_model,providers=["CPUExecutionProvider"])
+    txt_sess = ort.InferenceSession(args.text_model,providers=["CPUExecutionProvider"])
+
+    # ---- Tokenize & run text encoder ----
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_dir)
+    input_ids, attn_mask = tokenize(texts, tokenizer, args.max_seq_len)
+
+    txt_out = txt_sess.run(["text_features"], {
+        "input_ids": input_ids,
+        "attention_mask": attn_mask,
+    })
+    text_features = txt_out[0]  # (1, N, 768) — batch dim already included
+
+    # ---- Preprocess image & run image encoder ----
+    img_tensor, ratio, (pad_left, pad_top) = preprocess_image(args.image,
+                                                               args.image_size)
+
+    img_out = img_sess.run(None, {
+        "images": img_tensor,
+        "text_features": text_features.astype(np.float32),
+    })
+
+    # img_out = [cls_s8, bbox_s8, cls_s16, bbox_s16, cls_s32, bbox_s32]
+    scales = [
+        (img_out[0], img_out[1], 80, 80, 8),    # stride 8
+        (img_out[2], img_out[3], 40, 40, 16),   # stride 16
+        (img_out[4], img_out[5], 20, 20, 32),   # stride 32
+    ]
+
+    # ---- Decode all scales ----
+    all_boxes, all_scores, all_labels = [], [], []
+    for cls_scores, bbox_preds, h, w, stride in scales:
+        boxes, sc, _ = _decode_bboxes(cls_scores, bbox_preds, h, w, stride,
+                                      args.threshold)
+        if boxes.size == 0:
+            continue
+        labels = sc.argmax(axis=1)
+        scores = sc.max(axis=1)
+        all_boxes.append(boxes)
+        all_scores.append(scores)
+        all_labels.append(labels)
+
+    if not all_boxes:
+        print("No detections above threshold.")
+        return
+
+    boxes = np.concatenate(all_boxes)
+    scores = np.concatenate(all_scores)
+    labels = np.concatenate(all_labels)
+
+    # ---- NMS ----
+    keep = nms(boxes, scores, args.nms_threshold)
+    boxes, scores, labels = boxes[keep], scores[keep], labels[keep]
+
+    # ---- Rescale to original image coordinates ----
+    boxes[:, [0, 2]] -= pad_left
+    boxes[:, [1, 3]] -= pad_top
+    boxes /= ratio
+
+    # Clamp to image bounds
+    img = Image.open(args.image)
+    boxes[:, 0::2] = np.clip(boxes[:, 0::2], 0, img.size[0])
+    boxes[:, 1::2] = np.clip(boxes[:, 1::2], 0, img.size[1])
+
+    label_names = [texts[l] for l in labels]
+    print(f"Detections: {len(boxes)}")
+    for name, box, s in zip(label_names, boxes, scores):
+        print(f"  {name} {s:.3f}  ({box[0]:.0f}, {box[1]:.0f}, "
+              f"{box[2]:.0f}, {box[3]:.0f})")
+
+    # ---- Visualise ----
+    draw_boxes(args.image, boxes, label_names, scores, args.output,
+               font_path=args.font)
+
+
+if __name__ == "__main__":
+    main()

quant/attention_mask.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d576d7af2cf122f5600504e6454a54d3d83411bf5f8ed045353b2980f2dfeb80
+size 2092

quant/class_embedding_4cls.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0cb8759b85dd87945b3dd2c072141f5a4958a923257144783c9d3e4f835c8e19
+size 657358

quant/coco100.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:110af521bd5dec72b3626909f3bbb94fc6b91498117298c6f3d25a4f51b7cf43
+size 15039687

quant/image_encoder_4cls.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:632af6200435935d8e82fab173f7a56d5a5048236c582c0a0459a85a6e0c0695
+size 1283

quant/input_ids.tar.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f2135255f2a08c565fa91aa456b794a60909928426cca9f2514ec63c70ed791a
+size 3444

quant/text_encoder_4cls.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bfa19b3f84a5ce68c23b029528978cf9d9d41e626a84e3dc99589d7024f630f6
+size 1182

wqy-microhei.ttc

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2420e8078af796b19a3f6ef13de527a1a91c1e7171eea115926c614ced1009b3
+size 5177387

xlm-roberta-base/config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:618c004597964d75b68497181ca769a69ea5719524c812edb81cfd74d9014656
+size 662

xlm-roberta-base/sentencepiece.bpe.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cfc8146abe2a0488e9e2a0c56de7952f7c11ab059eca145a0a727afce0db2865
+size 5069051

xlm-roberta-base/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a898ea75433890f6610f4e470b8ebeb0c21dce5c8dd61f892eb09eb5919d2e2c
+size 9096718