代码
import os
import cv2
import glob
import matplotlib.pyplot as plt
def validate_yolov5_annotations(images_dir, annotations_dir, output_dir):
os.makedirs(output_dir, exist_ok=True)
annotation_files = glob.glob(os.path.join(annotations_dir, '*.txt'))
for annotation_file in annotation_files:
image_file = os.path.join(images_dir, os.path.splitext(os.path.basename(annotation_file))[0] + '.jpg')
image = cv2.imread(image_file)
# image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
with open(annotation_file, 'r') as f:
lines = f.readlines()
for line in lines:
data = line.strip().split()
points = [float(point) for point in data[1:]]
for i in range(0, len(points), 2):
x = int(points[i] * image.shape[1])
y = int(points[i + 1] * image.shape[0])
cv2.circle(image, (x, y), 2, (0, 255, 0), -1)
output_file = os.path.join(output_dir, os.path.basename(image_file))
cv2.imwrite(output_file, image)
# plt.imshow(image)
# plt.title(os.path.basename(image_file))
# plt.show()
if __name__ == '__main__':
images_dir = r"E:\treasure3\images"
annotations_dir = r"E:\treasure3\images"
output_dir = r"E:\treasure3\images\r"
validate_yolov5_annotations(images_dir, annotations_dir, output_dir)
代码解释
import os
: 导入Python内置的操作系统模块,用于与操作系统交互。import cv2
: 导入OpenCV库,用于处理图像和视频数据。import glob
: 导入glob模块,用于在文件系统中搜索具有特定模式的文件。-
import matplotlib.pyplot as plt
: 导入matplotlib库的pyplot模块,用于可视化数据。 -
def validate_yolov5_annotations(images_dir, annotations_dir, output_dir):
:定义名为validate_yolov5_annotations
的函数,它接受三个参数:images_dir
(图片目录),annotations_dir
(注释目录)和output_dir
(输出目录)。 os.makedirs(output_dir, exist_ok=True)
: 使用os.makedirs()
函数创建输出目录,如果目录已经存在,则不抛出异常。-
annotation_files = glob.glob(os.path.join(annotations_dir, '*.txt'))
: 使用glob.glob()
函数获取注释目录中所有以.txt
为扩展名的文件的路径,并将结果存储在annotation_files
列表中。 -
for annotation_file in annotation_files:
: 遍历annotation_files
列表中的每个注释文件。 image_file = os.path.join(images_dir, os.path.splitext(os.path.basename(annotation_file))[0] + '.jpg')
: 根据当前注释文件的文件名,通过os.path
模块构建相应的图片文件路径。image = cv2.imread(image_file)
: 使用OpenCV的cv2.imread()
函数读取图片文件并将其存储在image
变量中。with open(annotation_file, 'r') as f:
: 使用open()
函数打开注释文件,并将文件对象存储在f
变量中。使用with
语句可以确保文件在使用后被正确关闭。-
lines = f.readlines()
: 使用readlines()
方法将注释文件的内容逐行读取,并将结果存储在lines
列表中。 -
for line in lines:
: 遍历lines
列表中的每一行。 data = line.strip().split()
: 使用strip()
方法去除行首尾的空白符,然后使用split()
方法将行按空格分割,并将结果存储在data
列表中。-
points = [float(point) for point in data[1:]]
: 将data
列表中从第二个元素开始的每个元素转换为浮点数,并将转换结果存储在points
列表中。 -
for i in range(0, len(points), 2):
: 使用range()
函数生成从0到points
列表长度的整数序列,每次递增2。 x = int(points[i] * image.shape[1])
: 计算x坐标,将points
列表中第i
个元素乘以图片的宽度(image.shape[1]
)并取整。y = int(points[i + 1] * image.shape[0])
: 计算y坐标,将points
列表中第i+1
个元素乘以图片的高度(image.shape[0]
)并取整。-
cv2.circle(image, (x, y), 2, (0, 255, 0), -1)
: 在图片上绘制一个半径为2的绿色圆形,此圆心的坐标为(x, y)
。 -
output_file = os.path.join(output_dir, os.path.basename(image_file))
: 根据输出目录和当前图片文件的文件名,使用os.path
模块构建相应的输出文件路径。 cv2.imwrite(output_file, image)
: 将修改后的图片保存到输出文件。if __name__ == '__main__':
: 检查当前模块是否为主模块(而非被导入的模块)。images_dir = r"E:\treasure3\images"
: 设置图片目录路径。annotations_dir = r"E:\treasure3\images"
: 设置注释目录路径。output_dir = r"E:\treasure3\images\r"
: 设置输出目录路径。validate_yolov5_annotations(images_dir, annotations_dir, output_dir)
: 调用validate_yolov5_annotations
函数,传递图片目录、注释目录和输出目录作为参数。
注意事项
文件路径也要注意