部分小问题的修改

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')
 
         # 等待进程结束并获取返回码

main.py

@@ -27,14 +27,13 @@ from starlette.middleware.cors import CORSMiddleware
 def create_app():
     """
     启动项目
-
     docs_url：配置交互文档的路由地址，如果禁用则为None，默认为 /docs
     redoc_url： 配置 Redoc 文档的路由地址，如果禁用则为None，默认为 /redoc
     openapi_url：配置接口文件json数据文件路由地址，如果禁用则为None，默认为/openapi.json
     """
     app = FastAPI(
-        title="aicheckv2-api",
-        description="aicheckv2.0的后台api",
+        title="aicheck-v2-api",
+        description="aicheck-v2的后台api",
         version=settings.VERSION,
         lifespan=lifespan,
         docs_url=None,

utils/websocket_server.py

@@ -9,14 +9,17 @@ class SocketManager:
         if room not in self.rooms:
             self.rooms[room] = []
         self.rooms[room].append(websocket)
+        if room.startswith('pict_op'):
+            websocket.send_text(str(len(self.rooms[room])))
 
     async def remove_from_room(self, room: str, websocket: WebSocket):
         if room in self.rooms:
             self.rooms[room].remove(websocket)
             if len(self.rooms[room]) == 0:
                 del self.rooms[room]
-                if room.startswith('detect_rtsp_'):
-                    print()
+            else:
+                if room.startswith('pict_op'):
+                    websocket.send_text(str(len(self.rooms[room])))
 
     async def broadcast_to_room(self, room: str, message: str, exclude_websocket: WebSocket = None):
         if room in self.rooms: