modify: 多线程管理器

public/main.py

@@ -1,57 +1,34 @@
-import os
-import argparse
-from sophon_chakcy.get_detector import get_detector
+import time
+from sophon_chakcy.utils.task_manager import TaskManager, TaskArgs
 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阈值')
+    task_manager = TaskManager(4)
+    start_time = time.time()
+    print(start_time)
+    for i in range(1):
+        task_name = f"yolo_test_{i}"
+        args = TaskArgs(
+            task_name=task_name,
+            algorithms=["yolo_example", "yolo_example"],
+            input="rtsp://192.168.137.1:8554/stream",
+            is_vfs=True
+        )
+        task_manager.submit(args)
     
-    args = parser.parse_args()
+    # time.sleep(5)
 
-    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)
+    # for i in range(4):
+    #     task_name = f"yolo_test_{i}"
+    #     task_manager.stop_task(task_name)
+        
+    # print(time.time() - start_time)
 
-    # elif os.path.isdir(args.input):
-    #     # 图像目录
-    #     detector.detect_images_in_directory(args.input)
-    # else:
-    #     print(f"错误: {args.input} 既不是文件也不是目录")
 
 
 if __name__ == "__main__":