sun15684712291/aicheckv2-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

完成项目推理模块的接口测试

Browse Source
This commit is contained in:
sunyg
2025-04-23 16:38:55 +08:00
parent 0033746fe1
commit 5b38e91f61
8 changed files with 234 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

85
apps/business/detect/service.py
View File

@@ -21,13 +21,15 @@ async def before_detect(
detect_in: schemas.ProjectDetectLogIn,
detect: models.ProjectDetect,
train: train_models.ProjectTrain,
db: AsyncSession):
db: AsyncSession,
user_id: int):
"""
开始推理
:param detect:
:param detect_in:
:param train:
:param db:
:param user_id:
:return:
"""
# 推理版本
@@ -52,19 +54,33 @@ async def before_detect(
detect_log.pt_url = pt_url
detect_log.folder_url = img_url
detect_log.detect_folder_url = out_url
await crud.ProjectDetectLogDal(db).create_data(detect_log)
detect_log.user_id = user_id
await crud.ProjectDetectLogDal(db).create_model(detect_log)
return detect_log
def run_img_loop(
weights: str,
source: str,
project: str,
name: str,
detect_id: int,
is_gpu: str):
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
# 运行异步函数
loop.run_until_complete(run_detect_img(weights, source, project, name, detect_id, is_gpu))
# 可选: 关闭循环
loop.close()
async def run_detect_img(
weights: str,
source: str,
project: str,
name: str,
log_id: int,
detect_id: int,
db: AsyncSession,
rd: Redis):
is_gpu: str):
"""
执行yolov5的推理
:param weights: 权重文件
@@ -74,7 +90,7 @@ async def run_detect_img(
:param log_id: 日志id
:param detect_id: 推理集合id
:param db: 数据库session
:param rd: Redis
:param is_gpu: 是否gpu加速
:return:
"""
yolo_path = os.file_path(yolo_url, 'detect.py')
@@ -82,10 +98,9 @@ async def run_detect_img(
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"]
is_gpu = rd.get('is_gpu')
# 判断是否存在cuda版本
if is_gpu == 'True':
commend.append("--device", "0")
commend.append("--device=0")
# 启动子进程
with subprocess.Popen(
commend,
@@ -101,40 +116,51 @@ async def run_detect_img(
process.stdout.flush() # 刷新缓存，防止缓存过多造成卡死
if line != '\n':
await room_manager.send_to_room(room, line + '\n')
# 等待进程结束并获取返回码
return_code = process.wait()
if return_code != 0:
await room_manager.send_to_room(room, 'error')
else:
await room_manager.send_to_room(room, 'success')
detect_files = crud.ProjectDetectFileDal(db).get_data(
v_where=[models.ProjectDetectFile.detect_id == detect_id])
detect_log_files = []
for detect_file in detect_files:
detect_log_img = models.ProjectDetectLogFile()
detect_log_img.log_id = log_id
image_url = os.file_path(project, name, detect_file.file_name)
detect_log_img.image_url = image_url
detect_log_img.file_name = detect_file.file_name
detect_log_files.append(detect_log_img)
await crud.ProjectDetectLogFileDal(db).create_datas(detect_log_files)
async def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str, detect_id: int, rd: Redis):
async def update_sql(db: AsyncSession, detect_id: int, log_id: int, project, name):
"""
更新推理集合的状态
"""
detect_dal = crud.ProjectDetectDal(db)
detect = await detect_dal.get_data(detect_id)
detect.detect_version = detect.detect_version + 1
await detect_dal.put_data(data_id=detect_id, data=detect)
detect_files = await crud.ProjectDetectFileDal(db).get_datas(
limit=0,
v_where=[models.ProjectDetectFile.detect_id == detect_id],
v_return_objs=True,
v_return_count=False)
detect_log_files = []
for detect_file in detect_files:
detect_log_img = models.ProjectDetectLogFile()
detect_log_img.log_id = log_id
image_url = os.file_path(project, name, detect_file.file_name)
detect_log_img.file_url = image_url
detect_log_img.file_name = detect_file.file_name
detect_log_files.append(detect_log_img)
await crud.ProjectDetectLogFileDal(db).create_models(detect_log_files)
async def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str, detect_id: int, is_gpu: str):
"""
rtsp 视频流推理
:param detect_id: 训练集的id
:param weights_pt: 权重文件
:param rtsp_url: 视频流地址
:param data: yaml文件
:param rd: Redis :redis
:param is_gpu: 是否启用加速
:return:
"""
room = 'detect_rtsp_' + str(detect_id)
# 选择设备（CPU 或 GPU）
device = select_device('cpu')
is_gpu = rd.get('is_gpu')
# 判断是否存在cuda版本
if is_gpu == 'True':
device = select_device('cuda:0')
@@ -208,19 +234,10 @@ async def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str, detect_id:
break
def run_img_loop(weights: str, source: str, project: str, name: str, log_id: int, detect_id: int, db: AsyncSession):
def run_rtsp_loop(weights_pt: str, rtsp_url: str, data: str, detect_id: int, is_gpu: str):
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
# 运行异步函数
loop.run_until_complete(run_detect_img(weights, source, project, name, log_id, detect_id, db))
# 可选: 关闭循环
loop.close()
def run_rtsp_loop(weights_pt: str, rtsp_url: str, data: str, detect_id: int, rd: Redis):
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
# 运行异步函数
loop.run_until_complete(run_detect_rtsp(weights_pt, rtsp_url, data, detect_id, rd))
loop.run_until_complete(run_detect_rtsp(weights_pt, rtsp_url, data, detect_id, is_gpu))
# 可选: 关闭循环
loop.close()