diff --git a/app/api/business/project_detect_api.py b/app/api/business/project_detect_api.py index 3d7f304..82f8621 100644 --- a/app/api/business/project_detect_api.py +++ b/app/api/business/project_detect_api.py @@ -2,10 +2,10 @@ import threading import asyncio from typing import List from fastapi import APIRouter, Depends, UploadFile, File, Form -from fastapi.responses import StreamingResponse from fastapi.encoders import jsonable_encoder from sqlalchemy.orm import Session +from app.websocket.web_socket_server import room_manager from app.db.db_session import get_db from app.common import reponse_code as rc from app.model.crud import project_detect_crud as pdc @@ -52,6 +52,18 @@ def add_detect(detect_in: ProjectDetectIn, session: Session = Depends(get_db)): return rc.response_success(msg="新增成功", data=detect.id) +@detect.get("/del_detect/{detect_id}") +def del_detect(detect_id: int, session: Session = Depends(get_db)): + """ + 删除训练集合 + :param detect_id: + :param session: + :return: + """ + pds.del_detect(detect_id, session) + return rc.response_success(msg="删除成功") + + @detect.post("/get_img_list") def get_img_list(detect_img_pager: ProjectDetectImgPager, session: Session = Depends(get_db)): """ @@ -127,13 +139,15 @@ def run_detect_yolo(detect_log_in: ProjectDetectLogIn, session: Session = Depend detect_log.detect_version, detect_log.id, detect_log.detect_id, session,)) thread_train.start() elif detect.file_type == 'rtsp': - if detect_log_in.pt_type == 'best': - weights_pt = train.best_pt - else: - weights_pt = train.last_pt - thread_train = threading.Thread(target=run_rtsp_loop, - args=(weights_pt, detect.rtsp_url, train.train_data, detect.id)) - thread_train.start() + room = 'detect_rtsp_' + str(detect.id) + if not room_manager.rooms.get(room): + if detect_log_in.pt_type == 'best': + weights_pt = train.best_pt + else: + weights_pt = train.last_pt + thread_train = threading.Thread(target=run_rtsp_loop, + args=(weights_pt, detect.rtsp_url, train.train_data, detect.id)) + thread_train.start() return rc.response_success(msg="执行成功") diff --git a/app/model/crud/project_detect_crud.py b/app/model/crud/project_detect_crud.py index d38d6c0..28c6ef9 100644 --- a/app/model/crud/project_detect_crud.py +++ b/app/model/crud/project_detect_crud.py @@ -190,6 +190,17 @@ def get_log_list(detect_id: int, session: Session): return result +def get_logs(detect_id: int, session: Session): + """ + 获取推理记录 + :param detect_id: + :param session: + :return: + """ + query = session.query(ProjectDetectLog).filter_by(detect_id=detect_id).order_by(asc(ProjectDetectLog.id)) + return query.all() + + def get_log_pager(detect_log_pager: ProjectDetectLogPager, session: Session): """ 获取分页数据 diff --git a/app/service/project_detect_service.py b/app/service/project_detect_service.py index 073e2ba..1f6072f 100644 --- a/app/service/project_detect_service.py +++ b/app/service/project_detect_service.py @@ -1,3 +1,5 @@ +import time + from sqlalchemy.orm import Session from typing import List from fastapi import UploadFile @@ -37,6 +39,24 @@ def add_detect(detect_in: ProjectDetectIn, session: Session): return detect +def del_detect(detect_id: int, session: Session): + """ + 删除推理集合和推理记录 + :param detect_id: + :param session: + :return: + """ + folder_url = [] + detect = pdc.get_detect_by_id(detect_id, session) + session.delete(detect) + folder_url.append(detect.folder_url) + detect_logs = pdc.get_logs(detect_id, session) + for log in detect_logs: + folder_url.append(log.detect_folder_url) + os.create_folder(folder_url) + session.commit() + + def check_image_name(detect_id: int, files: List[UploadFile], session: Session): """ 校验上传的文件名称是否重复 @@ -187,7 +207,6 @@ async def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str, detect_id: :return: """ room = 'detect_rtsp_' + str(detect_id) - await room_manager.send_to_room(room, '开始推理rtsp视频流') # 选择设备（CPU 或 GPU） device = select_device('cpu') @@ -204,54 +223,60 @@ async def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str, detect_id: seen, windows, dt = 0, [], (Profile(device=device), Profile(device=device), Profile(device=device)) + time.sleep(3)# 等待3s，等待websocket进入 + for path, im, im0s, vid_cap, s in dataset: - with dt[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 - if model.xml and im.shape[0] > 1: - ims = torch.chunk(im, im.shape[0], 0) + if room_manager.rooms.get(room): + with dt[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 + if model.xml and im.shape[0] > 1: + ims = torch.chunk(im, im.shape[0], 0) - # Inference - with dt[1]: - if model.xml and im.shape[0] > 1: - pred = None - for image in ims: - if pred is None: - pred = model(image, augment=False, visualize=False).unsqueeze(0) - else: - pred = torch.cat((pred, model(image, augment=False, visualize=False).unsqueeze(0)), - dim=0) - pred = [pred, None] - else: - 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) + # Inference + with dt[1]: + if model.xml and im.shape[0] > 1: + pred = None + for image in ims: + if pred is None: + pred = model(image, augment=False, visualize=False).unsqueeze(0) + else: + pred = torch.cat((pred, model(image, augment=False, visualize=False).unsqueeze(0)), + dim=0) + pred = [pred, None] + else: + 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) - # 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() + # 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)) + # 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: + break - # 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) diff --git a/app/util/os_utils.py b/app/util/os_utils.py index 7f5cb01..4b1a9d1 100644 --- a/app/util/os_utils.py +++ b/app/util/os_utils.py @@ -89,3 +89,24 @@ def delete_file_if_exists(*file_paths: str): for path in file_paths: if os.path.exists(path): # 检查文件是否存在 os.remove(path) # 删除文件 + + +def delete_paths(paths): + """ + 删除给定路径数组中的每个路径及其包含的所有内容。 + :param paths: 文件或目录路径的列表 + """ + for path in paths: + if os.path.exists(path): + try: + if os.path.isfile(path) or os.path.islink(path): + # 如果是文件或符号链接，则删除 + os.remove(path) + print(f"Deleted file: {path}") + elif os.path.isdir(path): + # 如果是目录，则递归删除 + shutil.rmtree(path) + except Exception as e: + print(f"路径删除失败 {path}: {e}") + else: + print(f"路径不存在: {path}")