5.18.8.1. deepsort_tracker_controller

针对DeepSORT算法，通过处理检测的结果和提取的特征，实现对目标的跟踪。

5.18.8.1.1. __init__

接口形式:

def __init__(max_cosine_distance:float,
        nn_budget:int,
        k_feature_dim:int,
        max_iou_distance:float = 0.7,
        max_age:int = 30,
        n_init:int = 3)

参数说明:

• max_cosine_distance: float

输入参数。用于相似度计算的最大余弦距离阈值。

• nn_budget: int

输入参数。用于最近邻搜索的最大数量限制。

• k_feature_dim: int

输入参数。被检测的目标的特征维度。

• max_iou_distance: float

输入参数。模用于跟踪器中的最大交并比（IoU）距离阈值。

• max_age: int

输入参数。跟踪目标在跟踪器中存在的最大帧数。

• n_init: int

输入参数。跟踪器中的初始化帧数阈值。

5.18.8.1.2. process

处理接口

接口形式1:

def process(detected_objects:list[tuple[int, int, int, int, int, float]],
        feature:sail.Tensor)
        -> list[tuple[int, int, int, int, int, float, int]]

参数说明:

• detected_objects: list(tuple(left, top, right, bottom, class_id, score))

输入参数。检测出的物体框。

• feature: sail.Tensor

输入参数。检测出的物体的特征。

返回值说明:

• tracked_objects: list(tuple(left, top, right, bottom, class_id, score, track_id))

输出参数。被跟踪的物体。

接口形式2:

def process(detected_objects:list[tuple[int, int, int, int, int, float]],
        feature:list[numpy.array])
        -> list[tuple[int, int, int, int, int, float, int]]

参数说明:

• detected_objects: list(tuple(left, top, right, bottom, class_id, score))

输入参数。检测出的物体框。

• feature: list[numpy.array]

输入参数。检测出的物体的特征。

返回值说明:

• tracked_objects: list(tuple(left, top, right, bottom, class_id, score, track_id))

输出参数。被跟踪的物体。

示例代码:

# The example code relies on sophon-demo/sample/YOLOv5/python/yolov5_opencv.py and sophon-demo/sample/DeepSORT/python/deep_sort/deep/feature_extractor.py
import sophon.sail as sail
import cv2
import numpy as np
from python.yolov5_opencv import YOLOv5
from feature_extractor import Extractor
class yolov5_arg:
    def __init__(self, bmodel, dev_id, conf_thresh, nms_thresh):
        self.bmodel = bmodel
        self.dev_id = dev_id
        self.conf_thresh = conf_thresh
        self.nms_thresh = nms_thresh
if __name__ == '__main__':
    input = "data/test_car_person_1080P.mp4"
    bmodel_detector = "models/BM1684X/yolov5s_v6.1_3output_int8_1b.bmodel"
    bmodel_extractor = "models/BM1684X/extractor_int8_1b.bmodel"

    dev_id = 0
    conf = 0.4
    nms = 0.7

    yolov5_args = yolov5_arg(bmodel_detector, dev_id, conf, nms)
    yolov5 = YOLOv5(yolov5_args)
    extractor = Extractor(bmodel_extractor, dev_id)

    cap = cv2.VideoCapture(input)
    img_batch = []

    dstc = sail.deepsort_tracker_controller(max_cosine_distance=0.2, nn_budget=100, k_feature_dim=extractor.output_shape[1], max_iou_distance=0.7, max_age=70, n_init=3)

    track_res_all_numpy = np.array([])

    for i in range(15):
        _, im = cap.read()
        if im is None:
            break
        img_batch.append(im)
        results = yolov5(img_batch)
        det = results[0]

        im_crops = []
        for item in det:
            x1 = int(item[0])
            y1 = int(item[1])
            x2 = int(item[2])
            y2 = int(item[3])
            im_crop = im[y1:y2, x1:x2]
            im_crops.append(im_crop)

        ext_results = extractor(im_crops)

        #  The order of this API and the demo is inconsistent, and the class_id and score are reversed
        det[:, [4,5]] = det[:,[5,4]]
        img_batch.clear()

        det_tuple = [tuple(row) for row in det]

        # left, top, right, bottom, class_id, score, track_id
        track_res_numpy = dstc.process(det_tuple, ext_results)
        track_res_numpy = np.array(track_res_numpy)

        if i == 0:
            track_res_all_numpy = track_res_numpy
        else:
            track_res_all_numpy = np.concatenate((track_res_all_numpy, track_res_numpy), axis=0)

    cap.release()