import os
import argparse
from sophon_chakcy.get_detector import get_detector
from module_bank import PythonToSQLite
from multiprocessing import Process

# 安装导入器
packer = PythonToSQLite("plugins.mbank")
finder = packer.install_importer()

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)
    # 根据输入类型进行处理
    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} 既不是文件也不是目录")


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

    p = Process(target=detector_task,args=(args,))
    p.start()
    p.join()
    
    # 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)
    
    # # 根据输入类型进行处理
    # 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()