部分小问题的修改

apps/business/deepsort/service.py

@@ -3,12 +3,12 @@ from utils.websocket_server import room_manager
 
 import time
 import torch
+from deep_sort.deep_sort import DeepSort
+from deep_sort.utils.draw import draw_boxes
 from utils.yolov5.models.common import DetectMultiBackend
 from utils.yolov5.utils.torch_utils import select_device
 from utils.yolov5.utils.dataloaders import LoadImages, LoadStreams
 from utils.yolov5.utils.general import check_img_size, non_max_suppression, cv2, scale_coords, xyxy2xywh
-from deep_sort.deep_sort import DeepSort
-from deep_sort.utils.draw import draw_boxes
 
 
 def yolov5_to_deepsort_format(pred):
@@ -67,9 +67,16 @@ async def run_deepsort(
 
     model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *img_sz))
 
-    time.sleep(2)  # 等待2s，等待websocket进入
+    time.sleep(3)  # 等待3s，等待websocket进入
+
+    start_time = time.time()
 
     for path, im, im0s, vid_cap, s in dataset:
+        # 检查是否已经超过10分钟（600秒）
+        elapsed_time = time.time() - start_time
+        if elapsed_time > 600:  # 600 seconds = 10 minutes
+            print(room, "已达到最大执行时间，结束推理。")
+            break
         if room_manager.rooms.get(room):
             im0 = im0s[0]
             im = torch.from_numpy(im).to(model.device)

apps/business/detect/crud.py

@@ -6,15 +6,15 @@
 # @IDE            : PyCharm
 # @desc           : 数据访问层
 
-from sqlalchemy.ext.asyncio import AsyncSession
+from utils import status
 from core.crud import DalBase
 from . import schemas, models
-from utils.random_utils import random_str
 from utils import os_utils as os
-from application.settings import detect_url
 from utils.huawei_obs import MyObs
-from utils import status
+from utils.random_utils import random_str
 from core.exception import CustomException
+from application.settings import detect_url
+from sqlalchemy.ext.asyncio import AsyncSession
 
 from fastapi import UploadFile
 

apps/business/detect/service.py

@@ -114,7 +114,7 @@ async def run_detect_img(
         while process.poll() is None:
             line = process.stdout.readline()
             process.stdout.flush()  # 刷新缓存，防止缓存过多造成卡死
-            if line != '\n':
+            if line != '\n' and 'yolo' not in line:
                 await room_manager.send_to_room(room, line + '\n')
         # 等待进程结束并获取返回码
         return_code = process.wait()
@@ -178,7 +178,14 @@ async def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str, detect_id:
 
     time.sleep(3)  # 等待3s，等待websocket进入
 
+    start_time = time.time()
+
     for path, im, im0s, vid_cap, s in dataset:
+        # 检查是否已经超过10分钟（600秒）
+        elapsed_time = time.time() - start_time
+        if elapsed_time > 600:  # 600 seconds = 10 minutes
+            print(room, "已达到最大执行时间，结束推理。")
+            break
         if room_manager.rooms.get(room):
             im = torch.from_numpy(im).to(model.device)
             im = im.half() if model.fp16 else im.float()  # uint8 to fp16/32

apps/business/detect/views.py

@@ -118,7 +118,7 @@ async def run_detect_yolo(
         return ErrorResponse("训练权重不存在")
     file_count = await crud.ProjectDetectFileDal(auth.db).file_count(detect_log_in.detect_id)
     if file_count == 0 and detect.rtsp_url is None:
-        return ErrorResponse("推理集合中没有内容，请先到推理集合中上传图片")
+        return ErrorResponse("推理集合中没有文件，请先到推理集合中上传推理内容")
     is_gpu = await rd.get('is_gpu')
     # 判断一下是单纯的推理项目还是跟踪项目
     project_info = await ProjectInfoDal(auth.db).get_data(data_id=detect.project_id)
@@ -194,7 +194,7 @@ async def logs_download(
     # 检查源文件夹是否存在
     folder_path = os.path.join(detect_log.detect_folder_url, detect_log.detect_version)
     if not os.path.exists(folder_path):
-        return ErrorResponse(msg="推理结果文件夹不存在")
+        return ErrorResponse(msg="推理结果尚未生成，请稍后再下载")
 
     zip_filename = f"{detect_log.detect_name}_{detect_log.detect_version}.zip"
     zip_file_path = osu.zip_folder(folder_path, zip_filename)

apps/business/train/service.py

@@ -8,7 +8,6 @@ from apps.business.project import models as proj_models, crud as proj_crud
 import yaml
 import asyncio
 import subprocess
-from redis.asyncio import Redis
 from sqlalchemy.ext.asyncio import AsyncSession
 
 
@@ -171,7 +170,7 @@ async def run_commend(
         while process.poll() is None:
             line = process.stdout.readline()
             process.stdout.flush()  # 刷新缓存，防止缓存过多造成卡死
-            if line != '\n' and '0%' not in line:
+            if line != '\n' and '0%' not in line and 'yolo' not in line:
                 await room_manager.send_to_room(room, line + '\n')
 
         # 等待进程结束并获取返回码