from __future__ import annotations import argparse import colorsys from pathlib import Path import numpy as np from PIL import Image, ImageDraw, ImageFont try: import axengine as ort BACKEND = "axengine" print("Running on AXera NPU (axengine)...") except ImportError: import onnxruntime as ort BACKEND = "onnxruntime" print("Running on CPU/GPU (onnxruntime)...") MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) CLASSES = [ "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", ] COCO_IDS = [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90, ] CLASS_NAME_BY_ID = {cid: name for cid, name in zip(COCO_IDS, CLASSES)} def sigmoid(x: np.ndarray) -> np.ndarray: return 1.0 / (1.0 + np.exp(-np.clip(x, -88.0, 88.0))) def get_numpy_dtype(input_meta: object) -> np.dtype: if hasattr(input_meta, "dtype"): return np.dtype(input_meta.dtype) ort_type = getattr(input_meta, "type", "") mapping = { "tensor(float)": np.float32, "tensor(float16)": np.float16, "tensor(uint8)": np.uint8, "tensor(int8)": np.int8, "tensor(int32)": np.int32, "tensor(int64)": np.int64, } if ort_type not in mapping: raise ValueError(f"Unsupported input type: {ort_type}") return np.dtype(mapping[ort_type]) def preprocess( image_path: str, input_h: int, input_w: int, layout: str, dtype: np.dtype, letterbox: bool, ) -> tuple[np.ndarray, Image.Image, dict[str, float]]: raw_image = Image.open(image_path).convert("RGB") orig_w, orig_h = raw_image.size if letterbox: scale_x = min(input_w / orig_w, input_h / orig_h) scale_y = scale_x resized_w = max(1, int(round(orig_w * scale_x))) resized_h = max(1, int(round(orig_h * scale_y))) pad_x = (input_w - resized_w) // 2 pad_y = (input_h - resized_h) // 2 resized = raw_image.resize((resized_w, resized_h), Image.Resampling.BILINEAR) canvas = Image.new("RGB", (input_w, input_h), (0, 0, 0)) canvas.paste(resized, (pad_x, pad_y)) image = np.array(canvas) else: resized = raw_image.resize((input_w, input_h), Image.Resampling.BILINEAR) image = np.array(resized) pad_x = 0.0 pad_y = 0.0 scale_x = input_w / orig_w scale_y = input_h / orig_h if BACKEND == "axengine": if layout == "NHWC": tensor = image[None, ...].astype(dtype, copy=False) else: tensor = image.transpose(2, 0, 1)[None, ...].astype(dtype, copy=False) else: if layout == "NHWC": if dtype != np.uint8: raise ValueError(f"NHWC input only supports uint8 in this simple script, got {dtype}") tensor = image[None, ...].astype(np.uint8) else: tensor = image.astype(np.float32) / 255.0 tensor = (tensor - MEAN) / STD tensor = tensor.transpose(2, 0, 1)[None, ...].astype(dtype) meta = { "orig_w": float(orig_w), "orig_h": float(orig_h), "scale_x": float(scale_x), "scale_y": float(scale_y), "pad_x": float(pad_x), "pad_y": float(pad_y), "input_w": float(input_w), "input_h": float(input_h), } return tensor, raw_image, meta def decode(outputs: list[np.ndarray], meta: dict[str, float], thresh: float) -> list[tuple[np.ndarray, float, int, int]]: dets = outputs[0][0] labels = outputs[1][0] if dets.shape[-1] != 4: dets, labels = labels, dets boxes = dets logits = labels[:, :-1] probs = sigmoid(logits) input_w = meta["input_w"] input_h = meta["input_h"] orig_w = meta["orig_w"] orig_h = meta["orig_h"] scale_x = meta["scale_x"] scale_y = meta["scale_y"] pad_x = meta["pad_x"] pad_y = meta["pad_y"] flat = probs.reshape(-1) topk = min(300, flat.size) top_idx = np.argpartition(-flat, topk - 1)[:topk] top_idx = top_idx[np.argsort(-flat[top_idx])] num_classes = probs.shape[1] results = [] for rank, idx in enumerate(top_idx.tolist()): query_id = idx // num_classes label_id = idx % num_classes score = float(flat[idx]) if score < thresh: continue cx, cy, bw, bh = boxes[query_id] x1 = ((cx - bw / 2.0) * input_w - pad_x) / scale_x y1 = ((cy - bh / 2.0) * input_h - pad_y) / scale_y x2 = ((cx + bw / 2.0) * input_w - pad_x) / scale_x y2 = ((cy + bh / 2.0) * input_h - pad_y) / scale_y x1 = max(0.0, min(orig_w, x1)) y1 = max(0.0, min(orig_h, y1)) x2 = max(0.0, min(orig_w, x2)) y2 = max(0.0, min(orig_h, y2)) results.append((np.array([x1, y1, x2, y2]), score, label_id, query_id)) return results def color_for_label(label_id: int) -> tuple[int, int, int]: hue = (label_id * 0.61803398875) % 1.0 r, g, b = colorsys.hsv_to_rgb(hue, 0.75, 1.0) return int(r * 255), int(g * 255), int(b * 255) def draw(raw_img: Image.Image, results: list[tuple[np.ndarray, float, int, int]], output_path: str) -> None: draw_obj = ImageDraw.Draw(raw_img) try: font = ImageFont.truetype("DejaVuSans.ttf", 18) except OSError: font = ImageFont.load_default() for box, score, label_id, query_id in results: x1, y1, x2, y2 = box.tolist() name = CLASS_NAME_BY_ID.get(label_id, f"obj_{label_id}") text = f"{name} {score:.2f}" color = color_for_label(label_id) draw_obj.rectangle([x1, y1, x2, y2], outline=color, width=3) draw_obj.rectangle([x1, max(0, y1 - 22), x1 + 140, y1], fill=color) draw_obj.text((x1 + 2, max(0, y1 - 20)), text, fill="black", font=font) print(f"query={query_id:3d} class={name:<15} score={score:.4f} box=({x1:.1f}, {y1:.1f}, {x2:.1f}, {y2:.1f})") Path(output_path).parent.mkdir(parents=True, exist_ok=True) raw_img.save(output_path) def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--model", type=str, required=True) parser.add_argument("--img", type=str, required=True) parser.add_argument("--output", type=str, default="result.jpg") parser.add_argument("--thresh", type=float, default=0.3) parser.add_argument("--letterbox", action="store_true", help="use letterbox resize instead of direct resize") args = parser.parse_args() session = ort.InferenceSession(args.model) input_meta = session.get_inputs()[0] shape = [int(x) for x in input_meta.shape] dtype = get_numpy_dtype(input_meta) if shape[1] in (1, 3, 4): layout = "NCHW" input_h, input_w = shape[2], shape[3] else: layout = "NHWC" input_h, input_w = shape[1], shape[2] print(f"input_name={input_meta.name} shape={shape} dtype={dtype} layout={layout}") img_tensor, raw_img, meta = preprocess(args.img, input_h, input_w, layout, dtype, args.letterbox) outputs = session.run(None, {input_meta.name: img_tensor}) results = decode(outputs, meta, args.thresh) print(f"Detected {len(results)} objects.") draw(raw_img, results, args.output) print(f"Result saved to {args.output}") if __name__ == "__main__": main()