33 lines
1.2 KiB
Python
33 lines
1.2 KiB
Python
import os
|
|
import argparse
|
|
from sophon_chakcy.get_detector import get_detector
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="YOLOv8目标检测")
|
|
parser.add_argument('--name', type=str, required=True, help='算法名称')
|
|
parser.add_argument('--input', type=str, required=True, help='输入图像或目录路径')
|
|
parser.add_argument('--dev_id', type=int, default=0, help='设备ID')
|
|
parser.add_argument('--conf_thresh', type=float, default=0.25, help='置信度阈值')
|
|
parser.add_argument('--nms_thresh', type=float, default=0.7, help='NMS阈值')
|
|
|
|
args = parser.parse_args()
|
|
|
|
detector_class, model_config = get_detector(args.name)
|
|
|
|
# 初始化检测器
|
|
detector = detector_class(model_config, False, dev_id=args.dev_id, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh)
|
|
|
|
# 根据输入类型进行处理
|
|
if os.path.isfile(args.input):
|
|
# 单张图像
|
|
detector.detect_single_image(args.input)
|
|
|
|
elif os.path.isdir(args.input):
|
|
# 图像目录
|
|
detector.detect_images_in_directory(args.input)
|
|
else:
|
|
print(f"错误: {args.input} 既不是文件也不是目录")
|
|
|
|
if __name__ == "__main__":
|
|
main() |