完成推理模块的主体功能

app/service/project_detect_service.py

@@ -5,7 +5,7 @@ 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.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
@@ -23,6 +23,7 @@ def add_detect(detect_in: ProjectDetectIn, session: Session):
     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)
@@ -67,6 +68,22 @@ def upload_detect_imgs(detect: ProjectDetectOut, files: List[UploadFile], sessio
     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):
     """
     开始推理
@@ -94,25 +111,35 @@ def run_detect_yolo(detect_in: ProjectDetectLogIn, detect: ProjectDetect, train:
     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.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
 
 
 def run_commend(weights: str, source: str, project: str, name: str,
-                detect_log_id: int, session: Session):
+                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')
-
     yield f"stdout: 模型推理开始，请稍等。。。 \n"
     # 启动子进程
     with subprocess.Popen(
             ["python", '-u', yolo_path,
-             "--weights =" + weights,
-             "--source =" + source,
-             "--name=" + name,
-             "--project=" + project,
-             "--view-img"],
+             "--weights", weights,
+             "--source", source,
+             "--name", name,
+             "--project", project],
             bufsize=1,  # bufsize=0时，为不缓存；bufsize=1时，按行缓存；bufsize为其他正整数时，为按照近似该正整数的字节数缓存
             shell=False,
             stdout=subprocess.PIPE,
@@ -126,6 +153,24 @@ def run_commend(weights: str, source: str, project: str, name: str,
             if line != '\n':
                 yield line
 
+        # 等待进程结束并获取返回码
+        return_code = process.wait()
+        if return_code != 0:
+            pdc.update_detect_status(detect_id, -1, session)
+        else:
+            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)
+
+
 
 
 

app/service/project_train_service.py

@@ -212,7 +212,7 @@ def run_commend(data: str, project: str,
             train = ProjectTrain()
             train.project_id = project_id
             train.train_version = name
-            bast_pt_path = os.file_path(project, name, 'weights', 'bast.pt')
+            bast_pt_path = os.file_path(project, name, 'weights', 'best.pt')
             last_pt_path = os.file_path(project, name, 'weights', 'last.pt')
             train.best_pt = bast_pt_path
             train.last_pt = last_pt_path