增加删除推理集合接口

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="执行成功")
 
 

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):
     """
     获取分页数据

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)
 
 
 

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}")