#!/usr/bin/env python3 # Usage: python ax_infer.py --model-path yolov8n_640x640.axmodel --test-img bus.jpg --img-save-path result_yolov8_det.jpg --score-thres 0.25 --nms-thres 0.7 --providers AxEngineExecutionProvider import os import cv2 import numpy as np from time import time import argparse import logging import axengine as ort logging.basicConfig( level=logging.DEBUG, format='[%(name)s] [%(asctime)s.%(msecs)03d] [%(levelname)s] %(message)s', datefmt='%H:%M:%S' ) logger = logging.getLogger("YOLOv8-Det") def infer_hw_layout(shape): """Infer input height, width and layout from model input shape.""" shape = list(shape) if len(shape) == 4 and shape[-1] == 3: h = int(shape[1] or 640) w = int(shape[2] or 640) return h, w, "NHWC" if len(shape) == 4 and shape[1] == 3: h = int(shape[2] or 640) w = int(shape[3] or 640) return h, w, "NCHW" return 640, 640, "NCHW" def softmax(x, axis=-1): """Compute softmax along axis.""" e_x = np.exp(x - np.max(x, axis=axis, keepdims=True)) return e_x / np.sum(e_x, axis=axis, keepdims=True) def dfl_decode(box_pred, reg_max=16): """Decode DFL (Distribution Focal Loss) box predictions to ltrb distances.""" N = box_pred.shape[0] box_pred = box_pred.reshape(N, 4, reg_max) box_pred = softmax(box_pred, axis=-1) proj = np.arange(reg_max, dtype=np.float32) return np.sum(box_pred * proj, axis=-1) # (N, 4) def main(): ap = argparse.ArgumentParser(description='YOLOv8-Det Inference (AXERARuntime)') ap.add_argument('--model-path', type=str, default='yolov8n_640x640.axmodel') ap.add_argument('--test-img', type=str, default='bus.jpg') ap.add_argument('--img-save-path', type=str, default='result_yolov8_det.jpg') ap.add_argument('--score-thres', type=float, default=0.25) ap.add_argument('--nms-thres', type=float, default=0.7) ap.add_argument('--providers', type=str, default='AxEngineExecutionProvider') opt = ap.parse_args() if not os.path.exists(opt.model_path): logger.error(f"Model not found: {opt.model_path}") return t0 = time() providers = [p.strip() for p in opt.providers.split(",") if p.strip()] or None sess = ort.InferenceSession(opt.model_path, providers=providers) logger.debug(f"\033[1;31mLoad model time = {(time() - t0) * 1000:.2f} ms\033[0m") inp = sess.get_inputs()[0] input_name = inp.name m_h, m_w, layout = infer_hw_layout(inp.shape) img = cv2.imread(opt.test_img) if img is None: logger.error(f"Image not found or unreadable: {opt.test_img}") return # Preprocess t0 = time() orig_h, orig_w = img.shape[:2] scale = min(m_h / orig_h, m_w / orig_w) new_w, new_h = int(orig_w * scale), int(orig_h * scale) resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR) padded = cv2.copyMakeBorder( resized, 0, m_h - new_h, 0, m_w - new_w, cv2.BORDER_CONSTANT, value=(127, 127, 127) ) rgb = cv2.cvtColor(padded, cv2.COLOR_BGR2RGB) input_tensor = rgb[None, ...].astype(np.uint8) if layout == "NHWC" else np.transpose(rgb, (2, 0, 1))[None, ...].astype(np.uint8) logger.debug(f"\033[1;31mPre-process time = {(time() - t0) * 1000:.2f} ms\033[0m") # Inference t0 = time() ort_outputs = sess.run(None, {input_name: input_tensor}) out_metas = sess.get_outputs() logger.debug(f"\033[1;31mForward time = {(time() - t0) * 1000:.2f} ms\033[0m") # Post-process t0 = time() strides = (8, 16, 32) conf_raw = -np.log(1 / opt.score_thres - 1) detections = [] output_items = [] for meta, data in zip(out_metas, ort_outputs): shape = list(meta.shape) if any(s is None or isinstance(s, str) for s in shape): shape = list(data.shape) output_items.append((data, shape)) # Process each scale (6 outputs: 2 outputs per scale [box, cls]) for scale_idx, stride in enumerate(strides): box_idx = scale_idx * 2 cls_idx = scale_idx * 2 + 1 if box_idx >= len(output_items) or cls_idx >= len(output_items): continue box_data, box_shape = output_items[box_idx] cls_data, cls_shape = output_items[cls_idx] H, W = box_shape[1], box_shape[2] box_channels = box_shape[-1] # Determine if DFL is used (YOLOv8: 4*reg_max = 64) reg_max = None if box_channels > 4 and box_channels % 4 == 0: reg_max = box_channels // 4 box_data = box_data[0].reshape(-1, box_channels) cls_data = cls_data[0].reshape(-1, cls_shape[-1]) # Get max class scores cls_scores = np.max(cls_data, axis=1) cls_ids = np.argmax(cls_data, axis=1) valid_mask = cls_scores >= conf_raw if not np.any(valid_mask): continue v_box = box_data[valid_mask] v_score = 1 / (1 + np.exp(-cls_scores[valid_mask])) v_id = cls_ids[valid_mask] gy, gx = np.indices((H, W)) anchors = np.stack((gx.ravel(), gy.ravel()), axis=-1).astype(np.float32) + 0.5 anchors = anchors[valid_mask] # Decode boxes (DFL -> ltrb -> xyxy) if reg_max is not None: v_box = dfl_decode(v_box, reg_max) lt = v_box[:, :2] rb = v_box[:, 2:] x1y1 = anchors - lt x2y2 = anchors + rb boxes = np.hstack([x1y1, x2y2]) * stride for i in range(len(boxes)): detections.append([*boxes[i], v_score[i], v_id[i]]) logger.debug(f"\033[1;31mPost-process time = {(time() - t0) * 1000:.2f} ms\033[0m") if len(detections) == 0: logger.info("No detections found.") cv2.imwrite(opt.img_save_path, img) return detections = np.array(detections) # NMS xywh = detections[:, :4].copy() xywh[:, 2] = xywh[:, 2] - xywh[:, 0] # w xywh[:, 3] = xywh[:, 3] - xywh[:, 1] # h indices = cv2.dnn.NMSBoxes(xywh.tolist(), detections[:, 4].tolist(), opt.score_thres, opt.nms_thres) if len(indices) == 0: logger.info("No detections after NMS.") cv2.imwrite(opt.img_save_path, img) return indices = indices.flatten() final_dets = detections[indices] # Scale to original image coordinates final_dets[:, :4] = final_dets[:, :4] / scale final_dets[:, [0, 2]] = np.clip(final_dets[:, [0, 2]], 0, orig_w) final_dets[:, [1, 3]] = np.clip(final_dets[:, [1, 3]], 0, orig_h) # COCO class names coco_names = [ "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush" ] base_colors = [ (255, 42, 4), (235, 219, 11), (243, 243, 243), (183, 223, 0), (104, 31, 17), (221, 111, 255), (79, 68, 255), (0, 237, 204), (68, 243, 0), (255, 0, 189), (255, 180, 0), (186, 0, 221), (255, 255, 0), (0, 192, 38), (179, 255, 1), (255, 36, 125), (104, 0, 123), (108, 27, 255), (47, 109, 252), (11, 255, 162), ] logger.info(f"\033[1;32mDraw Results ({len(final_dets)} objects): \033[0m") for det in final_dets: box = det[:4].astype(int) conf = det[4] cls_id = int(det[5]) color = [int(c) for c in base_colors[cls_id % len(base_colors)]] cls_name = coco_names[cls_id] if cls_id < len(coco_names) else str(cls_id) logger.info(f"({box[0]}, {box[1]}, {box[2]}, {box[3]}) -> {cls_name}: {conf:.2f}") cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), color, 2) label = f"{cls_name} {conf:.2f}" cv2.putText(img, label, (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) cv2.imwrite(opt.img_save_path, img) logger.info(f"Saved to {opt.img_save_path}") if __name__ == "__main__": main()