跳过主要内容

Strong Sort (strong deep learning for simple online and realtime tracking, 强化深度学习网络的简单在线与实时跟踪) 操作符

Strong Sort (strong deep learning for simple online and realtime tracking, 强化深度学习网络的简单在线与实时跟踪) 使用深度学习 来唯一标识边界框以便通过图像流跟踪它们。

输入

  • 图像: 高 x 宽 x BGR array。
  • bbox: N_BBOX, X_MIN, X_MAX, Y_MIN, Y_MAX, CONDIDENCE, CLASS, array

输出

  • obstacles_id: x1, x2, y1, y2 track_id, class_id, conf

示例绘制 (跟踪相应于蓝色 # id )

图像

图描述

  - id: yolov5
operator:
outputs:
- obstacles_id
inputs:
image: webcam/image
bbox: yolov5/bbox
python: ../../operators/strong_sort_op.py

图可视化

方法

__init__()

源码
    def __init__(self):
model = StrongSORT(
"osnet_x0_25_msmt17.pt",
torch.device("cuda"),
False,
)
model.model.warmup()
self.model = model
self.frame = []


.on_event(...)

源码

def on_event(
self,
dora_event: dict,
send_output: Callable[[str, bytes], None],
) -> DoraStatus:
if dora_event["type"] == "INPUT":
return self.on_input(dora_event, send_output)
return DoraStatus.CONTINUE


.on_input(...)

源码

def on_input(
self,
dora_input: dict,
send_output: Callable[[str, bytes], None],
) -> DoraStatus:
if dora_input["id"] == "image":
frame = np.array(
dora_input["value"],
np.uint8,
).reshape((IMAGE_HEIGHT, IMAGE_WIDTH, 4))

self.frame = frame[:, :, :3]

elif dora_input["id"] == "obstacles_bbox" and len(self.frame) != 0:
obstacles = np.array(dora_input["value"]).reshape((-1, 6))
if obstacles.shape[0] == 0:
# self.model.increment_ages()
send_output(
"obstacles_id",
pa.array(np.array([]).ravel()),
dora_input["metadata"],
)
return DoraStatus.CONTINUE

# 后加工 yolov5
xywhs = xxyy2xywh(obstacles[:, 0:4])
confs = obstacles[:, 4]
clss = obstacles[:, 5]
with torch.no_grad():
outputs = np.array(
self.model.update(xywhs, confs, clss, self.frame)
).astype("int32")
if len(outputs) != 0:
outputs = outputs[
:, [0, 2, 1, 3, 4, 5, 6]
] # xyxy -> x1, x2, y1, y2 track_id, class_id, conf

send_output(
"obstacles_id",
pa.array(outputs.ravel()),
dora_input["metadata"],
)

return DoraStatus.CONTINUE