增加跟踪的一些逻辑代码

app/service/project_detect_service.py

@@ -12,7 +12,8 @@ from ultralytics.utils.plotting import Annotator, colors
 
 from app.model.crud import project_detect_crud as pdc
 from app.model.schemas.project_detect_schemas import ProjectDetectIn, ProjectDetectOut, ProjectDetectLogIn
-from app.model.bussiness_model import ProjectDetect, ProjectDetectImg, ProjectTrain, ProjectDetectLog, ProjectDetectLogImg
+from app.model.bussiness_model import ProjectDetect, ProjectDetectImg, ProjectTrain, ProjectDetectLog, \
+    ProjectDetectLogImg
 from app.util.random_utils import random_str
 from app.config.config_reader import detect_url
 from app.util import os_utils as os
@@ -146,7 +147,8 @@ def run_detect_yolo(detect_in: ProjectDetectLogIn, detect: ProjectDetect, train:
     return detect_log
 
 
-async def run_detect_img(weights: str, source: str, project: str, name: str, log_id: int, detect_id: int, session: Session):
+async def run_detect_img(weights: str, source: str, project: str, name: str, log_id: int, detect_id: int,
+                         session: Session):
     """
     执行yolov5的推理
     :param weights: 权重文件
@@ -161,7 +163,8 @@ async def run_detect_img(weights: str, source: str, project: str, name: str, log
     yolo_path = os.file_path(yolo_url, 'detect.py')
     room = 'detect_' + str(detect_id)
     await room_manager.send_to_room(room, f"AiCheck: 模型训练开始，请稍等。。。\n")
-    commend = ["python", '-u', yolo_path, "--weights", weights, "--source", source, "--name", name, "--project", project, "--save-txt", "--conf-thres", "0.4"]
+    commend = ["python", '-u', yolo_path, "--weights", weights, "--source", source, "--name", name, "--project",
+               project, "--save-txt", "--conf-thres", "0.4"]
     is_gpu = redis_conn.get('is_gpu')
     # 判断是否存在cuda版本
     if is_gpu == 'True':
@@ -231,7 +234,7 @@ async def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str, detect_id:
 
     seen, windows, dt = 0, [], (Profile(device=device), Profile(device=device), Profile(device=device))
 
-    time.sleep(3)# 等待3s，等待websocket进入
+    time.sleep(3)  # 等待3s，等待websocket进入
 
     for path, im, im0s, vid_cap, s in dataset:
         if room_manager.rooms.get(room):
@@ -283,8 +286,5 @@ async def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str, detect_id:
                     frame_data = jpeg.tobytes()
                     await room_manager.send_stream_to_room(room, frame_data)
         else:
+            print(room, '结束推理');
             break
-
-
-
-