sun15684712291/aicheckv2-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

增加目标追踪的模块

Browse Source
This commit is contained in:
sunyg
2025-05-07 14:17:36 +08:00
parent da1f5a874e
commit 939e29c2d6
14 changed files with 256 additions and 503 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
apps/business/deepsort/service.py
View File

@ -6,7 +6,7 @@ import torch
from utils.yolov5.models.common import DetectMultiBackend
from utils.yolov5.utils.torch_utils import select_device
from utils.yolov5.utils.dataloaders import LoadImages, LoadStreams
from utils.yolov5.utils.general import check_img_size, non_max_suppression, cv2, scale_coords, xyxy2xywh, Profile
from utils.yolov5.utils.general import check_img_size, non_max_suppression, cv2, scale_coords, xyxy2xywh
from deep_sort.deep_sort import DeepSort
from deep_sort.utils.draw import draw_boxes
@ -69,22 +69,17 @@ async def run_deepsort(
time.sleep(2) # 等待2s，等待websocket进入
dt = (Profile(device=device), Profile(device=device), Profile(device=device))
for path, im, im0s, vid_cap, s in dataset:
if room_manager.rooms.get(room):
with dt[0]:
im0 = im0s[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
with dt[1]:
pred = model(im, augment=False, visualize=False)
with dt[2]:
# NMS
pred = non_max_suppression(pred, 0.25, 0.45, None, False, max_det=1000)[0]
im0 = im0s[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
pred = model(im, augment=False, visualize=False)
# NMS
pred = non_max_suppression(pred, 0.25, 0.45, None, False, max_det=1000)[0]
pred[:, :4] = scale_coords(im.shape[2:], pred[:, :4], im0.shape).round()