Python 根据生成的txt对多目标跟踪结果画框可视化显示

Python 根据生成的txt对多目标跟踪结果画框可视化显示

2022.1.3更新

输入为视频文件与多目标跟踪结果生成的txt文件(格式与gt一致)
输出为视频文件
draw_mot.py

import os.path
import numpy as np
import cv2
import argparse

ap = argparse.ArgumentParser()
ap.add_argument("-t", "--input_txt", help="path to input txt", default="F:/ProgramData/object/videos/video2/video2.txt")
ap.add_argument("-v", "--input_video", help="path to input video", default="F:/ProgramData/object/videos/video2/detect-result2.mp4")
args = vars(ap.parse_args())
print(args)

# initialize a list of colors to represent each possible class label
np.random.seed(100)
COLORS = np.random.randint(0, 255, size=(200, 3), dtype="uint8")


def sort_output(txt_path):
    with open(txt_path, 'r') as f:
        list = []
        for line in f:
            list.append(line.strip())

    with open(txt_path, "w") as f:
        for item in sorted(list, key=lambda x: int(str(x).split(',')[0])):
            f.writelines(item)
            f.writelines('\n')
        f.close()


def draw_mot(source_file, bboxs, frame):
    for i in range(bboxs):
        line = source_file.readline()
        staff = line.split(',')
        # print(staff)
        id = staff[1]
        cls = staff[7]
        box = staff[2:6]
        # print(id, box)
        # draw_bbox
        color = [int(c) for c in COLORS[int(id) % len(COLORS)]]
        cv2.rectangle(frame, (int(float(box[0])), int(float(box[1]))),
                      (int(float(box[0])) + int(float(box[2])), int(float(box[1])) + int(float(box[3]))),
                      (color[0], color[1], color[2]), 3)
        # put_text
        cv2.putText(frame, 'ID-' + str(int(id)), (int(float(box[0])), int(float(box[1])-float(27))),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (color[0], color[1], color[2]), 3)
        cv2.putText(frame, 'person', (int(float(box[0])), int(float(box[1])-float(10))),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (color[0], color[1], color[2]), 2)

    return frame


def main():
    path_txt = args["input_txt"]
    path_video = args["input_video"]
    sort_output(path_txt)  # 这是一个对txt文本结果排序的代码，key=frame，根据帧数排序
    draw_txt = path_txt
    cap = cv2.VideoCapture(path_video)

    fps = cap.get(cv2.CAP_PROP_FPS)  # 帧率<每秒中展示多少张图片>
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 获取宽度
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 获取高度
    print(fps, width, height)
    fourcc = cv2.VideoWriter_fourcc(*'MJPG')
    out = cv2.VideoWriter('output.avi', fourcc, fps, (width, height), True)

    source_file = open(draw_txt)
    # 把frame存入列表img_names
    img_names = []
    for line in source_file:
        staff = line.split(',')
        img_name = staff[0]
        img_names.append(img_name)

    # 将每个frame的bbox数目存入字典
    name_dict = {}
    for i in img_names:
        if img_names.count(i):
            name_dict[i] = img_names.count(i)
    print(name_dict)
    source_file.close()

    source_file = open(draw_txt)
    # draw_mot
    if cap.isOpened():
        i = 0
        while True:
            frame_index = [idx for idx in name_dict]
            (flag, frame) = cap.read()  # 读取每一张 flag<读取是否成功> frame<内容>
            if not flag:
                break  # 当获取完最后一帧就结束
            try:
                frame_out = draw_mot(source_file, name_dict[str(frame_index[i])], frame)
                cv2.namedWindow("YOLOv5_DeepSORT", 0)
                cv2.resizeWindow('YOLOv5_DeepSORT', 1024, 768)
                cv2.imshow('YOLOv5_DeepSORT', frame_out)
                out.write(frame_out)
            except IndexError:
                cap.release()
                out.release()
                cv2.destroyAllWindows()

            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    else:
        print('视频打开失败！')


if __name__ == '__main__':
    main()

---

原版：

输入为图片序列与多目标跟踪结果生成的txt文件(格式与gt一致)
输出为图片序列

draw_mot.py

import os.path
import numpy as np
from sort_gt import sort_output
import cv2


def draw_mot(video_id):
    txt_name = 'D:/project/code/result/' + video_id + '.txt'  # txt文本内容
    file_path_img = 'D:/project/datasets/trainData/Multi-object-tracking-trainData/video/Multi-object-tracking-trainData/' + video_id + '/img'  # img图片路径
    # 生成新的文件夹来存储画了bbox的图片
    if not os.path.exists('./pic_draw/' + video_id):
        os.mkdir('./pic_draw/' + video_id)
        print('The ./pic_draw/' + video_id + '  have create!')
    save_file_path = './pic_draw/' + video_id
    sort_output(txt_name)  # 这是一个对txt文本结果排序的代码，key=frame，根据帧数排序

    source_file = open(txt_name)
    # 把frame存入列表img_names
    img_names = []
    for line in source_file:
        staff = line.split(',')
        img_name = staff[0]
        img_names.append(img_name)

    # 将每个frame的bbox数目存入字典
    name_dict = {}
    for i in img_names:
        if img_names.count(i):
            name_dict[i] = img_names.count(i)
    # print(name_dict)
    source_file.close()

    source_file = open(txt_name)
    for idx in name_dict:
        # print(str(idx).rjust(6, '0'))
        # 因为图片名称是000001格式的，所以需要str(idx).rjust(6, '0')进行填充
        img = cv2.imread(os.path.join(file_path_img, str(idx).rjust(6, '0') + '.jpg'))
        for i in range(name_dict[idx]):
            line = source_file.readline()
            staff = line.split(',')
            id = staff[1]
            cls = staff[7]
            box = staff[2:6]
            # print(id, box)
            # draw_bbox
            cv2.rectangle(img, (int(float(box[0])), int(float(box[1]))),
                          (int(float(box[0])) + int(float(box[2])), int(float(box[1])) + int(float(box[3]))),
                          (0, 255, 0), 2)
            # put_text
            cv2.putText(img, str(int(id)) + ' ' + str(int(cls)), (int(float(box[0])), int(float(box[1]))),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.3, (0, 0, 255), 1)
        # 保存图片
        cv2.imwrite(os.path.join(save_file_path, str(idx).rjust(6, '0') + '.jpg'), img)

    source_file.close()


if __name__ == '__main__':
    filename = os.listdir('D:/project/code/result')
    # print(filename)
    for name in filename:
        print('The video ' + name.split('.')[0] + ' begin!')
        draw_mot(name.split('.')[0])
        print('The video ' + name.split('.')[0] + ' Done!')

sort_gt.py

def sort_output(txt_path):
    with open(txt_path, 'r') as f:
        list = []
        for line in f:
            list.append(line.strip())

    with open(txt_path, "w") as f:
        for item in sorted(list, key=lambda x: int(str(x).split(',')[0])):
            f.writelines(item)
            f.writelines('\n')
        f.close()