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)