优化训练过程

app/service/project_detect_service.py

@@ -161,7 +161,7 @@ async def run_detect_img(weights: str, source: str, project: str, name: str, log
     yolo_path = os.file_path(yolo_url, 'detect.py')
     room = 'detect_' + str(detect_id)
     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"]
+    commend = ["python", '-u', yolo_path, "--weights", weights, "--source", source, "--name", name, "--project", project, "--save-txt", "--conf-thres", "0.4"]
     is_gpu = redis_conn.get('is_gpu')
     # 判断是否存在cuda版本
     if is_gpu == 'True':
@@ -259,7 +259,7 @@ async def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str, detect_id:
                     pred = model(im, augment=False, visualize=False)
             # NMS
             with dt[2]:
-                pred = non_max_suppression(pred, 0.25, 0.45, None, False, max_det=1000)
+                pred = non_max_suppression(pred, 0.45, 0.45, None, False, max_det=1000)
 
             # Process predictions
             for i, det in enumerate(pred):  # per image

app/service/project_train_service.py

@@ -14,12 +14,11 @@ from app.websocket.web_socket_server import room_manager
 from app.util.csv_utils import read_csv
 from app.common.redis_cli import redis_conn
 
-from sqlalchemy.orm import Session
-from typing import List
-from fastapi import UploadFile
 import yaml
 import subprocess
-import torch
+from typing import List
+from fastapi import UploadFile
+from sqlalchemy.orm import Session
 
 
 def add_project(info: ProjectInfoIn, session: Session, user_id: int):

app/util/yolov5/utils/general.py

@@ -513,7 +513,7 @@ def check_font(font=FONT, progress=False):
     font = Path(font)
     file = CONFIG_DIR / font.name
     if not font.exists() and not file.exists():
-        url = f"https://github.com/ultralytics/assets/releases/download/v0.0.0/{font.name}"
+        url = f"https://ultralytics.com/assets/{font.name}"
         LOGGER.info(f"Downloading {url} to {file}...")
         torch.hub.download_url_to_file(url, str(file), progress=progress)