提高版本到yolov11
This commit is contained in:
@ -5,10 +5,6 @@ import time
|
||||
import torch
|
||||
from deep_sort.deep_sort import DeepSort
|
||||
from deep_sort.utils.draw import draw_boxes
|
||||
from utils.yolov5.models.common import DetectMultiBackend
|
||||
from utils.yolov5.utils.torch_utils import select_device
|
||||
from utils.yolov5.utils.dataloaders import LoadImages, LoadStreams
|
||||
from utils.yolov5.utils.general import check_img_size, non_max_suppression, cv2, scale_coords, xyxy2xywh
|
||||
|
||||
|
||||
def yolov5_to_deepsort_format(pred):
|
||||
@ -17,11 +13,11 @@ def yolov5_to_deepsort_format(pred):
|
||||
:param pred: YOLOv5的预测结果
|
||||
:return: 转换后的bbox_xywh, confs, class_ids
|
||||
"""
|
||||
pred[:, :4] = xyxy2xywh(pred[:, :4])
|
||||
xywh = pred[:, :4].cpu().numpy()
|
||||
conf = pred[:, 4].cpu().numpy()
|
||||
cls = pred[:, 5].cpu().numpy()
|
||||
return xywh, conf, cls
|
||||
# pred[:, :4] = xyxy2xywh(pred[:, :4])
|
||||
# xywh = pred[:, :4].cpu().numpy()
|
||||
# conf = pred[:, 4].cpu().numpy()
|
||||
# cls = pred[:, 5].cpu().numpy()
|
||||
# return xywh, conf, cls
|
||||
|
||||
|
||||
async def run_deepsort(
|
||||
@ -37,99 +33,99 @@ async def run_deepsort(
|
||||
deep_sort追踪，先经过yolov5对目标进行识别，
|
||||
再调用deepsort对目标进行追踪
|
||||
"""
|
||||
room = 'deep_sort_' + str(detect_id)
|
||||
|
||||
room_count = 'deep_sort_count_' + str(detect_id)
|
||||
|
||||
# 选择设备（CPU 或 GPU）
|
||||
device = select_device('cuda:0')
|
||||
|
||||
model = DetectMultiBackend(weights_pt, device=device, dnn=False, data=data, fp16=False)
|
||||
|
||||
deepsort = DeepSort(
|
||||
model_path="deep_sort/deep/checkpoint/ckpt.t7", # ReID 模型路径
|
||||
max_dist=0.2, # 外观特征匹配阈值（越小越严格）
|
||||
max_iou_distance=0.7, # 最大IoU距离阈值
|
||||
max_age=70, # 目标最大存活帧数（未匹配时保留的帧数）
|
||||
n_init=3 # 初始确认帧数（连续匹配到n_init次后确认跟踪）
|
||||
)
|
||||
stride, names, pt = model.stride, model.names, model.pt
|
||||
img_sz = check_img_size((640, 640), s=stride) # check image size
|
||||
if sort_type == 'video':
|
||||
dataset = LoadImages(video_path, img_size=img_sz, stride=stride, auto=pt, vid_stride=1)
|
||||
else:
|
||||
dataset = LoadStreams(rtsp_url, img_size=img_sz, stride=stride, auto=pt, vid_stride=1)
|
||||
bs = len(dataset)
|
||||
|
||||
count_result = {}
|
||||
for key in idx_to_class.keys():
|
||||
count_result[key] = set()
|
||||
|
||||
model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *img_sz))
|
||||
|
||||
time.sleep(3) # 等待3s，等待websocket进入
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
for path, im, im0s, vid_cap, s in dataset:
|
||||
# 检查是否已经超过10分钟（600秒）
|
||||
elapsed_time = time.time() - start_time
|
||||
if elapsed_time > 600: # 600 seconds = 10 minutes
|
||||
print(room, "已达到最大执行时间，结束推理。")
|
||||
break
|
||||
if room_manager.rooms.get(room):
|
||||
im0 = im0s[0]
|
||||
im = torch.from_numpy(im).to(model.device)
|
||||
im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
|
||||
im /= 255 # 0 - 255 to 0.0 - 1.0
|
||||
if len(im.shape) == 3:
|
||||
im = im[None] # expand for batch dim
|
||||
pred = model(im, augment=False, visualize=False)
|
||||
# NMS
|
||||
pred = non_max_suppression(pred, 0.25, 0.45, None, False, max_det=1000)[0]
|
||||
|
||||
pred[:, :4] = scale_coords(im.shape[2:], pred[:, :4], im0.shape).round()
|
||||
|
||||
# 使用YOLOv5进行检测后得到的pred
|
||||
bbox_xywh, cls_conf, cls_ids = yolov5_to_deepsort_format(pred)
|
||||
|
||||
mask = cls_ids == 0
|
||||
|
||||
bbox_xywh = bbox_xywh[mask]
|
||||
bbox_xywh[:, 2:] *= 1.2
|
||||
cls_conf = cls_conf[mask]
|
||||
cls_ids = cls_ids[mask]
|
||||
|
||||
# 调用Deep SORT更新方法
|
||||
outputs, _ = deepsort.update(bbox_xywh, cls_conf, cls_ids, im0)
|
||||
|
||||
if len(outputs) > 0:
|
||||
bbox_xyxy = outputs[:, :4]
|
||||
identities = outputs[:, -1]
|
||||
cls = outputs[:, -2]
|
||||
names = [idx_to_class[str(label)] for label in cls]
|
||||
# 开始画框
|
||||
ori_img = draw_boxes(im0, bbox_xyxy, names, identities, None)
|
||||
|
||||
# 获取所有被确认过的追踪目标
|
||||
active_tracks = [
|
||||
track for track in deepsort.tracker.tracks
|
||||
if track.is_confirmed()
|
||||
]
|
||||
|
||||
for tark in active_tracks:
|
||||
class_id = str(tark.cls)
|
||||
count_result[class_id].add(tark.track_id)
|
||||
# 对应每个label进行计数
|
||||
result = {}
|
||||
for key in count_result.keys():
|
||||
result[idx_to_class[key]] = len(count_result[key])
|
||||
# 将帧编码为 JPEG
|
||||
ret, jpeg = cv2.imencode('.jpg', ori_img)
|
||||
if ret:
|
||||
jpeg_bytes = jpeg.tobytes()
|
||||
await room_manager.send_stream_to_room(room, jpeg_bytes)
|
||||
await room_manager.send_to_room(room_count, str(result))
|
||||
else:
|
||||
print(room, '结束追踪')
|
||||
break
|
||||
# room = 'deep_sort_' + str(detect_id)
|
||||
#
|
||||
# room_count = 'deep_sort_count_' + str(detect_id)
|
||||
#
|
||||
# # 选择设备（CPU 或 GPU）
|
||||
# device = select_device('cuda:0')
|
||||
#
|
||||
# model = DetectMultiBackend(weights_pt, device=device, dnn=False, data=data, fp16=False)
|
||||
#
|
||||
# deepsort = DeepSort(
|
||||
# model_path="deep_sort/deep/checkpoint/ckpt.t7", # ReID 模型路径
|
||||
# max_dist=0.2, # 外观特征匹配阈值（越小越严格）
|
||||
# max_iou_distance=0.7, # 最大IoU距离阈值
|
||||
# max_age=70, # 目标最大存活帧数（未匹配时保留的帧数）
|
||||
# n_init=3 # 初始确认帧数（连续匹配到n_init次后确认跟踪）
|
||||
# )
|
||||
# stride, names, pt = model.stride, model.names, model.pt
|
||||
# img_sz = check_img_size((640, 640), s=stride) # check image size
|
||||
# if sort_type == 'video':
|
||||
# dataset = LoadImages(video_path, img_size=img_sz, stride=stride, auto=pt, vid_stride=1)
|
||||
# else:
|
||||
# dataset = LoadStreams(rtsp_url, img_size=img_sz, stride=stride, auto=pt, vid_stride=1)
|
||||
# bs = len(dataset)
|
||||
#
|
||||
# count_result = {}
|
||||
# for key in idx_to_class.keys():
|
||||
# count_result[key] = set()
|
||||
#
|
||||
# model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *img_sz))
|
||||
#
|
||||
# time.sleep(3) # 等待3s，等待websocket进入
|
||||
#
|
||||
# start_time = time.time()
|
||||
#
|
||||
# for path, im, im0s, vid_cap, s in dataset:
|
||||
# # 检查是否已经超过10分钟（600秒）
|
||||
# elapsed_time = time.time() - start_time
|
||||
# if elapsed_time > 600: # 600 seconds = 10 minutes
|
||||
# print(room, "已达到最大执行时间，结束推理。")
|
||||
# break
|
||||
# if room_manager.rooms.get(room):
|
||||
# im0 = im0s[0]
|
||||
# im = torch.from_numpy(im).to(model.device)
|
||||
# im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
|
||||
# im /= 255 # 0 - 255 to 0.0 - 1.0
|
||||
# if len(im.shape) == 3:
|
||||
# im = im[None] # expand for batch dim
|
||||
# pred = model(im, augment=False, visualize=False)
|
||||
# # NMS
|
||||
# pred = non_max_suppression(pred, 0.25, 0.45, None, False, max_det=1000)[0]
|
||||
#
|
||||
# pred[:, :4] = scale_coords(im.shape[2:], pred[:, :4], im0.shape).round()
|
||||
#
|
||||
# # 使用YOLOv5进行检测后得到的pred
|
||||
# bbox_xywh, cls_conf, cls_ids = yolov5_to_deepsort_format(pred)
|
||||
#
|
||||
# mask = cls_ids == 0
|
||||
#
|
||||
# bbox_xywh = bbox_xywh[mask]
|
||||
# bbox_xywh[:, 2:] *= 1.2
|
||||
# cls_conf = cls_conf[mask]
|
||||
# cls_ids = cls_ids[mask]
|
||||
#
|
||||
# # 调用Deep SORT更新方法
|
||||
# outputs, _ = deepsort.update(bbox_xywh, cls_conf, cls_ids, im0)
|
||||
#
|
||||
# if len(outputs) > 0:
|
||||
# bbox_xyxy = outputs[:, :4]
|
||||
# identities = outputs[:, -1]
|
||||
# cls = outputs[:, -2]
|
||||
# names = [idx_to_class[str(label)] for label in cls]
|
||||
# # 开始画框
|
||||
# ori_img = draw_boxes(im0, bbox_xyxy, names, identities, None)
|
||||
#
|
||||
# # 获取所有被确认过的追踪目标
|
||||
# active_tracks = [
|
||||
# track for track in deepsort.tracker.tracks
|
||||
# if track.is_confirmed()
|
||||
# ]
|
||||
#
|
||||
# for tark in active_tracks:
|
||||
# class_id = str(tark.cls)
|
||||
# count_result[class_id].add(tark.track_id)
|
||||
# # 对应每个label进行计数
|
||||
# result = {}
|
||||
# for key in count_result.keys():
|
||||
# result[idx_to_class[key]] = len(count_result[key])
|
||||
# # 将帧编码为 JPEG
|
||||
# ret, jpeg = cv2.imencode('.jpg', ori_img)
|
||||
# if ret:
|
||||
# jpeg_bytes = jpeg.tobytes()
|
||||
# await room_manager.send_stream_to_room(room, jpeg_bytes)
|
||||
# await room_manager.send_to_room(room_count, str(result))
|
||||
# else:
|
||||
# print(room, '结束追踪')
|
||||
# break
|
||||
|
||||
@ -1,19 +1,11 @@
|
||||
from algo import YoloModel
|
||||
from utils import os_utils as os
|
||||
from . import models, crud, schemas
|
||||
from utils.websocket_server import room_manager
|
||||
from application.settings import yolo_url, detect_url
|
||||
from application.settings import detect_url
|
||||
from apps.business.train import models as train_models
|
||||
from utils.yolov5.utils.dataloaders import LoadStreams
|
||||
from utils.yolov5.utils.torch_utils import select_device
|
||||
from ultralytics.utils.plotting import Annotator, colors
|
||||
from utils.yolov5.models.common import DetectMultiBackend
|
||||
from apps.business.deepsort import service as deepsort_service
|
||||
from utils.yolov5.utils.general import check_img_size, non_max_suppression, cv2, scale_boxes
|
||||
|
||||
import time
|
||||
import torch
|
||||
import asyncio
|
||||
import subprocess
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
|
||||
@ -69,59 +61,30 @@ def run_img_loop(
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
# 运行异步函数
|
||||
loop.run_until_complete(run_detect_img(weights, source, project, name, detect_id, is_gpu))
|
||||
loop.run_until_complete(run_detect_folder(weights, source, project, name, detect_id, is_gpu))
|
||||
# 可选: 关闭循环
|
||||
loop.close()
|
||||
|
||||
|
||||
async def run_detect_img(
|
||||
def run_detect_folder(
|
||||
weights: str,
|
||||
source: str,
|
||||
project: str,
|
||||
name: str,
|
||||
detect_id: int,
|
||||
is_gpu: str):
|
||||
name: str):
|
||||
"""
|
||||
执行yolov5的推理
|
||||
:param weights: 权重文件
|
||||
:param source: 图片所在文件
|
||||
:param project: 推理完成的文件位置
|
||||
:param name: 版本名称
|
||||
:param log_id: 日志id
|
||||
:param detect_id: 推理集合id
|
||||
:param db: 数据库session
|
||||
:param is_gpu: 是否gpu加速
|
||||
:return:
|
||||
"""
|
||||
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", "--conf-thres", "0.6"]
|
||||
# 判断是否存在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 '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')
|
||||
model = YoloModel(weights)
|
||||
model.predict_folder(
|
||||
source=source,
|
||||
project=project,
|
||||
name=name
|
||||
)
|
||||
|
||||
|
||||
async def update_sql(db: AsyncSession, detect_id: int, log_id: int, project, name):
|
||||
@ -158,70 +121,70 @@ async def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str, detect_id:
|
||||
:param is_gpu: 是否启用加速
|
||||
:return:
|
||||
"""
|
||||
room = 'detect_rtsp_' + str(detect_id)
|
||||
# 选择设备（CPU 或 GPU）
|
||||
device = select_device('cpu')
|
||||
# 判断是否存在cuda版本
|
||||
if is_gpu == 'True':
|
||||
device = select_device('cuda:0')
|
||||
|
||||
# 加载模型
|
||||
model = DetectMultiBackend(weights_pt, device=device, dnn=False, data=data, fp16=False)
|
||||
|
||||
stride, names, pt = model.stride, model.names, model.pt
|
||||
img_sz = check_img_size((640, 640), s=stride) # check image size
|
||||
|
||||
dataset = LoadStreams(rtsp_url, img_size=img_sz, stride=stride, auto=pt, vid_stride=1)
|
||||
bs = len(dataset)
|
||||
|
||||
model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *img_sz))
|
||||
|
||||
time.sleep(3) # 等待3s，等待websocket进入
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
for path, im, im0s, vid_cap, s in dataset:
|
||||
# 检查是否已经超过10分钟（600秒）
|
||||
elapsed_time = time.time() - start_time
|
||||
if elapsed_time > 600: # 600 seconds = 10 minutes
|
||||
print(room, "已达到最大执行时间，结束推理。")
|
||||
break
|
||||
if room_manager.rooms.get(room):
|
||||
im = torch.from_numpy(im).to(model.device)
|
||||
im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
|
||||
im /= 255 # 0 - 255 to 0.0 - 1.0
|
||||
if len(im.shape) == 3:
|
||||
im = im[None] # expand for batch dim
|
||||
|
||||
# Inference
|
||||
pred = model(im, augment=False, visualize=False)
|
||||
# NMS
|
||||
pred = non_max_suppression(pred, 0.25, 0.45, None, False, max_det=1000)
|
||||
|
||||
# Process predictions
|
||||
for i, det in enumerate(pred): # per image
|
||||
p, im0, frame = path[i], im0s[i].copy(), dataset.count
|
||||
annotator = Annotator(im0, line_width=3, example=str(names))
|
||||
if len(det):
|
||||
# Rescale boxes from img_size to im0 size
|
||||
det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
|
||||
|
||||
# Write results
|
||||
for *xyxy, conf, cls in reversed(det):
|
||||
c = int(cls) # integer class
|
||||
label = None if False else (names[c] if False else f"{names[c]} {conf:.2f}")
|
||||
annotator.box_label(xyxy, label, color=colors(c, True))
|
||||
|
||||
# Stream results
|
||||
im0 = annotator.result()
|
||||
# 将帧编码为 JPEG
|
||||
ret, jpeg = cv2.imencode('.jpg', im0)
|
||||
if ret:
|
||||
frame_data = jpeg.tobytes()
|
||||
await room_manager.send_stream_to_room(room, frame_data)
|
||||
else:
|
||||
print(room, '结束推理')
|
||||
break
|
||||
# room = 'detect_rtsp_' + str(detect_id)
|
||||
# # 选择设备（CPU 或 GPU）
|
||||
# device = select_device('cpu')
|
||||
# # 判断是否存在cuda版本
|
||||
# if is_gpu == 'True':
|
||||
# device = select_device('cuda:0')
|
||||
#
|
||||
# # 加载模型
|
||||
# model = DetectMultiBackend(weights_pt, device=device, dnn=False, data=data, fp16=False)
|
||||
#
|
||||
# stride, names, pt = model.stride, model.names, model.pt
|
||||
# img_sz = check_img_size((640, 640), s=stride) # check image size
|
||||
#
|
||||
# dataset = LoadStreams(rtsp_url, img_size=img_sz, stride=stride, auto=pt, vid_stride=1)
|
||||
# bs = len(dataset)
|
||||
#
|
||||
# model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *img_sz))
|
||||
#
|
||||
# time.sleep(3) # 等待3s，等待websocket进入
|
||||
#
|
||||
# start_time = time.time()
|
||||
#
|
||||
# for path, im, im0s, vid_cap, s in dataset:
|
||||
# # 检查是否已经超过10分钟（600秒）
|
||||
# elapsed_time = time.time() - start_time
|
||||
# if elapsed_time > 600: # 600 seconds = 10 minutes
|
||||
# print(room, "已达到最大执行时间，结束推理。")
|
||||
# break
|
||||
# if room_manager.rooms.get(room):
|
||||
# im = torch.from_numpy(im).to(model.device)
|
||||
# im = im.half() if model.fp16 else im.float() # uint8 to fp16/32
|
||||
# im /= 255 # 0 - 255 to 0.0 - 1.0
|
||||
# if len(im.shape) == 3:
|
||||
# im = im[None] # expand for batch dim
|
||||
#
|
||||
# # Inference
|
||||
# pred = model(im, augment=False, visualize=False)
|
||||
# # NMS
|
||||
# pred = non_max_suppression(pred, 0.25, 0.45, None, False, max_det=1000)
|
||||
#
|
||||
# # Process predictions
|
||||
# for i, det in enumerate(pred): # per image
|
||||
# p, im0, frame = path[i], im0s[i].copy(), dataset.count
|
||||
# annotator = Annotator(im0, line_width=3, example=str(names))
|
||||
# if len(det):
|
||||
# # Rescale boxes from img_size to im0 size
|
||||
# det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
|
||||
#
|
||||
# # Write results
|
||||
# for *xyxy, conf, cls in reversed(det):
|
||||
# c = int(cls) # integer class
|
||||
# label = None if False else (names[c] if False else f"{names[c]} {conf:.2f}")
|
||||
# annotator.box_label(xyxy, label, color=colors(c, True))
|
||||
#
|
||||
# # Stream results
|
||||
# im0 = annotator.result()
|
||||
# # 将帧编码为 JPEG
|
||||
# ret, jpeg = cv2.imencode('.jpg', im0)
|
||||
# if ret:
|
||||
# frame_data = jpeg.tobytes()
|
||||
# await room_manager.send_stream_to_room(room, frame_data)
|
||||
# else:
|
||||
# print(room, '结束推理')
|
||||
# break
|
||||
|
||||
|
||||
def run_rtsp_loop(weights_pt: str, rtsp_url: str, data: str, detect_id: int, is_gpu: str):
|
||||
|
||||
@ -8,7 +8,6 @@
|
||||
|
||||
from utils import os_utils as osu
|
||||
from core.dependencies import IdList
|
||||
from core.database import redis_getter
|
||||
from utils.websocket_server import room_manager
|
||||
from . import schemas, crud, params, service, models
|
||||
from apps.business.train.crud import ProjectTrainDal
|
||||
@ -20,7 +19,6 @@ from utils.response import SuccessResponse, ErrorResponse
|
||||
import os
|
||||
import shutil
|
||||
import threading
|
||||
from redis.asyncio import Redis
|
||||
from fastapi.responses import FileResponse
|
||||
from fastapi import Depends, APIRouter, Form, UploadFile, BackgroundTasks
|
||||
|
||||
@ -106,8 +104,7 @@ async def delete_file(
|
||||
@app.post("/start", summary="开始推理")
|
||||
async def run_detect_yolo(
|
||||
detect_log_in: schemas.ProjectDetectLogIn,
|
||||
auth: Auth = Depends(AllUserAuth()),
|
||||
rd: Redis = Depends(redis_getter)):
|
||||
auth: Auth = Depends(AllUserAuth())):
|
||||
detect_dal = crud.ProjectDetectDal(auth.db)
|
||||
train_dal = ProjectTrainDal(auth.db)
|
||||
detect = await detect_dal.get_data(detect_log_in.detect_id)
|
||||
@ -119,21 +116,29 @@ async def run_detect_yolo(
|
||||
file_count = await crud.ProjectDetectFileDal(auth.db).file_count(detect_log_in.detect_id)
|
||||
if file_count == 0 and detect.rtsp_url is None:
|
||||
return ErrorResponse("推理集合中没有文件，请先到推理集合中上传推理内容")
|
||||
is_gpu = await rd.get('is_gpu')
|
||||
# 判断一下是单纯的推理项目还是跟踪项目
|
||||
project_info = await ProjectInfoDal(auth.db).get_data(data_id=detect.project_id)
|
||||
if project_info.type_code == 'yolo':
|
||||
if detect.file_type == 'img' or detect.file_type == 'video':
|
||||
detect_log = await service.before_detect(detect_log_in, detect, train, auth.db, auth.user.id)
|
||||
thread_train = threading.Thread(target=service.run_img_loop,
|
||||
args=(detect_log.pt_url, detect_log.folder_url,
|
||||
detect_log.detect_folder_url, detect_log.detect_version,
|
||||
detect_log.detect_id, is_gpu))
|
||||
thread_train = threading.Thread(
|
||||
target=service.run_detect_folder,
|
||||
args=(
|
||||
detect_log.pt_url,
|
||||
detect_log.folder_url,
|
||||
detect_log.detect_folder_url,
|
||||
detect_log.detect_version
|
||||
)
|
||||
)
|
||||
thread_train.start()
|
||||
await service.update_sql(
|
||||
auth.db, detect_log.detect_id,
|
||||
detect_log.id, detect_log.detect_folder_url,
|
||||
detect_log.detect_version)
|
||||
auth.db,
|
||||
detect_log.detect_id,
|
||||
detect_log.id,
|
||||
detect_log.detect_folder_url,
|
||||
detect_log.detect_version
|
||||
)
|
||||
return SuccessResponse(msg="执行成功", data=file_count)
|
||||
elif detect.file_type == 'rtsp':
|
||||
room = 'detect_rtsp_' + str(detect.id)
|
||||
if not room_manager.rooms.get(room):
|
||||
@ -141,9 +146,18 @@ async def run_detect_yolo(
|
||||
weights_pt = train.best_pt
|
||||
else:
|
||||
weights_pt = train.last_pt
|
||||
thread_train = threading.Thread(target=service.run_rtsp_loop,
|
||||
args=(weights_pt, detect.rtsp_url, train.train_data, detect.id, is_gpu,))
|
||||
thread_train = threading.Thread(
|
||||
target=service.run_rtsp_loop,
|
||||
args=(
|
||||
weights_pt,
|
||||
detect.rtsp_url,
|
||||
train.train_data,
|
||||
detect.id,
|
||||
None
|
||||
)
|
||||
)
|
||||
thread_train.start()
|
||||
return SuccessResponse(msg="执行成功")
|
||||
elif project_info.type_code == 'deepsort':
|
||||
room = 'deep_sort_' + str(detect.id)
|
||||
if not room_manager.rooms.get(room):
|
||||
@ -166,7 +180,7 @@ async def run_detect_yolo(
|
||||
args=(detect.id, train.best_pt, train.train_data,
|
||||
idx_to_class, 'video', detect.folder_url, None))
|
||||
threading_main.start()
|
||||
return SuccessResponse(msg="执行成功")
|
||||
return SuccessResponse(msg="执行成功")
|
||||
|
||||
|
||||
###########################################################
|
||||
|
||||
@ -18,7 +18,7 @@ from typing import Optional
|
||||
|
||||
class ProjectTrainIn(BaseModel):
|
||||
project_id: Optional[int] = Field(..., description="项目id")
|
||||
weights_id: Optional[str] = Field(None, description="权重文件")
|
||||
weights_id: Optional[int] = Field(None, description="权重文件")
|
||||
weights_name: Optional[str] = Field(None, description="权重文件名称")
|
||||
epochs: Optional[int] = Field(50, description="训练轮数")
|
||||
patience: Optional[int] = Field(20, description="早停的耐心值")
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -25,8 +25,7 @@ app = APIRouter()
|
||||
@app.post("/start", summary="执行训练")
|
||||
async def run_train(
|
||||
train_in: schemas.ProjectTrainIn,
|
||||
auth: Auth = Depends(AllUserAuth()),
|
||||
rd: Redis = Depends(redis_getter)):
|
||||
auth: Auth = Depends(AllUserAuth())):
|
||||
proj_id = train_in.project_id
|
||||
proj_dal = ProjectInfoDal(auth.db)
|
||||
proj_img_dal = ProjectImageDal(auth.db)
|
||||
@ -48,14 +47,14 @@ async def run_train(
|
||||
if val_label_count > 0:
|
||||
return ErrorResponse("验证图片中存在未标注的图片")
|
||||
data, project, name = await service.before_train(proj_info, auth.db)
|
||||
is_gpu = await rd.get('is_gpu')
|
||||
train_info = None
|
||||
if train_in.weights_id is not None:
|
||||
train_info = await crud.ProjectTrainDal(auth.db).get_data(train_in.weights_id)
|
||||
# 异步执行操作，操作过程通过websocket进行同步
|
||||
thread_train = threading.Thread(
|
||||
target=service.run_event_loop,
|
||||
args=(data, project, name, train_in, proj_id, train_info, is_gpu))
|
||||
target=service.run_commend,
|
||||
args=(data, project, name, train_in.epochs, train_in.patience, train_info)
|
||||
)
|
||||
thread_train.start()
|
||||
await service.add_train(auth.db, proj_id, name, project, data, train_in, auth.user.id)
|
||||
return SuccessResponse(msg="执行成功")
|
||||
|
||||
Reference in New Issue
Block a user