deep_sort 提升版本

apps/business/detect/views.py

@@ -11,10 +11,11 @@ from core.dependencies import IdList
 from utils.websocket_server import room_manager
 from . import schemas, crud, params, service, models
 from apps.business.train.crud import ProjectTrainDal
-from apps.business.project.crud import ProjectInfoDal, ProjectLabelDal
 from apps.vadmin.auth.utils.current import AllUserAuth
 from apps.vadmin.auth.utils.validation.auth import Auth
 from utils.response import SuccessResponse, ErrorResponse
+from apps.business.deepsort import service as deep_sort_service
+from apps.business.project.crud import ProjectInfoDal, ProjectLabelDal
 
 import os
 import shutil
@@ -147,13 +148,11 @@ async def run_detect_yolo(
                 else:
                     weights_pt = train.last_pt
                 thread_train = threading.Thread(
-                    target=service.run_rtsp_loop,
+                    target=service.run_detect_rtsp,
                     args=(
                         weights_pt,
                         detect.rtsp_url,
-                        train.train_data,
-                        detect.id,
-                        None
+                        room
                     )
                 )
                 thread_train.start()
@@ -164,22 +163,16 @@ async def run_detect_yolo(
             # 查询项目所属标签，返回两个 id，name一一对应的数组
             label_id_list, label_name_list = await ProjectLabelDal(auth.db).get_label_for_train(project_info.id)
             idx_to_class = {str(i): name for i, name in enumerate(label_name_list)}
-            # if detect_log_in.pt_type == 'best':
-            #     weights_pt = train.best_pt
-            # else:
-            #     weights_pt = train.last_pt
-            if detect.file_type == 'rtsp':
-                threading_main = threading.Thread(
-                    target=service.run_deepsort_loop,
-                    args=(detect.id, train.best_pt, train.train_data,
-                          idx_to_class, 'rtsp', None, detect.rtsp_url))
-                threading_main.start()
-            elif detect.file_type == 'video':
-                threading_main = threading.Thread(
-                    target=service.run_deepsort_loop,
-                    args=(detect.id, train.best_pt, train.train_data,
-                          idx_to_class, 'video', detect.folder_url, None))
-                threading_main.start()
+            threading_main = threading.Thread(
+                target=deep_sort_service.run_deepsort,
+                args=(
+                    detect.id,
+                    train.best_pt,
+                    idx_to_class,
+                    detect.rtsp_url
+                )
+            )
+            threading_main.start()
         return SuccessResponse(msg="执行成功")