提高版本到yolov11

apps/business/train/service.py

@@ -1,13 +1,12 @@
+from algo import YoloModel
 from utils import os_utils as osu
 from application.settings import *
 from . import schemas, models, crud
-from utils.websocket_server import room_manager
 from apps.business.project import models as proj_models, crud as proj_crud
 
 
 import yaml
 import asyncio
-import subprocess
 from sqlalchemy.ext.asyncio import AsyncSession
 
 
@@ -118,19 +117,16 @@ def run_event_loop(
     # 运行异步函数，开始训练
     loop_run = asyncio.new_event_loop()
     asyncio.set_event_loop(loop_run)
-    loop_run.run_until_complete(run_commend(data, project, name, train_in.epochs, train_in.patience,
-                                            project_id, train_info, is_gup))
+    loop_run.run_until_complete(run_commend(data, project, name, train_in.epochs, train_in.patience, train_info))
 
 
-async def run_commend(
+def run_commend(
         data: str,
         project: str,
         name: str,
         epochs: int,
         patience: int,
-        project_id: int,
-        train_info: models.ProjectTrain,
-        is_gpu: str):
+        train_info: models.ProjectTrain):
     """
     执行训练
     :param data: 训练数据集
@@ -138,47 +134,14 @@ async def run_commend(
     :param name: 实验名称
     :param epochs: 训练轮数
     :param patience: 早停耐心值
-    :param weights: 权重文件
-    :param project_id: 项目id
     :param train_info: 训练信息
-    :param is_gpu: 是否是gpu环境
     :return:
     """
-    yolo_path = osu.file_path(yolo_url, 'train.py')
-    room = 'train_' + str(project_id)
-    await room_manager.send_to_room(room, f"AiCheckV2.0: 模型训练开始，请稍等。。。\n")
-    commend = ["python", '-u', yolo_path, "--data=" + data, "--project=" + project, "--name=" + name,
-               "--epochs=" + str(epochs), "--batch-size=8", "--exist-ok", "--patience=" + str(patience)]
-
-    # 增加权重文件，在之前训练的基础上重新训练
-    if train_info is not None:
-        commend.append("--weights=" + train_info.best_pt)
-
-    # 判断是否存在cuda版本
-    if is_gpu == 'True':
-        commend.append("--device=0")
-    # 启动子进程
-    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' and '0%' not in line and 'yolo' not in line:
-                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')
+    if train_info is None:
+        model = YoloModel()
+    else:
+        model = YoloModel(train_info.best_pt)
+    model.train(data=data, epochs=epochs, project=project, name=name, patience=patience)
 
 
 async def add_train(