【mask数据转yolov8数据格式】二值图转yolov8 txt 标签数据集

【mask数据转yolov8数据格式】二值图转yolov8 txt 标签数据集

import cv2
import numpy as np
import os
from pathlib import Path

def mask_to_yolo(mask_path, output_dir, max_points=100):
    """
    将掩膜数据转为yolo格式
    Convert binary mask to YOLO segmentation format.
    YOLO format: class_id x1 y1 x2 y2 ... xn yn
    Coordinates are normalized to [0, 1]
    """
    # Read the binary mask
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    if mask is None:
        print(f"Error reading mask: {mask_path}")
        return
    
    height, width = mask.shape
    
    # Ensure mask is binary
    _, mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
    
    # Find contours in the binary mask
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    # Get the filename without extension
    mask_filename = Path(mask_path).stem
    output_path = os.path.join(output_dir, f"{mask_filename}.txt")
    
    with open(output_path, 'w') as f:
        for contour in contours:
            # Calculate appropriate epsilon based on contour length and max points
            target_points = min(max_points, len(contour))
            if len(contour) > target_points:
                epsilon = 0.01 * cv2.arcLength(contour, True)
                while True:
                    approx = cv2.approxPolyDP(contour, epsilon, True)
                    if len(approx) <= target_points:
                        break
                    epsilon *= 1.1
            else:
                approx = contour
            
            # Convert to normalized coordinates
            normalized_coords = []
            for point in approx:
                x, y = point[0]
                # Ensure coordinates are strictly within [0, 1]
                x_norm = max(0, min(1, x / width))
                y_norm = max(0, min(1, y / height))
                normalized_coords.extend([x_norm, y_norm])
            
            # Only write if we have valid coordinates
            if normalized_coords and len(normalized_coords) >= 6:  # At least 3 points
                coords_str = ' '.join([f"{coord:.6f}" for coord in normalized_coords])
                f.write(f"0 {coords_str}\n")

def process_directory(input_dir, output_dir):
    """Process all mask files in a directory"""
    input_path = Path(input_dir)
    for mask_file in input_path.glob('*.png'):  # Adjust file extension if needed
        mask_to_yolo(str(mask_file), output_dir)

if __name__ == "__main__":
    # Example usage
    input_directory = "path/to/masks"	#存放masks图片文件目录，png格式
    output_directory = "path/to/output"  #存放转换的txt文件目录
    process_directory(input_directory, output_directory)