一般而言,按照YOLO的数据分布方式及标签文件排布方式划分好数据集后,训练集中的图片与训练集中的标签是一一对应的,所以我在这里直接读取了训练集图片文件和训练集标签文件,利用图片名称解析标签文件,而没有使用for loop来搜索,这样会使程序的运行效率得到较大提升。当然,如果自己的数据集不是YOLO排布,仅需要在标签路径这里搜索一下即可。
另外,txt文件中存储的xywh是经过归一的,所以需要将归一化后的坐标化为原图中的绝对坐标,也就是程序中的
x_max=int((x+0.5*w)*width)
x_min=int((x-0.5*w)*width)
y_max=int((y+0.5*h)*height)
y_min=int((y-0.5*h)*height)
其具体推导过程如下图所示:
总体的代码如下:
import cv2
import os
img_file=r'train_dataset'
img_dir=os.listdir(img_file)
for image in img_dir:
# 解析出图片名称,后续直接用名称取标签
image_name=image.split('.')[0]
# 图片路径
image_path=img_file+'\\'+image
# 读取图片
img = cv2.imread(image_path)
# 图片的尺寸
height,width,_=img.shape
# 标签路径
labelpath = r'train_label'+'\\'+image_name+'.txt'
# 读取标签
f=open(labelpath,'r')
# 将换行符去除
data=[x.strip() for x in f.readlines() if x.strip()]
# 关闭文件,防止内存溢出
f.close()
for line in data:
# 类型转换,str->float
label,x,y,w,h=line = list(map(float, line.split(' ')))
# 得到制作BB的坐标
x_max=int((x+0.5*w)*width)
x_min=int((x-0.5*w)*width)
y_max=int((y+0.5*h)*height)
y_min=int((y-0.5*h)*height)
# 5个参数的含义(图片,矩形框左上角点坐标,矩形框右下角点坐标,矩形框的颜色,线框粗细)
cv2.rectangle(img, (x_min, y_max), (x_max, y_min), (0, 255, 0), 1)
# 添加文本信息
font = cv2.FONT_HERSHEY_SIMPLEX
text = str(int(label))
# 7个参数的含义(图片,文本信息,放置位置,字体,字体大小,字体颜色,粗细)
cv2.putText(img, text, (x_min, y_min), font, 1, (0, 0, 255), 2)
img_demo = cv2.resize(img,(800,800))
cv2.moveWindow("ship_dataset", 100, 200)
cv2.imshow('ship_dataset', img_demo)
cv2.waitKey(0)
cv2.destroyAllWindows()