数据集制作——语义分割前png、jpg格式标签图转yolo格式.txt文件（附代码）

💪 专业从事且热爱图像处理，图像处理专栏更新如下👇：
📝《图像去噪》
📝《超分辨率重建》
📝《语义分割》
📝《风格迁移》
📝《目标检测》
📝《暗光增强》
📝《模型优化》
📝《模型实战部署》

在这里插入图片描述

一、YOLO格式

YOLO格式的数据集通常包含两部分：图像文件和对应的文本标注文件。每个文本标注文件中包含了图像中每个物体的类别和位置信息。每一行代表一个物体，格式如下：

<class_id> <x_center> <y_center> <width> <height>

其中，<class_id>是物体类别的ID，<x_center>和<y_center>是物体中心点的坐标，和是物体的宽度和高度。所有的坐标和尺寸都需要被归一化，即除以图像的宽度和高度，因此它们的值都在0到1之间。

二、实现步骤

要将.png格式的标签图转换为YOLO格式的.txt文件，需要以下步骤：

（1）读取.png标签图，每个物体应该被标记为不同的颜色；

（2）解析标签图，对每种颜色进行遍历，找出所有像素点的坐标；

（3）对每种颜色的像素点坐标进行分析，计算出对应的bounding box（通过找到最小和最大的x，y坐标来实现）；

（4）将bounding box的坐标和尺寸归一化，然后保存为.txt文件。

三、代码

3.1 参数修改

在这里插入图片描述

3.2 代码

注：.png格式个标签图像，必须是单通道图像。

import os
import cv2
import numpy as np


def convert_segmentation_to_yolo(img_path, output_path, num_classes):
    # 读取标签图
    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
    height, width = img.shape

    # 创建用于存储YOLO格式的列表
    yolo_labels = []

    # 遍历每个类别
    for class_id in range(num_classes):
        # 找到当前类别的所有像素位置
        class_pixels = np.where(img == class_id)

        # 如果当前类别不存在，跳过
        if len(class_pixels[0]) == 0:
            continue

        # 找到类别的最小和最大边界
        x_min = np.min(class_pixels[1])
        x_max = np.max(class_pixels[1])
        y_min = np.min(class_pixels[0])
        y_max = np.max(class_pixels[0])

        # 计算中心点和宽高，并归一化
        x_center = (x_min + x_max) / 2 / width
        y_center = (y_min + y_max) / 2 / height
        bbox_width = (x_max - x_min) / width
        bbox_height = (y_max - y_min) / height

        # 保存YOLO格式的标签
        yolo_labels.append(f"{class_id} {x_center} {y_center} {bbox_width} {bbox_height}")

    # 将YOLO标签写入.txt文件
    txt_file = os.path.splitext(os.path.basename(img_path))[0] + ".txt"
    with open(os.path.join(output_path, txt_file), "w") as f:
        for label in yolo_labels:
            f.write(label + "\n")


# 示例调用
# img_folder = 'path/to/your/png/folder'
# output_folder = 'path/to/your/txt/folder'
img_folder = 'Images/Segment_Images/image_png'
output_folder = 'Images/Segment_Images/label_txt'
# num_classes = 21  # 假设有21个类别
num_classes = 2  # 假设有21个类别

if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# 遍历标签图文件夹并转换
for img_file in os.listdir(img_folder):
    if img_file.endswith('.png'):
        img_path = os.path.join(img_folder, img_file)
        convert_segmentation_to_yolo(img_path, output_folder, num_classes)