基于语义分割Ground Truth(GT)转换yolov5目标检测标签(路面积水检测例子)
概述
许多目标检测的数据是通过直接标注或者公开平台获得,如果存在语义分割Ground Truth的标签文件,怎么样实现yolov5的目标检测格式转换呢?查遍全网没有很好的方法,因此使用opencv自己写了一个,检验效果还不错。
这里的例子是基于极市平台的路面积水检测给出的数据集完成,由于平台只给了分割的示例数据,因此想使用yolo进行目标检测,需要自己进行标签的转换.已有的数据集有原图和label,这里的label是PNG格式的图片,如下所示:
数据集包含原图片以及相对应分割后的图片(标注文件),标注文件的格式为PNG,并且为单通道灰度图。在本任务中,为了直观地观察输出的PNG图片,使用如下灰度值:
- 背景:0
- 积水:1
可以直接使用该数据集进行分割模型的训练和验证,例如unet、deeplab等模型
流程
由于分割图像相对目标检测具有更精细的标注,因此只需要三个步骤即可完成整体转换流程
1、找到图中的分割块(目标),并找到其最小矩形
2、得到的最小矩形进行坐标转换,OpenCV进行边缘检测,得到坐标
3、根据YOLO坐标归一化方法,完成坐标转换,并存储*.txt
注意事项:数据量比较大,转换的时候会使用批量转换。由于是使用边缘检测获取最小矩形框,因此在训练过程中会出现误差,影响不大。
读取图像
使用os读取文件夹下的标签图像,并使用OpenCV读取图片
path = "./mask/train/"
files = os.listdir(path)
for file in files:
img = cv2.imread(path+file)
为了方便查看标签,可以先将标签可视化,原始分割标签是0,1,。。。,n的数字,所以在图像上看都是黑色的不明显,不同类别使用不同颜色显示出来,
#1->255####
print(img.shape)
for i in range(len(img)):
for j in range(len(img[0])):
if img[i][j] == 1:
img[i][j]print(file) = 255
print(img[i][j])
cv2.imwrite("./mask/{}".format(file),img)
#get_bbox
path = "./mask/ponding_sample_103.png"
img = cv2.imread(path)
img_w = img.shape[1]
img_h = img.shape[0]
转换并获取最小外接矩形
OpenCV常规读取是BGR格式,标签是单通道图形,需要转换为单通道的灰度图,方便后面检测,使用cv2.cvtColor转换为灰度图,或者读取时cv2.imread(path,0)
,0模式是灰度图。OpenCV给了很方便的边缘检测方法contours, hierarchy = cv.findContours( image, mode, method[, contours[, hierarchy[, offset]]] )
,
image:输入为二值图像,黑色为背景,白色为目标
单通道图像矩阵,可以是灰度图,但更常用的是二值图像,一般是经过Canny、拉普拉斯等边缘检测算子处理过的二值图像;该函数会修改原图像,因此若想保留原图像在,则需拷贝一份,在拷贝图里修改。
mode:
method:
offset:轮廓点的偏移量,格式为tuple,如(-10,10)表示轮廓点沿X负方向偏移10个像素点,沿Y正方向偏移10个像素点
返回值
contours:轮廓点。列表格式,每一个元素为一个3维数组(其形状为(n,1,2),其中n表示轮廓点个数,2表示像素点坐标),表示一个轮廓
hierarchy:轮廓间的层次关系,为三维数组,形状为(1,n,4),其中n表示轮廓总个数,4指的是用4个数表示各轮廓间的相互关系。第一个数表示同级轮廓的下一个轮廓编号,第二个数表示同级轮廓的上一个轮廓的编号,第三个数表示该轮廓下一级轮廓的编号,第四个数表示该轮廓的上一级轮廓的编号。
通过cv2.minAreaRect
和cv2.contourArea
获得每个目标的最小外接矩形
img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
cnts,_ = cv2.findContours(img.copy(),cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
for i, contour in enumerate(cnts):
area = cv2.contourArea(contour) # 计算包围形状的面积
rect = cv2.minAreaRect(contour) # 检测轮廓最小外接矩形,得到最小外接矩形的(中心(x,y), (宽,高), 旋转角度)
转换yolo坐标
yolo的标签形式是四个点,四个点分别代表bbox的中心坐标(x,y)和宽、高,但是需要根据原图的长宽进行归一化,既x/原图宽,y/原图高,w/原图宽,h/原图高,根据计算并记录即可
boxs = []
for i, contour in enumerate(cnts):
area = cv2.contourArea(contour) # 计算包围形状的面积
rect = cv2.minAreaRect(contour) # 检测轮廓最小外接矩形,得到最小外接矩形的(中心(x,y), (宽,高), 旋转角度)
temp =[0]
temp = [0]
temp.append(rect[0][0])
temp.append(rect[0][1])
temp.append(rect[1][0])
temp.append(rect[1][1])
temp[1] /= img_w
temp[2] /= img_h
temp[3] /= img_w
temp[4] /= img_h
# box = np.int0(cv2.boxPoints(rect))
boxs.append(temp) # 最后剩下的有用的框
注意:这里由于只有一个类别,因此我直接初始化写入的时候temp=[0],如果是多种目标对应修改即可
转存txt文件
存储在boxs数组中的数据就是需要保存的类别和对应的四个坐标
f = open("./labels/train2017/{}.txt".format(file.split(".")[0]), "w+")
for line in boxs:
line = str(line)[1:-2].replace(",","")
print(line)
f.write(line+"\n")
f.close()
转换之后的txt标签如下图,可直接在yolov5中训练使用
完整代码地址:https://github.com/magau123/CSDN/blob/master/GT2yolo.py