补全yolov5的代码

app/api/business/project_detect_api.py

@@ -120,13 +120,20 @@ def run_detect_yolo(detect_log_in: ProjectDetectLogIn, session: Session = Depend
     detect_img_count = pdc.check_detect_img(detect_log_in.detect_id, session)
     if detect_img_count == 0:
         return rc.response_error("推理集合中没有内容，请先到推理集合中上传图片")
-    detect_log = pds.run_detect_yolo(detect_log_in, detect, train, session)
-    thread_train = threading.Thread(target=run_event_loop, args=(detect_log.pt_url,
-                                             detect_log.folder_url,
-                                             detect_log.detect_folder_url,
-                                             detect_log.detect_version,
-                                             detect_log.id, detect_log.detect_id, session,))
-    thread_train.start()
+    if detect.file_type == 'img' or detect.file_type == 'video':
+        detect_log = pds.run_detect_yolo(detect_log_in, detect, train, session)
+        thread_train = threading.Thread(target=run_event_loop,
+                                        args=(detect_log.pt_url, detect_log.folder_url, detect_log.detect_folder_url,
+                                              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=pds.run_detect_rtsp,
+                                        args=(weights_pt, detect.rtsp_url, train.train_data,))
+        thread_train.start()
     return rc.response_success(msg="执行成功")
 
 

app/model/bussiness_model.py

@@ -103,6 +103,7 @@ class ProjectDetect(DbCommon):
     detect_status: Mapped[int] = mapped_column(Integer)
     file_type: Mapped[str] = mapped_column(String(10))
     folder_url: Mapped[str] = mapped_column(String(255))
+    rtsp_url: Mapped[str] = mapped_column(String(255))
 
 
 class ProjectDetectImg(DbCommon):

app/model/schemas/project_detect_schemas.py

@@ -24,6 +24,7 @@ class ProjectDetectOut(BaseModel):
     detect_version: Optional[int]
     file_type: Optional[str]
     folder_url: Optional[str]
+    rtsp_url: Optional[str]
     create_time: Optional[datetime]
 
     class Config:

app/service/project_detect_service.py

@@ -2,6 +2,14 @@ from sqlalchemy.orm import Session
 from typing import List
 from fastapi import UploadFile
 import subprocess
+from yolov5.models.common import DetectMultiBackend
+from yolov5.utils.torch_utils import select_device
+from yolov5.utils.dataloaders import LoadStreams
+from yolov5.utils.general import check_img_size, Profile, non_max_suppression, cv2, scale_boxes
+import torch
+from pathlib import Path
+from ultralytics.utils.plotting import Annotator, colors, save_one_box
+import platform
 
 from app.model.crud import project_detect_crud as pdc
 from app.model.schemas.project_detect_schemas import ProjectDetectIn, ProjectDetectOut, ProjectDetectLogIn
@@ -172,8 +180,96 @@ async def run_commend(weights: str, source: str, project: str, name: str,
             pdc.add_detect_imgs(detect_log_imgs, session)
 
 
-def run_detect_rtsp():
-    return None
+def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str):
+    """
+    rtsp 视频流推理
+    :param weights_pt: 权重文件
+    :param rtsp_url: 视频流地址
+    :param data: yaml文件
+    :return:
+    """
+    # 选择设备（CPU 或 GPU）
+    device = select_device('cpu')
+
+    # 加载模型
+    model = DetectMultiBackend(weights_pt, device=device, dnn=False, data=data, fp16=False)
+
+    stride, names, pt = model.stride, model.names, model.pt
+    imgsz = check_img_size((640, 640), s=stride)  # check image size
+
+    dataset = LoadStreams(rtsp_url, img_size=imgsz, stride=stride, auto=pt, vid_stride=1)
+    bs = len(dataset)
+
+    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz))
+
+    seen, windows, dt = 0, [], (Profile(device=device), Profile(device=device), Profile(device=device))
+
+    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)
+
+        # 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
+            seen += 1
+            p, im0, frame = path[i], im0s[i].copy(), dataset.count
+            s += f"{i}: "
+
+            p = Path(p)  # to Path
+            s += "{:g}x{:g} ".format(*im.shape[2:])  # print string
+            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
+            imc = im0.copy() if False else im0  # for save_crop
+            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()
+
+                # Print results
+                for c in det[:, 5].unique():
+                    n = (det[:, 5] == c).sum()  # detections per class
+                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
+
+                # Write results
+                for *xyxy, conf, cls in reversed(det):
+                    c = int(cls)  # integer class
+                    label = names[c] if False else f"{names[c]}"
+                    confidence = float(conf)
+                    confidence_str = f"{confidence:.2f}"
+
+                    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()
+            if platform.system() == "Linux" and p not in windows:
+                windows.append(p)
+                cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO)  # allow window resize (Linux)
+                cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])
+            cv2.imshow(str(p), im0)
+            cv2.waitKey(1)  # 1 millisecond