【目标检测】YOLO数据集标注可视化

可视化YOLO标注：

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

alphabet = alphabet.split("\n")
alphabet = ["label1","label2","label3"]

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


def paint(label_file, image_file):
    # Read labels
    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')
    # Read images
    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)

    cv2.imwrite('img.png', img)


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