import os
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt

# 设置图像和标签的目录
images_dir = 'dataset/enhance_yawn/images'
labels_dir = 'dataset/enhance_yawn/labels'
'''
显示image和对应的标记点
'''
# 选择一个图像文件
for image_file in os.listdir(images_dir):
    image_path = os.path.join(images_dir, image_file)

    # 找到对应的标签文件
    label_file = os.path.splitext(image_file)[0] + '.txt'
    label_path = os.path.join(labels_dir, label_file)

    # 读取图像
    image = Image.open(image_path)
    draw = ImageDraw.Draw(image)

    # 读取标签文件并绘制标注框
    if os.path.exists(label_path):
        with open(label_path, 'r') as file:
            for line in file.readlines():
                # YOLO格式: class x_center y_center width height
                class_id, x_center, y_center, width, height = map(float, line.split()[:5])
                # 转换为图像坐标系
                img_width, img_height = image.size
                x_center *= img_width
                y_center *= img_height
                width *= img_width
                height *= img_height
                x_min = x_center - width / 2
                y_min = y_center - height / 2
                # 绘制矩形框
                draw.rectangle(((x_min, y_min), (x_min + width, y_min + height)), outline='red', width=2)

                keypoints = line.split()[5:]  # 剩余的部分是关键点数据
                for i in range(0, len(keypoints), 3):
                    x, y, v = map(float, keypoints[i:i + 3])
                    if v != 1:  # 如果关键点可见
                        x *= img_width
                        y *= img_height
                        draw.ellipse((x - 3, y - 3, x + 3, y + 3), fill='blue')
            image.show()