sophon_ai_project/main.py

import argparse
from sophon_chakcy.get_detector import get_detector
from sophon_chakcy.utils.detect_video_stream import DetectVideoStream
from threading import Thread


def detector_task(args):
    detector_class, model_config = get_detector(args.name, True)
    detector = detector_class(model_config, True, dev_id=args.dev_id, conf_thresh=args.conf_thresh, nms_thresh=args.nms_thresh)
    detect_video_stream = DetectVideoStream([detector], args.input)
    t = Thread(target=detect_video_stream.run_task)
    t.start()
    import time
    time.sleep(5)
    detect_video_stream.stop_task()
    t.join()


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_task(args)


if __name__ == "__main__":
    main()