File size: 1,793 Bytes
aa718fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | 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 YOLO26l-OBB QNN model with AidLite.")
parser.add_argument(
"--target_model",
type=str,
default="./models/QCS8550/FP16/yolo26l-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())
|