yolo格式分割标签可视化,coco-seg数据集
import cv2
import numpy as np
import os
def visualize_mask_on_image(image_path, label_path, output_path):
# 读取图像
image = cv2.imread(image_path)
height, width, _ = image.shape
# 读取标签文件
try:
with open(label_path, 'r') as f:
lines = f.readlines()
except FileNotFoundError:
print(f"Label file {label_path} not found. Skipping...")
return
for line in lines:
parts = line.strip().split()
class_id = int(parts[0])
# 提取归一化的坐标
normalized_coords = [float(x) for x in parts[1:]]
# 将归一化坐标转换为实际坐标
coords = []
for i in range(0, len(normalized_coords), 2):
x = int(normalized_coords[i] * width)
y = int(normalized_coords[i + 1] * height)
coords.append([x, y])
# 将坐标转换为 numpy 数组
pts = np.array(coords, np.int32)
pts = pts.reshape((-1, 1, 2))
# 随机生成颜色
color = tuple(np.random.randint(0, 255, 3).tolist())
# 在图像上绘制掩码
cv2.polylines(image, [pts], True, color, 2)
cv2.fillPoly(image, [pts], color)
# 保存图像
cv2.imwrite(output_path, image)
def process_images(image_dir, label_dir, output_dir):
# 确保输出目录存在
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# 遍历图像目录中的所有文件
for filename in os.listdir(image_dir):
if filename.endswith(('.jpg', '.jpeg', '.png')):
image_path = os.path.join(image_dir, filename)
# 构造对应的标签文件路径
label_filename = os.path.splitext(filename)[0] + '.txt'
label_path = os.path.join(label_dir, label_filename)
# 构造输出文件路径
output_path = os.path.join(output_dir, filename)
visualize_mask_on_image(image_path, label_path, output_path)
# 请替换为实际的图像、标签和输出目录路径
image_dir = 'coco128-seg/images/train2017/'
label_dir = 'labels/train2017/'
output_dir = 'results'
process_images(image_dir, label_dir, output_dir)