基于语义分割Ground Truth（GT）转换yolov5目标检测标签（路面积水检测例子）

Posted 2023-01-30 麻花地

基于语义分割Ground Truth（GT）转换yolov5目标检测标签（路面积水检测例子）

概述

许多目标检测的数据是通过直接标注或者公开平台获得，如果存在语义分割Ground Truth的标签文件，怎么样实现yolov5的目标检测格式转换呢？查遍全网没有很好的方法，因此使用opencv自己写了一个，检验效果还不错。

这里的例子是基于极市平台的路面积水检测给出的数据集完成，由于平台只给了分割的示例数据，因此想使用yolo进行目标检测，需要自己进行标签的转换.已有的数据集有原图和label，这里的label是PNG格式的图片，如下所示：

数据集包含原图片以及相对应分割后的图片（标注文件），标注文件的格式为PNG，并且为单通道灰度图。在本任务中，为了直观地观察输出的PNG图片，使用如下灰度值：

1. 背景：0
2. 积水：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