语义分割标签转化为YOLO格式标签

数据集

数据集包含原图片以及相对应分割后的图片（标注文件）

思路

1、找到图中的分割块（目标），并找到其最小矩形

2、得到的最小矩形进行坐标转换，OpenCV进行边缘检测，得到坐标

3、根据YOLO坐标归一化方法，完成坐标转换，并存储*.txt

代码

import cv2
import os
import numpy as np

# 定义类别标签
background_label = 1
object_label = 0

# 输入文件夹和输出文件夹
input_folder = "./data/label/"
output_folder = "./data/labels/"
count = 1
# 获取输入文件夹中的所有标签图像文件名
image_files = [f for f in os.listdir(input_folder) if f.endswith('.png')]

def extract_number(filename):
    # 提取文件名中的数字部分
    return int(''.join(filter(str.isdigit, filename)))

image_files = sorted(os.listdir(input_folder), key=extract_number)
# 遍历每个标签图像文件
for image_file in image_files:
    print(image_file)
    # 读取语义分割图像
    image_path = os.path.join(input_folder, image_file)
    semantic_image = cv2.imread(image_path, cv2.IMREAD_COLOR)

    # 将红色区域标记为目标区域
    object_mask = (semantic_image[:, :, 2] == 128) & (semantic_image[:, :, 0] == 0) & (semantic_image[:, :, 1] == 0)
    # cv2.imshow("object_mask", object_mask.astype(np.uint8)*255)
    # cv2.waitKey()
    
    # print(object_mask)
    # 寻找目标区域的轮廓
    contours, _ = cv2.findContours(object_mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    print(contours)
    # 定义输出文件路径
    output_path = "{}img{}.txt".format(output_folder, count)
    count += 1

    # 将目标信息保存到txt文件中
    with open(output_path, "w") as f:
        for contour in contours:
            # 计算目标的边界框
            x, y, w, h = cv2.boundingRect(contour)

            # 计算边界框的中心坐标和归一化坐标
            center_x = (x + w / 2) / semantic_image.shape[1]
            center_y = (y + h / 2) / semantic_image.shape[0]
            normalized_w = w / semantic_image.shape[1]
            normalized_h = h / semantic_image.shape[0]

            # 写入YOLO格式的标签信息到txt文件
            line = f"{object_label} {center_x:.6f} {center_y:.6f} {normalized_w:.6f} {normalized_h:.6f}\n"
            f.write(line)

    print(f"Processed: {image_file}")
    
print("Batch processing completed.")

验证

将转换后的标签显示在原图像上查看是否准确

import cv2
import numpy as np
 
# 读入图像
image = cv2.imread('label11.png')    # 标签
image2 = cv2.imread('img11.png')    # 原图
semantic_image = image

# 将红色区域标记为目标区域
object_mask = (semantic_image[:, :, 2] == 128) & (semantic_image[:, :, 0] == 0) & (semantic_image[:, :, 1] == 0)
 
# 查找图像中的轮廓
contours, _ = cv2.findContours(object_mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
# 边界框
x, y, w, h = cv2.boundingRect(contours[0])

# 绘制方框
brcnt = np.array([[[x,y]] ,[[x+w , y]] , [[x+w , y+h]] , [[x,y+h]]] )
cv2.drawContours( image2 , [brcnt] , -1 ,(255,255,255) , 2)

# 在原图像上绘制轮廓
cv2.drawContours(image2, contours, -1, (255, 0, 0), 3)
 
# 显示图像
cv2.namedWindow('Contours',0)
cv2.imshow("Contours", image2)
cv2.waitKey(0)
cv2.destroyAllWindows()

参考文章

基于语义分割Ground Truth（GT）转换yolov5目标检测标签（路面积水检测例子）_路面积水数据集_麻花地的博客-CSDN博客