目标检测数据集绘制GroundTurth

目标检测数据集绘制GroundTurth

以下代码的使用，读者需要注意以下几点：

• 读者需要修改srcFolder、destFolder和indexToLabel

import cv2
import numpy as np
import os

# 源始数据集路径与绘制后存放的路径
srcFolder = r"E:\Datasets\ExDark"
destFolder = r"E:\Datasets\ExDark_GroundTruth"
# 数据集标签
indexToLabel = ['Bycicle', 'Boat', 'Bottle', 'Bus', 'Car', 'Cat',
                'Chair', 'Cup', 'Dog', 'Motorbike', 'People', 'Table']
# 图片后缀，以适配任何图片
imgSuffix = ['.jpg', '.png', '.jpeg', '.JPEG', '.JPG',
             '.PNG', '.bmp', '.BMP']


def get_axisinfo(filepath, w, h):
    lines = []  # 存放txt文件的所有行
    bboxesinfo = []  # 存放所有bboxes的信息
    with open(filepath, 'r') as f:
        lines = f.readlines()

    for line in lines:
        line = line.strip()
        info = line.split(' ')
        # 每一行 为一个锚框
        label_index, x_center, y_center, width, height = int(info[0]), \
            float(info[1]), float(info[2]), float(info[3]), float(info[4])
        # 反归一化以计算坐标
        x1 = int((x_center - width / 2.0) * w)
        y1 = int((y_center - height / 2.0) * h)
        x2 = int((x_center + width / 2.0) * w)
        y2 = int((y_center + height / 2.0) * h)

        bboxesinfo.append([label_index, x1, y1, x2, y2])
    return bboxesinfo


def draw_bboxes(image, bboxes):
    for bbox in bboxes:
        label_index, x1, y1, x2, y2 = bbox
        label = indexToLabel[label_index]
        # 画框
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 1)  # 绿色框
        #  写标签
        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX,
                    0.9, (255, 255, 255), 2)
    return image


def AutoProcess(srcFolder, indexToLabel):
    for root, dirs, files in os.walk(srcFolder):
        #  获取标签文件夹
        labelfolder = os.path.join(*root.split('\\')[:-1], 'labels')
        for file in files:
            #  获取图片名及扩展名
            name, exname = os.path.splitext(file)
            if exname in imgSuffix:
                print("Processing:", file)
                # 获取图片路径
                imgfilepath = os.path.join(root, file)
                #  获取标签路径
                labelfilepath = os.path.join(labelfolder,
                                             name + ".txt")
                img = cv2.imread(imgfilepath)
                # opencv读的图片形状是（高，宽，通道数）
                bboxes = get_axisinfo(labelfilepath, img.shape[1],
                                      img.shape[0])
                result = draw_bboxes(img, bboxes)
                # 保存
                cv2.imwrite(os.path.join(destFolder, file), result)


if __name__ == '__main__':
    AutoProcess(srcFolder, indexToLabel)