From 86c6669593afdf07b9c81d860927b92a5b8257eb Mon Sep 17 00:00:00 2001 From: sunyugang Date: Tue, 20 May 2025 17:03:14 +0800 Subject: [PATCH] =?UTF-8?q?=E9=83=A8=E5=88=86=E5=B0=8F=E9=97=AE=E9=A2=98?= =?UTF-8?q?=E7=9A=84=E4=BF=AE=E6=94=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/business/deepsort/service.py | 13 ++++++++++--- apps/business/detect/crud.py | 8 ++++---- apps/business/detect/service.py | 9 ++++++++- apps/business/detect/views.py | 4 ++-- apps/business/train/service.py | 3 +-- main.py | 5 ++--- utils/websocket_server.py | 7 +++++-- 7 files changed, 32 insertions(+), 17 deletions(-) diff --git a/apps/business/deepsort/service.py b/apps/business/deepsort/service.py index 0378509..2bc82ae 100644 --- a/apps/business/deepsort/service.py +++ b/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) diff --git a/apps/business/detect/crud.py b/apps/business/detect/crud.py index 44657d1..0a2a6ba 100644 --- a/apps/business/detect/crud.py +++ b/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 diff --git a/apps/business/detect/service.py b/apps/business/detect/service.py index d39b88f..5731b70 100644 --- a/apps/business/detect/service.py +++ b/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 diff --git a/apps/business/detect/views.py b/apps/business/detect/views.py index 87b09af..61fc7d5 100644 --- a/apps/business/detect/views.py +++ b/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) diff --git a/apps/business/train/service.py b/apps/business/train/service.py index b71aaff..5b26263 100644 --- a/apps/business/train/service.py +++ b/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') # 等待进程结束并获取返回码 diff --git a/main.py b/main.py index 0752dae..6862d7c 100644 --- a/main.py +++ b/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, diff --git a/utils/websocket_server.py b/utils/websocket_server.py index 2a19c9c..8f9e231 100644 --- a/utils/websocket_server.py +++ b/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: