vis_dataset.py
yolov6\data\vis_dataset.py
目录
vis_dataset.py
1.所需的库和模块
2.def main(args):
3.if __name__ == '__main__':
1.所需的库和模块
# coding=utf-8
# 描述：可视化 yolo 标签图像。
# Description: visualize yolo label image.import argparse
import os
import cv2
import numpy as np# IMG_FORMATS 是一个包含图像文件扩展名的列表。它首先被定义为包含一些常见的图像文件格式的扩展名，如 bmp 、 jpg 、 jpeg 、 png 等。
IMG_FORMATS = ["bmp", "jpg", "jpeg", "png", "tif", "tiff", "dng", "webp", "mpo"]
# 接下来， extend 方法被用来将 IMG_FORMATS 列表中的每个扩展名转换为大写，并添加到列表的末尾。
# 执行这段代码后， IMG_FORMATS 列表将包含每种图像格式的扩展名及其大写版本。例如，如果原始列表是 ["bmp", "jpg", "jpeg", "png", "tif"] ，执行这段代码后，列表将变为 ["bmp", "jpg", "jpeg", "png", "tif", "BMP", "JPG", "JPEG", "PNG", "TIF"] 。
IMG_FORMATS.extend([f.upper() for f in IMG_FORMATS])
2.def main(args):
def main(args):# 从 args 对象中提取图像目录、标签目录和类别名称。img_dir, label_dir, class_names = args.img_dir, args.label_dir, args.class_names# 创建一个空字典，用于存储类别 ID 和类别名称的映射。label_map = dict()# 遍历类别名称列表，同时获取类别的索引（ID）和名称。for class_id, classname in enumerate(class_names):# 将类别 ID 和类别名称添加到 label_map 字典中。label_map[class_id] = classname# os.listdir(path)# os.listdir(path) 是 Python 的标准库 os 模块中的一个函数，它用于列出指定目录下的所有文件和子目录的名字。这个函数不会递归地列出子目录下的内容，仅返回指定路径下一级的文件和目录名。# 参数 ：# path ： 字符串，表示要列出内容的目录路径。如果省略该参数或设置为 None ，则默认为当前工作目录。# 返回值 ：# 返回一个列表，其中包含了指定路径下的文件和目录的名字。这些名字不包含路径信息，仅仅是文件或目录的名称。返回的列表中文件名的顺序是不确定的。# 异常 ：# 如果指定的路径不存在或无法访问，可能会抛出 FileNotFoundError 或 OSError 。# 遍历图像目录中的所有文件。for file in os.listdir(img_dir):# 检查文件扩展名是否为支持的图像格式。if file.split('.')[-1] not in IMG_FORMATS:print(f'[Warning]: Non-image file {file}') # [警告]：非图像文件 {file}。continue# 构建完整的图像文件路径。img_path = os.path.join(img_dir, file)# 构建对应的标签文件路径。# file[: file.rindex('.')] ：这部分代码获取文件名中最后一个点（ . ）之前的部分，即文件的基本名称（不包括扩展名）。# rindex 方法用于从字符串的末尾开始查找指定字符（在这个例子中是点），并返回该字符最后一次出现的索引。然后通过切片操作 [:] 获取从字符串开头到这个索引之前的所有字符。# str.rindex(sub[, start[, end]])# 在 Python 中，字符串（ str ）对象的 rindex() 方法用于从右侧开始查找子字符串（或字符）在字符串中最后一次出现的索引，并返回该索引。如果子字符串不存在于字符串中，会抛出一个 ValueError 。# 参数 ：# sub ：要查找的子字符串。# start （可选）：开始搜索的起始位置，默认为字符串的末尾。# end （可选）：结束搜索的位置，默认为字符串的开头。# 返回值 ：# 返回子字符串最后一次出现的索引。# 异常 ：# 如果 sub 子字符串不存在于字符串中，抛出 ValueError 。label_path = os.path.join(label_dir, file[: file.rindex('.')] + '.txt')# 开始一个 try-except 块，用于捕获并处理可能发生的异常。try:# 使用 OpenCV 读取图像数据。# cv2.imread(filename, flags=1)# cv2.imread 函数是 OpenCV 库中用于读取图像文件的函数。# filename ：字符串类型，指定要读取的图像文件的相对地址或完整路径。# flags ：这是一个可选参数，表示读取图像的方式，默认为1。该参数可以有以下几种取值：# cv2.IMREAD_COLOR(1) ：默认方式，3通道BGR彩色图像。# cv2.IMREAD_GRAYSCALE(0) ：单通道灰度图像。# cv2.IMREAD_UNCHANGED(-1) ：原样返回（使用Alpha通道）。# cv2.IMREAD_ANYDEPTH(2) ：输入具有相应深度时返回16位/32位图像，否则转为8位。# cv2.IMREAD_ANYCOLOR(4) ：以任何可能的颜色格式读取。img_data = cv2.imread(img_path)# 获取图像的高度、宽度和通道数。height, width, _ = img_data.shape# 为每个类别生成一个随机颜色。color = [tuple(np.random.choice(range(256), size=3)) for i in class_names]# 设置绘制边界框和文本的厚度。thickness = 2# 打开标签文件进行读取。# open(file, mode='r', buffering=-1, encoding=None, errors=None, newline=None, closefd=True, opener=None)# 在 Python 中， open 函数是内置的，用于打开一个文件，并返回一个文件对象。通过这个文件对象，你可以读取或写入文件。# file ：要打开的文件的路径。# mode （默认为 'r' ） ：文件打开模式，常见的模式有：# 'r' ：只读模式。这是默认值。# 'w' ：写入模式。如果文件已存在，则覆盖文件；如果文件不存在，则创建新文件。# 'x' ：独占创建模式。尝试创建一个新文件，如果文件已存在，则抛出 FileExistsError 。# 'a' ：追加模式。如果文件已存在，则写入数据会被追加到文件末尾；如果文件不存在，则创建新文件。# 'b' ：二进制模式。用于二进制文件的读写。# 't' ：文本模式（默认）。用于文本文件的读写。# '+' ：更新模式。用于读写文件。可以与 'r' 、 'w' 或 'a' 组合使用，例如 'r+' 、 'w+' 或 'a+' 。with open(label_path, 'r') as f:# 遍历标签文件中的每个边界框。for bbox in f:# 解析边界框数据，包括类别 ID、中心点坐标xy、宽度和高度。cls, x_c, y_c, w, h = [float(v) if i > 0 else int(v) for i, v in enumerate(bbox.split('\n')[0].split(' '))]# 计算边界框左上角的 x 坐标。x_tl = int((x_c - w / 2) * width)# 计算边界框左上角的 y 坐标。y_tl = int((y_c - h / 2) * height)# 在图像上绘制边界框。# cv2.rectangle(img, pt1, pt2, color[, thickness[, lineType[, shift]]]) -> img# 绘制矩形框# img ：指定一张图片，在这张图片的基础上进行绘制；（img相当于一个画板）。# pt1 ：由（x_min，x_min）组成，为绘制的边框的左上角。# pt2 ：由（x_max, y_max）坐标，为绘制的边框的右下角，示意如下。# color ：指定边框的颜色，由（B,G,R）组成，当为（255,0，0）时为绿色，可以自由设定。# thinkness ：线条的粗细值，为正值时代表线条的粗细（以像素为单位），为负值时边框实心。cv2.rectangle(img_data, (x_tl, y_tl), (x_tl + int(w * width), y_tl + int(h * height)), tuple([int(x) for x in color[cls]]), thickness)# 在边界框上方绘制类别名称。# cv2.putText(img, text, org, fontFace, fontScale, color, thickness)# 可以在OpenCV python中很容易地在图像上添加文本。# img ：它是要写上文字的图像。# text ：它是需要写在图像上的文字。# org ：图像中文本字符串的左下角。# fontFace ：文本的字体：点击这里查看OpenCV中可用的字体类型。# fontScale ：这个值通过乘以它的基本尺寸来扩展文本的大小。# color ：文本的颜色。# thickness ：文本行的厚度。cv2.putText(img_data, label_map[cls], (x_tl, y_tl - 10), cv2.FONT_HERSHEY_COMPLEX, 1, tuple([int(x) for x in color[cls]]), thickness)# 显示带有边界框和类别标签的图像。cv2.imshow('image', img_data)# 等待用户按键，0 表示无限等待。cv2.waitKey(0)except Exception as e:# 如果在处理图像时发生异常，打印错误信息。print(f'[Error]: {e} {img_path}')# 处理完所有图像后，打印完成信息。print('======All Done!======')3.if __name__ == '__main__':
if __name__ == '__main__':# parser = argparse.ArgumentParser()# parser.add_argument()# opt = parser.parse_args()# argparse 是一个Python模块：命令行选项、参数和子命令解析器。通过命令行运行Python脚本时，可以通过 ArgumentParser 来高效地接受并解析命令行参数。# 1. 创建一个 ArgumentParser 对象，这个对象会保存我们定义的所有命令行参数信息。# parser = argparse.ArgumentParser ( description=' 这是一个示例程序') 这里的 description 参数是可选的，它用于描述这个命令行程序的主要功能。# add_argument() 方法用于添加一个你希望程序接受的命令行参数。# parser.add_argument('input', type=str, help='输入文件的路径')# 参数类型：# *位置参数：这些参数是必须的，使用时需按正确的顺序提供。*可选参数：通常以-或--开始，不需要按特定顺序。# add_argument 的常见参数：# help ：描述参数作用的字符串。# type ：命令行参数应转换成的Python类型。# default ：如果命令行中未提供参数，则使用的默认值。# required ：可选参数是否可以省略（仅对可选参数有效）。# choices ：参数值限制为特定的选项。# action ：当参数在命令行中存在时采取的行动。# 可以使用 choices 参数来限制用户只能选择特定的值。# 例子：# parser.add_argument('-v', '--verbose', action='store_true', help='显示详细输出')# -v 和 --verbose 是同一个可选参数的两种形式，它们的类型是布尔值，当用户在命令行中指定这个参数时，它的值为True，否则为False。# action='store_true' 表示当指定这个参数时，将其值设置为True。# 2. 使用 parse_args() 方法来解析命令行参数。# args = parser.parse_args()# 3. 然后，我们就可以通过args.参数名的方式来访问这些参数的值了。# 创建一个新的参数解析器对象。parser = argparse.ArgumentParser()# 向解析器添加一个命令行参数 --img_dir ，它用于指定图像目录的路径。如果用户没有在命令行中提供这个参数，它将默认为 'VOCdevkit/voc_07_12/images' 。parser.add_argument('--img_dir', default='VOCdevkit/voc_07_12/images')# 向解析器添加另一个命令行参数 --label_dir ，用于指定标签文件目录的路径。默认值为 'VOCdevkit/voc_07_12/labels' 。parser.add_argument('--label_dir', default='VOCdevkit/voc_07_12/labels')# 添加一个命令行参数 --class_names ，用于指定类别名称的列表。如果用户没有提供这个参数，它将默认为一个包含多个常见类别名称的列表。parser.add_argument('--class_names', default=['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog','horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'])# 解析命令行参数，并将解析后的参数存储在 args 对象中。args = parser.parse_args()# 打印 args 对象，这将显示所有解析后的命令行参数及其值。print(args)# 调用 main 函数，并传入 args 对象。 main 函数将使用这些参数来执行一些操作，例如加载图像和标签，进行图像处理等。main(args)
