import os
from PIL import Image, ImageDraw

def draw_centered_red_box(image, percentage=0.8, color=(255, 0, 0)):
    """
    在图片上绘制一个红色框，框住指定比例的面积。
    image: PIL Image 对象
    percentage: 正方形框的面积占图片面积的比例（0-1之间）
    color: 框的颜色 (R, G, B)
    """
    w, h = image.size
    # 计算正方形的边长
    box_size = int((w * h * percentage) ** 0.5)
    half_box = box_size // 2

    # 确定框的中心点
    center_x, center_y = w // 2, h // 2

    # 计算正方形的边界
    left = max(center_x - half_box, 0)
    top = max(center_y - half_box, 0)
    right = min(center_x + half_box, w)
    bottom = min(center_y + half_box, h)

    # 在图片上绘制红色框
    draw = ImageDraw.Draw(image)
    draw.rectangle([left, top, right, bottom], outline=color, width=3)  # 边框宽度为 3 像素

    return image

# 文件夹路径
input_folder = "output10"
output_folder = "output10_with_red_box"
os.makedirs(output_folder, exist_ok=True)

# 遍历文件夹中的图片
for filename in os.listdir(input_folder):
    if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
        file_path = os.path.join(input_folder, filename)
        img = Image.open(file_path)

        # 添加红色框
        img_with_box = draw_centered_red_box(img, percentage=0.8)

        # 保存结果
        output_path = os.path.join(output_folder, filename)
        img_with_box.save(output_path)

print("处理完成，已保存带框的图片到 'output10_with_red_box' 文件夹。")