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