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

您所在的位置：网站首页 › 显示跟踪结果 › Python 根据生成的txt对多目标跟踪结果画框可视化显示

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

2023-10-22 12:49| 来源: 网络整理| 查看: 265

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()

【本文地址】

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

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

今日新闻

推荐新闻