| 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 YOLO26n-OBB QNN model with AidLite.") |
| parser.add_argument( |
| "--target_model", |
| type=str, |
| default="./models/QCS8550/FP16/yolo26n-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()) |
|
|
