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