写一个python脚本根据yolov5的分割标签生成对应图片的分割掩码图片

<think>好的，我现在需要帮用户写一个Python脚本，根据YOLOv5的分割标签生成对应的分割掩码图片。首先，我得确认用户的需求是什么。YOLOv5的分割标签通常是文本文件，里面包含了多边形坐标，每个对象一行，用归一化的坐标表示。而分割掩码图片应该是二值图像，缺陷区域用白色（255）表示，背景是黑色（0）。 接下来，我得考虑如何解析这些标签文件。每个标签文件对应一张图片，所以需要遍历标签目录中的每个.txt文件。对于每个文件，读取里面的每一行，解析出类别和归一化的多边形坐标。然后需要将这些归一化的坐标转换为实际图像的像素坐标，因为生成掩码需要具体的像素位置。 然后，需要确定原始图像的尺寸。用户可能没有提到图片的路径，所以可能需要用户指定图片目录，或者假设所有图片都是同样的大小？但更灵活的方式应该是根据对应的图片文件来获取尺寸。例如，如果图片存放在另一个目录，可以通过图片文件名找到对应的图片，并读取其尺寸。不过，如果图片不可用，可能需要用户输入默认的尺寸，或者从其他地方获取。但用户可能希望脚本独立运行，不需要依赖原始图片，这时候可能需要用户在脚本中指定图片的宽度和高度，或者在标签文件名中包含尺寸信息。不过通常YOLOv5的标签文件不包含尺寸信息，所以可能需要用户提供图片目录，以便读取图片尺寸。 但用户可能希望不依赖原始图片，所以可能需要另一种方式。或者，用户可能已经知道图片的尺寸，可以在脚本中设置。例如，用户可能所有的图片都是同一尺寸，比如640x640，这时候可以在脚本中硬编码。但为了灵活性，最好让用户通过参数指定图片的宽度和高度，或者从图片文件中读取。 不过，根据YOLOv5的规范，标签文件中的坐标是归一化的，即相对于图片宽度和高度的比例。因此，必须要有原始图片的尺寸才能将归一化坐标转换为实际坐标。因此，脚本需要访问原始图片来获取尺寸。或者，用户可能已经将图片尺寸信息保存在其他地方，但通常来说，最直接的方式是通过图片文件本身获取。 所以，脚本需要输入图片目录和标签目录，然后对于每个标签文件，找到对应的图片，读取其尺寸，然后处理标签。但这样会增加脚本的复杂性，因为需要处理图片文件。或者，用户可能已经将图片尺寸统一，可以在脚本中指定，比如通过命令行参数传递--img-width和--img-height。 但考虑到用户可能希望自动化处理，最好的做法是从图片文件中读取尺寸。因此，脚本需要两个目录：图片目录和标签目录，并且假设图片和标签的文件名一一对应（除了扩展名）。例如，图片是.jpg，标签是.txt，同名。 因此，步骤大致如下： 1. 遍历标签目录中的所有.txt文件。 2. 对于每个标签文件，找到对应的图片文件（比如将.txt替换为.jpg）。 3. 读取图片文件的尺寸（宽度和高度）。 4. 解析标签文件中的每一行，提取类别和多边形坐标。 5. 将归一化的多边形坐标转换为实际像素坐标。 6. 使用OpenCV或PIL创建一个空白掩码图像，尺寸与图片相同。 7. 在掩码图像上绘制多边形，填充为白色（255）。 8. 保存掩码图像到输出目录，保持与标签相同的文件名，但扩展名为.png。 现在需要处理可能的问题，比如图片文件不存在，或者尺寸无法读取。这时候可能需要异常处理，或者跳过这些标签文件。此外，YOLOv5的分割标签格式是怎样的？每个行的格式应该是：class_id x1 y1 x2 y2 ... xn yn，其中坐标是归一化的，按多边形顺序排列。 例如，一行可能是：0 0.5 0.5 0.6 0.5 0.6 0.6 0.5 0.6，表示类别0，四个点构成一个矩形。 因此，在解析每一行时，首先分割成列表，第一个元素是类别，后面是交替的x和y坐标。然后，将这些归一化的坐标转换为实际的x和y坐标，乘以图片宽度和高度，并取整。 接下来，使用OpenCV的fillPoly函数来填充多边形区域。需要注意的是，一个标签文件中可能有多个对象（多行），每个对象对应一个多边形，需要将所有多边形绘制到同一个掩码图像上。 此外，掩码图像应该是单通道的，每个像素为0或255，表示背景和前景。 现在，代码的大致结构： - 导入必要的库：os, cv2, numpy等。 - 定义输入输出路径：图片目录，标签目录，输出目录。 - 遍历标签目录中的每个.txt文件。 - 对于每个文件，获取对应的图片路径。 - 读取图片并获取尺寸。 - 创建空白掩码图像。 - 解析标签文件的每一行，转换为多边形坐标。 - 使用cv2.fillPoly绘制多边形。 - 保存掩码图像。 需要考虑的问题： - 如何处理不同通道数的图片？比如，如果图片是RGB，但掩码是单通道。但读取图片时只需要尺寸，所以用cv2.imread的IMREAD_UNCHANGED或IMREAD_COLOR，然后获取shape的前两个值作为高度和宽度。 - 需要确保标签文件和图片文件一一对应，且文件名相同（除了扩展名）。 - 需要处理可能的多边形闭合问题，但fillPoly会自动处理。 现在，编写代码的大纲： 导入库： import os import cv2 import numpy as np 设置路径： image_dir = 'path/to/images' label_dir = 'path/to/labels' output_dir = 'path/to/masks' 创建输出目录： os.makedirs(output_dir, exist_ok=True) 遍历标签文件： for label_file in os.listdir(label_dir): if not label_file.endswith('.txt'): continue # 获取对应的图片文件名 image_name = os.path.splitext(label_file)[0] + '.jpg' # 假设图片是jpg格式 image_path = os.path.join(image_dir, image_name) # 读取图片尺寸 if not os.path.exists(image_path): print(f"Image {image_path} not found, skipping...") continue img = cv2.imread(image_path) if img is None: print(f"Failed to read image {image_path}, skipping...") continue img_h, img_w = img.shape[:2] # 创建掩码 mask = np.zeros((img_h, img_w), dtype=np.uint8) # 解析标签文件 label_path = os.path.join(label_dir, label_file) with open(label_path, 'r') as f: lines = f.readlines() for line in lines: parts = line.strip().split() if len(parts) < 3 or len(parts) % 2 != 1: print(f"Invalid line in {label_file}: {line}") continue class_id = int(parts[0]) coords = list(map(float, parts[1:])) # 转换为像素坐标 polygon = [] for i in range(0, len(coords), 2): x = coords[i] * img_w y = coords[i+1] * img_h polygon.append([x, y]) polygon = np.array(polygon, dtype=np.int32).reshape((-1, 1, 2)) # 绘制多边形 cv2.fillPoly(mask, [polygon], color=255) # 保存掩码 mask_filename = os.path.splitext(label_file)[0] + '.png' mask_path = os.path.join(output_dir, mask_filename) cv2.imwrite(mask_path, mask) 这样应该可以生成对应的掩码图片。需要注意的点： - 假设图片是.jpg格式，如果用户使用其他格式（如.png），需要相应修改。 - 处理坐标转换时，需要确保x和y的顺序正确，YOLO的坐标是归一化的中心坐标吗？或者分割标签是多边形的归一化坐标？根据YOLOv5的分割标签格式，每个点是一对归一化的(x,y)坐标，按顺序排列，所以上述处理是正确的。 - 需要处理可能的异常，比如无法读取图片，或者标签文件中的坐标数量不对。 - 生成的掩码是单通道的PNG图片，背景为0，目标区域为255。 可能还需要用户调整的部分，比如图片扩展名、路径设置等，应该在脚本中作为参数，或者通过命令行参数传入，但用户可能需要一个简单易用的脚本，所以可以写成函数，让用户自行修改路径。 此外，可以添加注释，说明如何运行脚本，以及必要的依赖安装（如OpenCV）。 最后，测试一下这个脚本是否能够正确处理标签文件。例如，一个标签文件中的一行是：0 0.1 0.1 0.1 0.9 0.9 0.9 0.9 0.1，这应该生成一个四边形，对应的掩码区域。</think>以下是将YOLOv5分割标签转换为掩码图像的Python实现，包含详细注释和异常处理： ```python import os import cv2 import numpy as np def yolo_seg_to_mask(image_dir, label_dir, output_dir, img_ext='.jpg'): """ 将YOLOv5分割标签转换为二值掩码图像 参数： image_dir: 原始图像目录路径 label_dir: YOLO标签目录路径 output_dir: 掩码输出目录路径 img_ext: 原始图像扩展名（默认.jpg） """ # 创建输出目录 os.makedirs(output_dir, exist_ok=True) # 遍历所有标签文件 for label_file in os.listdir(label_dir): if not label_file.endswith('.txt'): continue # 构建对应图像路径 base_name = os.path.splitext(label_file)[0] image_path = os.path.join(image_dir, base_name + img_ext) # 检查图像是否存在 if not os.path.exists(image_path): print(f"警告：找不到对应图像 {image_path}，已跳过") continue # 读取图像获取尺寸 img = cv2.imread(image_path) if img is None: print(f"错误：无法读取图像 {image_path}，已跳过") continue img_h, img_w = img.shape[:2] mask = np.zeros((img_h, img_w), dtype=np.uint8) # 创建空白掩码 # 解析标签文件 label_path = os.path.join(label_dir, label_file) try: with open(label_path, 'r') as f: lines = f.readlines() except Exception as e: print(f"错误：无法读取标签文件 {label_path} ({str(e)})") continue # 处理每个多边形 for line in lines: parts = line.strip().split() if len(parts) < 3 or (len(parts)-1)%2 != 0: print(f"无效行格式：{line} （文件：{label_file}）") continue try: # 转换坐标到绝对像素值 coords = np.array([float(x) for x in parts[1:]]) polygon = coords.reshape((-1, 2)) * [img_w, img_h] polygon = polygon.astype(np.int32).reshape((-1, 1, 2)) # 绘制填充多边形 cv2.fillPoly(mask, [polygon], color=255) except Exception as e: print(f"坐标转换错误：{line} （文件：{label_file}）") continue # 保存掩码图像 output_path = os.path.join(output_dir, f"{base_name}_mask.png") cv2.imwrite(output_path, mask) print(f"生成掩码：{output_path}") if __name__ == "__main__": # 使用示例 - 修改为实际路径 yolo_seg_to_mask( image_dir='path/to/images', label_dir='path/to/labels', output_dir='path/to/masks', img_ext='.jpg' # 根据实际情况修改图像格式 ) ``` --- ### 关键功能说明 1. **坐标转换逻辑**： ```python polygon = coords.reshape((-1, 2)) * [img_w, img_h] ``` YOLO格式的归一化坐标 $(x_i, y_i)$ 转换为实际像素坐标： $$x_{\text{pixel}} = x_i \times \text{width}$$ $$y_{\text{pixel}} = y_i \times \text{height}$$ 2. **掩码生成原理**： - 创建与原始图像同尺寸的全黑画布 - 对每个多边形执行 `cv2.fillPoly()` 填充白色(255) - 多个多边形会自动合并（适用于重叠缺陷） --- ### 使用准备 1. 安装依赖： ```bash pip install opencv-python numpy ``` 2. 文件结构要求： ``` ├── images/ # 原始图像 │ ├── 0001.jpg │ └── 0002.jpg ├── labels/ # YOLO标签 │ ├── 0001.txt │ └── 0002.txt └── masks/ # 输出目录（自动创建） ``` 3. YOLO标签格式示例 (`0001.txt`)： ```text 0 0.30 0.45 0.60 0.45 0.60 0.80 0.30 0.80 ``` - 第一个数字为类别ID - 后续为交替排列的归一化x,y坐标 --- ### 效果验证 输入图像 (缺陷示例) | 生成的掩码 ---|---  |  --- ### 常见问题处理 1. **坐标点数量异常**： - 自动跳过无法构成多边形的无效行 - 输出警告信息定位错误位置 2. **图像-标签匹配问题**： - 文件名必须严格一致（仅扩展名不同） - 自动跳过无对应图像的标签文件 3. **大尺寸图像处理**： - 添加内存优化（可选）： ```python mask = cv2.createMat(img_h, img_w, cv2.CV_8UC1) ```