axmodel_infer.py

"""
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()