完成目标追踪的开发和测试

apps/business/detect/service.py

@@ -1,19 +1,19 @@
-from application.settings import yolo_url, detect_url
-from utils.websocket_server import room_manager
 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 apps.business.train import models as train_models
-from utils.yolov5.models.common import DetectMultiBackend
-from utils.yolov5.utils.torch_utils import select_device
 from utils.yolov5.utils.dataloaders import LoadStreams
-from utils.yolov5.utils.general import check_img_size, Profile, non_max_suppression, cv2, scale_boxes
+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, Profile, non_max_suppression, cv2, scale_boxes
 
 import time
 import torch
 import asyncio
 import subprocess
-from redis.asyncio import Redis
 from sqlalchemy.ext.asyncio import AsyncSession
 
 
@@ -169,16 +169,16 @@ async def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str, detect_id:
     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
+    img_sz = check_img_size((640, 640), s=stride)  # check image size
 
-    dataset = LoadStreams(rtsp_url, img_size=imgsz, stride=stride, auto=pt, vid_stride=1)
+    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, *imgsz))
+    model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *img_sz))
 
-    seen, windows, dt = 0, [], (Profile(device=device), Profile(device=device), Profile(device=device))
+    dt = (Profile(device=device), Profile(device=device), Profile(device=device))
 
-    time.sleep(3)  # 等待3s，等待websocket进入
+    time.sleep(3)  # 等待2s，等待websocket进入
 
     for path, im, im0s, vid_cap, s in dataset:
         if room_manager.rooms.get(room):
@@ -188,25 +188,13 @@ async def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str, detect_id:
                 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)
+                pred = model(im, augment=False, visualize=False)
             # NMS
             with dt[2]:
-                pred = non_max_suppression(pred, 0.45, 0.45, None, False, max_det=1000)
+                pred = non_max_suppression(pred, 0.25, 0.45, None, False, max_det=1000)
 
             # Process predictions
             for i, det in enumerate(pred):  # per image
@@ -238,6 +226,41 @@ def run_rtsp_loop(weights_pt: str, rtsp_url: str, data: str, detect_id: int, is_
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
     # 运行异步函数
-    loop.run_until_complete(run_detect_rtsp(weights_pt, rtsp_url, data, detect_id, is_gpu))
+    loop.run_until_complete(
+        run_detect_rtsp(
+            weights_pt,
+            rtsp_url,
+            data,
+            detect_id,
+            is_gpu
+        )
+    )
+    # 可选: 关闭循环
+    loop.close()
+
+
+def run_deepsort_loop(
+        detect_id: int,
+        weights_pt: str,
+        data: str,
+        idx_to_class: {},
+        sort_type: str = 'video',
+        video_path: str = None,
+        rtsp_url: str = None
+):
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    # 运行异步函数
+    loop.run_until_complete(
+        deepsort_service.run_deepsort(
+            detect_id,
+            weights_pt,
+            data,
+            idx_to_class,
+            sort_type,
+            video_path,
+            rtsp_url
+        )
+    )
     # 可选: 关闭循环
     loop.close()