ByteTrack多目标跟踪算法使用yolox官方模型进行多类别跟踪

ByteTrack官方提供的demo只能检测人,无法多类别检测。

所以尝试直接下载 yolox 的官方预训练模型作为检测模型进行多目标跟踪。

yolox官方预训练模型下载地址:

https://github.com/Megvii-BaseDetection/YOLOX/

ByteTrack多目标跟踪算法使用yolox官方模型进行多类别跟踪_第1张图片

我下载YOLOX-m模型,放在ByteTrack/pretrained目录下。

编辑./yolox/data/data_augment.py文件,将preproc 函数替换为原 yolox 项目中的 preproc 函数即可。

def preproc(img, input_size, swap=(2, 0, 1)):
    if len(img.shape) == 3:
        padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
    else:
        padded_img = np.ones(input_size, dtype=np.uint8) * 114

    r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
    resized_img = cv2.resize(
        img,
        (int(img.shape[1] * r), int(img.shape[0] * r)),
        interpolation=cv2.INTER_LINEAR,
    ).astype(np.uint8)
    padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img

    padded_img = padded_img.transpose(swap)
    padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
    return padded_img, r

修改demo_track.py文件,将img, ratio = preproc(img, self.test_size, self.rgb_means, self.std)替换为:

img, ratio = preproc(img, self.test_size)

运行查看效果:

python tools/demo_track.py video -f exps/default/yolox_m.py -c pretrained/yolox_m.pth --fp16 --fuse --save_result

改为使用yolox官方提供的模型后,识别帧率居然能高不少!

你可能感兴趣的:(多目标跟踪,目标跟踪,人工智能,计算机视觉)