增加删除推理集合接口

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)