Mask标注转YOLO格式txt文件

因为要用yolov5训练，但是下载的数据集只有image和mask文件，需要把其转换成yolo格式。

在网上找的教程很少又看不懂，找了一个简单的又转换的不对，转出来的txt里面的值有多个。

我们知道，txt的值包括{类别；x；y；w；h}，否则在labelimg打开查看的时候会出问题。

就自己搞了一个代码，可以直接转换，分享给大家，也作为自己工作的一个记录。

import os
import cv2
import numpy as np

# Function to convert mask annotations to YOLO format for a single image-mask pair
def mask_to_yolo_single(img_path, mask_path, output_path):
    # Load image and mask
    img = cv2.imread(img_path)
    if img is None:
        print(f"Error: Unable to read image at {img_path}")
        return
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    if mask is None:
        print(f"Error: Unable to read mask at {mask_path}")
        return

    # Determine image dimensions
    img_height, img_width, _ = img.shape

    # Find contours in the mask
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Open a file to write YOLO format annotations
    with open(output_path, 'w') as f:
        for contour in contours:
            # Get bounding box coordinates
            x, y, w, h = cv2.boundingRect(contour)

            # Calculate normalized coordinates
            x_normalized = x / img_width
            y_normalized = y / img_height
            w_normalized = w / img_width
            h_normalized = h / img_height

            # Write to file in YOLO format: <object-class> <x> <y> <width> <height>
            f.write(f"0 {x_normalized} {y_normalized} {w_normalized} {h_normalized}\n")

    print(f"Converted {img_path} and saved annotations to {output_path}")

# Function to process all images and masks in a directory
def process_folder(img_folder, mask_folder, output_folder):
    # Ensure output folder exists
    os.makedirs(output_folder, exist_ok=True)

    # Get list of files in the input folders
    img_files = os.listdir(img_folder)
    mask_files = os.listdir(mask_folder)

    # Sort files to ensure corresponding images and masks match
    img_files.sort()
    mask_files.sort()

    # Process each pair of image and mask
    for img_file, mask_file in zip(img_files, mask_files):
        img_path = os.path.join(img_folder, img_file)
        mask_path = os.path.join(mask_folder, mask_file)

        # Generate output path for labels
        label_file = os.path.splitext(img_file)[0] + '.txt'
        output_path = os.path.join(output_folder, label_file)

        # Convert masks to YOLO format
        mask_to_yolo_single(img_path, mask_path, output_path)

# Example usage
img_folder = 'image文件夹位置'
mask_folder = 'mask文件夹位置'
output_folder = '需要保存的txt文件夹位置'

process_folder(img_folder, mask_folder, output_folder)

转换之后我用labelimg检查了一下，标签的位置是对的，举个例子如下图所示
原始图片：

mask：

再labelimg里面的标签显示：