133 lines
4.3 KiB
Python
133 lines
4.3 KiB
Python
from sqlalchemy.orm import Session
|
||
from typing import List
|
||
from fastapi import UploadFile
|
||
import subprocess
|
||
|
||
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
|
||
from app.util.random_utils import random_str
|
||
from app.config.config_reader import detect_url
|
||
from app.util import os_utils as os
|
||
from app.util import random_utils as ru
|
||
from app.config.config_reader import yolo_url
|
||
|
||
|
||
def add_detect(detect_in: ProjectDetectIn, session: Session):
|
||
"""
|
||
新增训练集合信息,并创建文件夹
|
||
:param detect_in:
|
||
:param session:
|
||
:return:
|
||
"""
|
||
detect = ProjectDetect(**detect_in.dict())
|
||
detect.detect_no = random_str(6)
|
||
detect.detect_version = 0
|
||
url = os.create_folder(detect_url, detect.detect_no, 'images')
|
||
detect.folder_url = url
|
||
detect = pdc.add_detect(detect, session)
|
||
return detect
|
||
|
||
|
||
def check_image_name(detect_id: int, files: List[UploadFile], session: Session):
|
||
"""
|
||
校验上传的文件名称是否重复
|
||
:param detect_id:
|
||
:param files:
|
||
:param session:
|
||
:return:
|
||
"""
|
||
for file in files:
|
||
if not pdc.check_img_name(detect_id, file.filename, session):
|
||
return False, file.filename
|
||
return True, None
|
||
|
||
|
||
def upload_detect_imgs(detect: ProjectDetectOut, files: List[UploadFile], session: Session):
|
||
"""
|
||
上传推理集合的照片，保存原图，并生成缩略图
|
||
:param detect:
|
||
:param files:
|
||
:param session:
|
||
:return:
|
||
"""
|
||
images = []
|
||
for file in files:
|
||
image = ProjectDetectImg()
|
||
image.detect_id = detect.id
|
||
image.file_name = file.filename
|
||
# 保存原图
|
||
path = os.save_images(detect.folder_url, file=file)
|
||
image.image_url = path
|
||
# 生成缩略图
|
||
thumb_image_url = os.file_path(detect.folder_url, 'thumb', ru.random_str(10) + ".jpg")
|
||
os.create_thumbnail(path, thumb_image_url)
|
||
image.thumb_image_url = thumb_image_url
|
||
images.append(image)
|
||
pdc.add_detect_imgs(images, session)
|
||
|
||
|
||
def run_detect_yolo(detect_in: ProjectDetectLogIn, detect: ProjectDetect, train: ProjectTrain, session: Session):
|
||
"""
|
||
开始推理
|
||
:param detect:
|
||
:param detect_in:
|
||
:param train:
|
||
:param session:
|
||
:return:
|
||
"""
|
||
# 推理版本
|
||
version_path = 'v' + str(detect.detect_version + 1)
|
||
|
||
# 权重文件
|
||
pt_url = train.best_pt if detect_in.pt_type == 'best' else train.last_pt
|
||
|
||
# 推理集合文件路径
|
||
img_url = detect.folder_url
|
||
|
||
out_url = os.file_path(detect_url, detect.detect_no, 'detect')
|
||
|
||
# 构建推理记录数据
|
||
detect_log = ProjectDetectLog()
|
||
detect_log.detect_id = detect.id
|
||
detect_log.detect_version = version_path
|
||
detect_log.train_id = train.id
|
||
detect_log.train_version = train.train_version
|
||
detect_log.pt_type = detect_in.pt_type
|
||
detect_log.folder_url = detect.folder_url
|
||
detect_log.detect_folder_url = out_url
|
||
detect_log = pdc.add_detect_log(detect_log, session)
|
||
return detect_log
|
||
|
||
|
||
def run_commend(weights: str, source: str, project: str, name: str,
|
||
detect_log_id: int, session: Session):
|
||
yolo_path = os.file_path(yolo_url, 'detect.py')
|
||
|
||
yield f"stdout: 模型推理开始，请稍等。。。 \n"
|
||
# 启动子进程
|
||
with subprocess.Popen(
|
||
["python", '-u', yolo_path,
|
||
"--weights =" + weights,
|
||
"--source =" + source,
|
||
"--name=" + name,
|
||
"--project=" + project,
|
||
"--view-img"],
|
||
bufsize=1, # bufsize=0时，为不缓存；bufsize=1时，按行缓存；bufsize为其他正整数时，为按照近似该正整数的字节数缓存
|
||
shell=False,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT, # 这里可以显示yolov5训练过程中出现的进度条等信息
|
||
text=True, # 缓存内容为文本，避免后续编码显示问题
|
||
encoding='utf-8',
|
||
) as process:
|
||
while process.poll() is None:
|
||
line = process.stdout.readline()
|
||
process.stdout.flush() # 刷新缓存，防止缓存过多造成卡死
|
||
if line != '\n':
|
||
yield line
|
||
|
||
|
||
|
||
|
||
|