Yolo标签可视化

我的目录结构是这样的，yolov5的数据集目录结构

希望能把标签数据绘制到图片上看看效果，代码如下

from pathlib import Path
from alphabet import alphabet
import pandas as pd
import cv2
import matplotlib.pyplot as plt

alphabet = alphabet.split("\n")

label_root = Path("labels/test")
image_root = Path("images/test")


def paint(label_file,image_file):
    #读取标签
    df = pd.read_csv(label_file,sep=" ",names=['id','center-x','center-y','w','h'])
    df['id'] = df['id'].apply(lambda x:alphabet[x])
    df = df.sort_values(by='center-x')
    #读取图片
    img = cv2.imread(str(image_root/image_file))
    h,w = img.shape[:2]
    
    df[['center-x','w']] = df[['center-x','w']].apply(lambda x:x*w)
    df[['center-y','h']] = df[['center-y','h']].apply(lambda x:x*h)
    
    df['x1'] = df['center-x']-df['w']/2
    df['x2'] = df['center-x']+df['w']/2
    df['y1'] = df['center-y']-df['h']/2
    df['y2'] = df['center-y']+df['h']/2

    df[['x1','x2','y1','y2']] = df[['x1','x2','y1','y2']].astype('int')
    
    points = zip(df['x1'],df['y1'],df['x2'],df['y2'])
    for point in points:
        img = cv2.rectangle(img,point[:2],point[2:],color=(0, 255, 0),thickness=1)
    plt.imshow(img)
    plt.show()

for label_file in label_root.iterdir():
    image_file = label_file.name.replace(".txt",".jpg")
    paint(label_file,image_file)
    break