以息肉分割的数据集为例,开源的数据集通常会包含两个文件,一个是imges,一个是masks,且里面均为png、jpg等图片格式,但是我们做yolo的实验,需要txt格式的标注文件。所以需要进行一个数据转换步骤,将你的 masks(分割掩码)转换为 YOLO 所需的边界框标注。

在masks和images的同一目录下创建一个py文件,把以下的代码粘贴上去。

import os
import cv2
import numpy as np

def convert_masks_to_yolo_labels(images_dir, masks_dir, output_dir, class_id=0):
    """
    将图像的分割掩码转换为 YOLO 格式的边界框标注文件。

    Args:
        images_dir (str): 原始图像文件所在的目录路径。
        masks_dir (str): 对应的二值掩码文件所在的目录路径。
        output_dir (str): 生成 YOLO 标注文件 (.txt) 的输出目录。
        class_id (int): 目标类别ID,默认为0 (例如,对于息肉)。
    """

    os.makedirs(output_dir, exist_ok=True) # 确保输出目录存在

    image_files = sorted([f for f in os.listdir(images_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))])

    for image_name in image_files:
        image_path = os.path.join(images_dir, image_name)
        # 假设掩码文件名和图像文件名相同,只是可能扩展名不同
        # 尝试寻找 .png, .jpg, .jpeg 扩展名的掩码
        mask_name_without_ext = os.path.splitext(image_name)[0]
        mask_path = None
        for ext in ['.png', '.jpg', '.jpeg']: # 常见的图片扩展名
            temp_mask_path = os.path.join(masks_dir, mask_name_without_ext + ext)
            if os.path.exists(temp_mask_path):
                mask_path = temp_mask_path
                break
        
        if mask_path is None:
            print(f"警告: 未找到 {image_name} 对应的掩码文件,跳过。")
            continue

        try:
            image = cv2.imread(image_path)
            mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) # 读取为灰度图

            if image is None:
                print(f"错误: 无法加载图像 {image_path},跳过。")
                continue
            if mask is None:
                print(f"错误: 无法加载掩码 {mask_path},跳过。")
                continue
            
            # 确保掩码是二值图像 (0或255)
            _, binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)

            # 查找轮廓
            # cv2.findContours 在 OpenCV 4.x 版本中返回两个值:contours, hierarchy
            # 在 OpenCV 3.x 及更早版本中返回三个值:image, contours, hierarchy
            contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            img_height, img_width, _ = image.shape
            label_lines = []

            for contour in contours:
                # 过滤掉过小的轮廓,避免噪声
                if cv2.contourArea(contour) < 10: # 可以根据需要调整这个阈值
                    continue

                # 计算最小外接矩形
                x, y, w, h = cv2.boundingRect(contour)

                # 归一化坐标
                x_center = (x + w / 2) / img_width
                y_center = (y + h / 2) / img_height
                width_norm = w / img_width
                height_norm = h / img_height

                # 格式化为 YOLO 标签行
                label_lines.append(f"{class_id} {x_center:.6f} {y_center:.6f} {width_norm:.6f} {height_norm:.6f}")

            # 写入 YOLO .txt 文件
            output_label_filename = os.path.splitext(image_name)[0] + '.txt'
            output_label_path = os.path.join(output_dir, output_label_filename)

            with open(output_label_path, 'w') as f:
                for line in label_lines:
                    f.write(line + '\n')
            
            print(f"已处理 {image_name},生成标注文件:{output_label_filename}")

        except Exception as e:
            print(f"处理 {image_name} 时发生错误: {e}")

# --- 如何使用 ---
if __name__ == "__main__":
    # 请根据你的实际情况修改以下路径!
    # 原始图像文件夹
    images_directory = 'your_dataset/images' 
    # 对应掩码文件夹
    masks_directory = 'your_dataset/masks'   
    # YOLO 格式标注文件输出目录
    output_labels_directory = 'your_dataset/labels' 

    # 运行转换
    convert_masks_to_yolo_labels(images_directory, masks_directory, output_labels_directory)
    print("\n转换完成!请检查输出目录。")

在相应地方改成自己的路径即可。

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐