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

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

概述

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

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

基于语义分割Ground Truth(GT)转换yolov5目标检测标签(路面积水检测例子)_第1张图片

数据集包含原图片以及相对应分割后的图片(标注文件),标注文件的格式为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

基于语义分割Ground Truth(GT)转换yolov5目标检测标签(路面积水检测例子)_第2张图片

method

基于语义分割Ground Truth(GT)转换yolov5目标检测标签(路面积水检测例子)_第3张图片

offset:轮廓点的偏移量,格式为tuple,如(-10,10)表示轮廓点沿X负方向偏移10个像素点,沿Y正方向偏移10个像素点

返回值
contours:轮廓点。列表格式,每一个元素为一个3维数组(其形状为(n,1,2),其中n表示轮廓点个数,2表示像素点坐标),表示一个轮廓

hierarchy:轮廓间的层次关系,为三维数组,形状为(1,n,4),其中n表示轮廓总个数,4指的是用4个数表示各轮廓间的相互关系。第一个数表示同级轮廓的下一个轮廓编号,第二个数表示同级轮廓的上一个轮廓的编号,第三个数表示该轮廓下一级轮廓的编号,第四个数表示该轮廓的上一级轮廓的编号。

通过cv2.minAreaRectcv2.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中训练使用

image-20230128171139014
完整代码地址:https://github.com/magau123/CSDN/blob/master/GT2yolo.py

你可能感兴趣的:(目标检测,深度学习,python,目标检测,计算机视觉,深度学习)