| from __future__ import annotations | |
| import argparse | |
| from pathlib import Path | |
| import cv2 | |
| from utils import draw_obb_res | |
| from yolo_obb import Yolo26Obb | |
| def parse_args(): | |
| parser = argparse.ArgumentParser(description="Run YOLO26x-OBB QNN model with AidLite.") | |
| parser.add_argument( | |
| "--target_model", | |
| type=str, | |
| default="./models/QCS8550/FP16/yolo26x-obb_qcs8550_fp16.qnn248.ctx.bin", | |
| help=( | |
| "ctx.bin 路径，可换成 ./models/QCS8550/W8A16/ 或 ./models/QCS8550/W8A8/ 下的量化模型" | |
| ), | |
| ) | |
| parser.add_argument("--imgs", type=str, default="./code/python/boats.jpg") | |
| parser.add_argument("--height", type=int, default=1024) | |
| parser.add_argument("--width", type=int, default=1024) | |
| parser.add_argument("--invoke_nums", type=int, default=10) | |
| parser.add_argument("--model_type", type=str, default="QNN") | |
| parser.add_argument("--conf_thres", type=float, default=0.25) | |
| parser.add_argument("--iou_thres", type=float, default=0.45) | |
| parser.add_argument("--save_path", type=str, default="./code/python/result.jpg") | |
| return parser.parse_args() | |
| def main() -> int: | |
| args = parse_args() | |
| frame = cv2.imread(args.imgs) | |
| if frame is None: | |
| raise FileNotFoundError(args.imgs) | |
| model = Yolo26Obb(args.target_model, args.width, args.height, args.model_type) | |
| detections = model(frame, args.invoke_nums, args.conf_thres, args.iou_thres) | |
| print("===================") | |
| print(f"Detect {len(detections)} targets.") | |
| result = draw_obb_res(frame, detections) | |
| Path(args.save_path).parent.mkdir(parents=True, exist_ok=True) | |
| cv2.imwrite(args.save_path, result) | |
| return 0 | |
| if __name__ == "__main__": | |
| raise SystemExit(main()) | |