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/detect/service.py
View File

@ -8,7 +8,7 @@ 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
from utils.yolov5.utils.general import check_img_size, non_max_suppression, cv2, scale_boxes
import time
import torch
@ -97,7 +97,7 @@ async def run_detect_img(
room = 'detect_' + str(detect_id)
await room_manager.send_to_room(room, f"AiCheck: 模型训练开始，请稍等。。。\n")
commend = ["python", '-u', yolo_path, "--weights", weights, "--source", source, "--name", name, "--project",
project, "--save-txt", "--conf-thres", "0.4"]
project, "--save-txt", "--conf-thres", "0.6"]
# 判断是否存在cuda版本
if is_gpu == 'True':
commend.append("--device=0")
@ -176,25 +176,20 @@ async def run_detect_rtsp(weights_pt: str, rtsp_url: str, data: str, detect_id:
model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *img_sz))
dt = (Profile(device=device), Profile(device=device), Profile(device=device))
time.sleep(3) # 等待2s，等待websocket进入
time.sleep(3) # 等待3s，等待websocket进入
for path, im, im0s, vid_cap, s in dataset:
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
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
# Inference
with dt[1]:
pred = model(im, augment=False, visualize=False)
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)
pred = non_max_suppression(pred, 0.25, 0.45, None, False, max_det=1000)
# Process predictions
for i, det in enumerate(pred): # per image