Python实用脚本:可视化分割txt标签数据
可视化原始图像以及对应txt分割标签,校验txt生成是否正确。
import cv2
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollectiondef visualize_masks(image_path, label_path, output_path=None):# 读取图像image = cv2.imread(image_path)image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)height, width = image.shape[:2]# 创建图形fig, ax = plt.subplots(1)ax.imshow(image)# 读取标签文件with open(label_path, 'r') as f:lines = f.readlines()patches = []colors = []# 预定义一些颜色用于不同类别color_map = {0: (1, 0, 0, 0.5), # 红色1: (0, 1, 0, 0.5), # 绿色2: (0, 0, 1, 0.5), # 蓝色3: (1, 1, 0, 0.5), # 黄色4: (1, 0, 1, 0.5), # 紫色}for line in lines:parts = line.strip().split()if len(parts) < 3: # 至少类别+1个点continueclass_id = int(parts[0])# 将相对坐标转换为绝对坐标points = np.array([float(x) for x in parts[1:]]).reshape(-1, 2)points[:, 0] *= width # x坐标points[:, 1] *= height # y坐标# 创建多边形polygon = Polygon(points, closed=True)patches.append(polygon)# 为不同类别分配颜色color = color_map.get(class_id % len(color_map), (0.5, 0.5, 0.5, 0.5))colors.append(color)# 添加所有多边形到图中p = PatchCollection(patches, alpha=0.4)p.set_color(colors)ax.add_collection(p)# 添加图例from matplotlib.lines import Line2Dlegend_elements = [Line2D([0], [0], marker='o', color='w', label=f'Class {k}',markerfacecolor=v[:3], markersize=10) for k, v in color_map.items()]ax.legend(handles=legend_elements, loc='upper right')plt.axis('off')if output_path:plt.savefig(output_path, bbox_inches='tight', pad_inches=0)print(f"可视化结果已保存到 {output_path}")else:plt.show()# 使用示例
image_path = "test.jpg" # 替换为你的图像路径
label_path = "test.txt" # 替换为你的标签路径
output_path = "test.png" # 可选: 输出文件路径visualize_masks(image_path, label_path, output_path)