YOLOv5+双目实现三维跟踪（python）

news2024/11/15 13:02:05

YOLOv5+双目实现三维跟踪（python）

1. 目标跟踪
2. 测距模块
- 2.1 测距原理
- 2.2 添加测距
3. 细节修改（可忽略）
4. 实验效果

相关链接
1. YOLOV5 + 双目测距（python）
2. YOLOV7 + 双目测距（python）
3. YOLOv7+双目实现三维跟踪（python）
4. 具体实现效果已在Bilibili发布，点击跳转

1. 目标跟踪

用yolov5实现跟踪步骤比较简单，去官网下载deepsort源码，这里有个版本对应关系
DeepSort v3.0 ~YOLOv5 v5.0-------------------DeepSort v4.0 ~ YOLOv5 v6.1
后续有机会的话会特意写一下跟踪原理…
在这里插入图片描述

下载完DeepSort之后去YOLO官网下载相应的YOLO版本，然后把下载的YOLO拖进DeepSort文件夹里，并把YOLO文件夹改名为yolov5，接下来把环境装好，然后运行代码 track.py ，此时如果不出问题就完成了普通检测
也可以用终端运行命令python track.py --source 1.mp4 --show-vid --save-vid --yolo_weights yolov5/weights/yolov5s.pt
这里有几个常用知识需要注意的，我直接在以下代码作了注释

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--yolo_weights', type=str, default='yolov5/weights/yolov5s.pt', help='model.pt path')
    parser.add_argument('--deep_sort_weights', type=str, default='deep_sort_pytorch/deep_sort/deep/checkpoint/ckpt.t7', help='ckpt.t7 path')    # file/folder, 0 for webcam
    #parser.add_argument('--source', type=str, default='0', help='source')# 改成0可以调用摄像头
    parser.add_argument('--source', type=str, default='1.mp4', help='source')
    parser.add_argument('--output', type=str, default='output', help='output folder')  # output folder
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
    parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--show-vid', action='store_true', help='display tracking video results')   # 显示检测画面
    parser.add_argument('--save-vid', action='store_true', help='save video tracking results')  # 保存检测后的画面
    parser.add_argument('--save-txt', action='store_true', help='save MOT compliant results to *.txt')
    # class 0 is person, 1 is bycicle, 2 is car... 79 is oven
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class')  # 检测类别
    #parser.add_argument('--classes', nargs='+', default=[0], type=int, help='filter by class')  # default=[0]代表只检测coco数据集里的类别0，即person，同理可换成别的类别

请添加图片描述

2. 测距模块

2.1 测距原理

测距原理详见双目三维测距（python）

2.2 添加测距

接下来调用测距代码到主代码 track.py 文件中，先在代码开头导入库，添加

from stereo import stereoconfig
from stereo.stereo import stereo_40
from stereo.stereo import stereo_threading, MyThread
from yolov5.utils.plots import plot_one_box

我们需要将立体匹配等代码写进跟踪模块里，具体写法在我之前开源的 YOLOv5+双目测距（python）这片文章里已经提及，这里就不再细讲，最后计算得到目标框的中心点坐标和距离对其进行显示，具体如下

for *xyxy, conf, cls in det:
    # to deep sort format
    x_c, y_c, bbox_w, bbox_h = xyxy_to_xywh(*xyxy)
    xywh_obj = [x_c, y_c, bbox_w, bbox_h]
    xywh_bboxs.append(xywh_obj)
    confs.append([conf.item()])

    if (0 < xyxy[2] < 1280):
            x_center = (xyxy[0] + xyxy[2]) / 2
            y_center = (xyxy[1] + xyxy[3]) / 2
            x_0 = int(x_center)
            y_0 = int(y_center)
            if (0 < x_0 < 1280):
                x1 = xyxy[0]
                x2 = xyxy[2]
                y1 = xyxy[1]
                y2 = xyxy[3]

                if (accel_frame % fps_set == 0):
                    t3 = time.time()  # stereo time end
                    thread.join()
                    points_3d = thread.get_result()
                   # gol.set_value('points_3d', points_3d)
                    t4 = time.time()  # stereo time end
                    print(f'{s}Stereo Done. ({t4 - t3:.3f}s)')
                a = points_3d[int(y_0), int(x_0), 0] / 1000
                b = points_3d[int(y_0), int(x_0), 1] / 1000
                c = points_3d[int(y_0), int(x_0), 2] / 1000
                distance = ((a**2+b**2+c**2)**0.5)

                if (distance != 0):  ## Add bbox to image
                    label = f'{names[int(cls)]} {conf:.2f} '
                    text_xy_0 = "*"
                    print('点 (%d, %d) 的 %s 距离左摄像头的相对距离为 %0.2f m' % (x_center, y_center, label, distance))
                    text_dis_avg = "dis:%0.2fm" % distance
                    cv2.rectangle(im0, (int(x1 + (x2 - x1)), int(y1)),(int(x1 + (x2 - x1) + 5 + 100), int(y1 + 12)), colors[int(cls)],-1)  # 画框存三维坐标
                    cv2.putText(im0, text_dis_avg, (int(x1 + (x2 - x1) + 5), int(y1 + 10)),cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 2)

3. 细节修改（可忽略）

下边是一些小细节修改，可以忽略不看
为了实时显示画面，对运行的py文件点击编辑配置，在形参那里输入–view-img --save-txt
在这里插入图片描述
但实时显示画面太大，我们对显示部分做了修改，这部分也可以不要，具体是把代码

if view_img:
      cv2.imshow(str(p), im0)
      cv2.waitKey(1)  # 1 millisecond

替换成

if view_img:
     cv2.namedWindow("Webcam", cv2.WINDOW_NORMAL)
     cv2.resizeWindow("Webcam", 1280, 720)
     cv2.moveWindow("Webcam", 0, 100)
     cv2.imshow("Webcam", im0)
     cv2.waitKey(1)