aicheckv2/app/service/project_detect_service.py

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, ProjectDetectLogImg
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
from app.websocket.web_socket_server import room_manager


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
    detect.detect_status = '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 del_detect_img(detect_img_id: int, session: Session):
    """
    删除训练集合图片
    :param detect_img_id:
    :param session:
    :return:
    """
    detect_img = session.query(ProjectDetectImg).filter_by(id=detect_img_id).first()
    if detect_img is None:
        return 0
    os.delete_file_if_exists(detect_img.image_url, detect_img.thumb_image_url)
    session.delete(detect_img)
    session.commit()
    return 1


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_name = detect.detect_name
    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.pt_url = pt_url
    detect_log.folder_url = img_url
    detect_log.detect_folder_url = out_url
    detect_log = pdc.add_detect_log(detect_log, session)
    return detect_log


async def run_commend(weights: str, source: str, project: str, name: str,
                log_id: int, detect_id: int, session: Session):
    """
    执行yolov5的推理
    :param weights: 权重文件
    :param source: 图片所在文件
    :param project: 推理完成的文件位置
    :param name: 版本名称
    :param log_id: 日志id
    :param detect_id: 推理集合id
    :param session:
    :return:
    """
    yolo_path = os.file_path(yolo_url, 'detect.py')
    room = 'detect_' + str(detect_id)
    await room_manager.send_to_room(room, f"stdout: 模型训练开始，请稍等。。。\n")
    commend = ["python", '-u', yolo_path, "--weights", weights, "--source", source, "--name", name, "--project", project, "--save-txt"]
    # 启动子进程
    with subprocess.Popen(
            commend,
            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':
                await room_manager.send_to_room(room, line + '\n')

        # 等待进程结束并获取返回码
        return_code = process.wait()
        if return_code != 0:
            pdc.update_detect_status(detect_id, -1, session)
        else:
            await room_manager.send_to_room(room, 'success')
            pdc.update_detect_status(detect_id, 2, session)
            detect_imgs = pdc.get_img_list(detect_id, session)
            detect_log_imgs = []
            for detect_img in detect_imgs:
                detect_log_img = ProjectDetectLogImg()
                detect_log_img.log_id = log_id
                image_url = os.file_path(project, name, detect_img.file_name)
                detect_log_img.image_url = image_url
                detect_log_img.file_name = detect_img.file_name
                detect_log_imgs.append(detect_log_img)
            pdc.add_detect_imgs(detect_log_imgs, session)