可视化YOLO模型的检测框和关键点的python代码

本篇文章将yolo格式标注文件(txt)中的数据绘制到原图上，实现可视化。文中代码参考了yolov5-face的绘图代码。
https://github.com/deepcam-cn/yolov5-face

从yolov5-fac项目中导入需要的文件

from utils.general import xywh2xyxy
from utils.plots import colors,  plot_one_box

完整代码

import torch
import cv2
from utils.general import xywh2xyxy
from utils.plots import colors,  plot_one_box

if __name__ == '__main__':
    colormap = [(255, 255, 255), (56, 56, 255), (151, 157, 255), (31, 112, 255), (29, 178, 255)]   # 不是RGB,是BGR

    # 将yolo格式的标注好的图片展示出来
    im0 = cv2.imread("./face_img/000000191874.jpg")
    f = open("./face_img/000000191874.txt", "r")
    img_save_path = './face_img/'
    det_list = []
    for line in f.readlines():
        # line = line.strip('\n')  #去掉列表中每一个元素的换行符
        det_list = list(map(float, line.split(' ')))
        print("det_list",det_list)
        f.close()
        class_label = int(det_list[0])
        # 前5个参数为 0类型 xyxy
        xywh = det_list[1:5]
        # 将xywh转化为xyxy格式
        gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]
        print(torch.tensor(xywh).view(1, 4))
        xyxy = (xywh2xyxy(torch.tensor(xywh).view(1, 4)) * gn).view(-1).tolist()
        print("xyxy", xyxy)
        # 后10个参数为 5个人脸关键点(x, y)数据 官方的是百分比表示法，需要乘宽高
        kpts = det_list[5:]
        # 原图的宽和高
        weight = im0.shape[1]
        height = im0.shape[0]

        for index, i in enumerate(kpts):
            if index % 2 == 0:
                kpts[index] = kpts[index] * weight
            if index % 2 == 1:
                kpts[index] = kpts[index] * height
        for i in range(len(kpts)):
            print("kpts", kpts[i])
        
          # 画人脸关键点
        print("\n face_keypoint信息:")
        for i in range(0, len(kpts), 2):
            x,y = kpts[i:i+2]
            cv2.circle(im0, (int(x), int(y)), 2, (255,255,0), -1) #画圆点
            
    
        # 画人体框和人脸框
        xmin,ymin,xmax,ymax =  xyxy[0:4]
        p1, p2 = (int(xmin), int(ymin)), (int(xmax), int(ymax))
        # 绘制矩形框
        cv2.rectangle(im0, p1, p2, colormap[1], thickness=2, lineType=cv2.LINE_AA)
        cv2.putText(im0, str(class_label), (int(xmin), int(ymin - 2)), fontFace = cv2.FONT_HERSHEY_SIMPLEX,fontScale=0.5,color=(0,255,0),thickness=2, lineType=cv2.LINE_AA)

        
            
    cv2.imwrite(img_save_path + 'face_dot_02.jpg', im0)
    cv2.waitKey(0)  # 1 millisecond