YOLOv5可视化训练集的标注框和类别

导入必要的库

#os：包含了多种能够操作路径的函数，创建删除文件夹，判断文件夹是否存在等一系列操作
import os
#shutil：是os库的一些补充，也是包含了对路径的各种操作
import shutil
from pathlib import Path
import numpy as np
import cv2
#tqdm：一种读取数据的库，可以将读取过程显式的用进度条的方式呈现
from tqdm import tqdm
#argparse：argparse 模块是 Python 内置的用于命令项选项与参数解析的模块，
#argparse 模块可以让人轻松编写用户友好的命令行接口，
#能够帮助程序员为模型定义参数。
import argparse
1
2
3
4
5
6
7
8
9
10
11
12
13

设置超参数

parser = argparse.ArgumentParser()
parser.add_argument('--imgs', type=str, default='', help='图像路径')
parser.add_argument('--labels', type=str, default='', help='标签路径')
parser.add_argument('--save', type=str, default='runs/GT', help='保存路径')
args = parser.parse_args()

# 修改输入图片文件夹
img_folder = args.imgs
img_list = os.listdir(img_folder)
img_list.sort()
# 修改输入标签文件夹
label_folder = args.labels
label_list = os.listdir(label_folder)
label_list.sort()
# 输出图片文件夹位置
output_folder = args.save

labels = ['']  # 这里修改为自己的类别

# 色盘，可根据类别添加新颜色，注意第一个别动，这是字体的颜色也就是白色，往后面推
colormap = [(255, 255, 255)]  # 不是RGB,是BGR

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

标签坐标反归一化

# 坐标转换
def xywh2xyxy(x, w1, h1, img):
    label, x_, y_, w_, h_ = x

    label = int(label)
    label_ind = label

    # 边界框反归一化（还原成图像尺寸下的坐标）
    x_t = x_ * w1
    y_t = y_ * h1
    w_t = w_ * w1
    h_t = h_ * h1

    # 计算坐标 从xywh->xyxy(左上xy，右下xy)
    top_left_x = x_t - w_t / 2
    top_left_y = y_t - h_t / 2
    bottom_right_x = x_t + w_t / 2
    bottom_right_y = y_t + h_t / 2

    p1, p2 = (int(top_left_x), int(top_left_y)), (int(bottom_right_x), int(bottom_right_y))
    # 绘制矩形框
    cv2.rectangle(img, p1, p2, colormap[label_ind + 1], thickness=2, lineType=cv2.LINE_AA)
    label = labels[label_ind]
    if label:
        w, h = cv2.getTextSize(label, 0, fontScale=2 / 3, thickness=2)[0]  # text width, height
        outside = p1[1] - h - 3 >= 0  # label fits outside box
        p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
        # 绘制矩形框填充
        cv2.rectangle(img, p1, p2, colormap[label_ind + 1], -1, cv2.LINE_AA)
        # 绘制标签
        cv2.putText(img, label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), 0, 2 / 3, colormap[0],
                    thickness=2, lineType=cv2.LINE_AA)
    return img
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

主函数

 # 创建输出文件夹
 #如果路径存在（之前保存的结果），删除路径下的所有文件
    if Path(output_folder).exists():
        shutil.rmtree(output_folder)
#重新创建文件夹
    os.mkdir(output_folder)
#通过tqdm来遍历图片路径，并实时显示进度条
    for img in tqdm(img_list):
        if not 'png' in img:
            continue
        image_path = img_folder + "/" + img
        label_path = label_folder + "/" + img.split('.')[0] + '.txt'
        # 读取图像文件
        # print(img)
        img = cv2.imread(str(image_path))

        h, w = img.shape[:2]
        # 读取 labels
        with open(label_path, 'r') as f:
            lb = np.array([x.split() for x in f.read().strip().splitlines()], dtype=np.float32)
        # 绘制每一个目标
        for x in lb:
            # 反归一化并得到左上和右下坐标，画出矩形框
            img = xywh2xyxy(x, w, h, img)
        cv2.imwrite(output_folder + '/' + '{}.png'.format(image_path.split('/')[-1][:-4]), img)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25