import cv2
import os


def draw_boxes(image, labels, w, h):
    for line in labels:
        parts = line.strip().split()
        class_id = int(parts[0])
        x_center = float(parts[1])
        y_center = float(parts[2])
        width = float(parts[3])
        height = float(parts[4])

        # 计算边框的左上角和右下角坐标
        x_center *= w
        y_center *= h
        width *= w
        height *= h
        x_min = int(x_center - width / 2)
        y_min = int(y_center - height / 2)
        x_max = int(x_center + width / 2)
        y_max = int(y_center + height / 2)

        # 画出边框
        color = (0, 255, 0)  # 绿色边框
        thickness = 10
        image = cv2.rectangle(image, (x_min, y_min), (x_max, y_max), color, thickness)
    return image


def process_folder(image_folder, label_folder, output_folder):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    image_files = [f for f in os.listdir(image_folder) if f.endswith(('.jpg', '.jpeg', '.png'))]

    for image_file in image_files:
        image_path = os.path.join(image_folder, image_file)
        label_file = os.path.splitext(image_file)[0] + '.txt'
        label_path = os.path.join(label_folder, label_file)

        if os.path.exists(label_path):
            # 读取图片
            image = cv2.imread(image_path)
            h, w, _ = image.shape

            # 读取标签文件
            with open(label_path, 'r') as file:
                labels = file.readlines()

            # 绘制边框
            image_with_boxes = draw_boxes(image, labels, w, h)

            # 保存结果图片
            output_path = os.path.join(output_folder, image_file)
            cv2.imwrite(output_path, image_with_boxes)
            print(f'Processed {image_file}')
        else:
            print(f'No label file found for {image_file}')


# 示例用法
image_folder = 'testTiaoJuan/cutJPGOnlyTiaoJuan'
label_folder = 'testTiaoJuan/CutedLabels_tiaojuan'
output_folder = 'runs/test/cuted_conf0.035'

process_folder(image_folder, label_folder, output_folder)