# -*- coding:utf-8 -*-
# @Author len
# @Create 2023/10/31 20:24


from PIL import Image
import os
import random

# 设置路径
background_folder = r'D:\Waste\嵌入式\数据集\背景'
overlay_image_path = r'D:\Waste\嵌入式\数据集\交通标志\禁止通行.png'
output_folder = r'D:\Waste\嵌入式\数据集\交通标志\002\禁止通行'

# 读取叠加图片
overlay = Image.open(overlay_image_path)
overlay = overlay.convert("RGBA")

# 确保输出目录存在
if not os.path.exists(output_folder):
    os.makedirs(output_folder)

# 遍历背景图片文件夹中的所有文件
for background_image_name in os.listdir(background_folder):
    if background_image_name.endswith(('.png', '.jpg', '.jpeg')):  # 检查文件扩展名
        background_path = os.path.join(background_folder, background_image_name)
        background = Image.open(background_path)
        background = background.convert("RGBA")

        # 随机选择一个缩放因子在90%到110%之间
        scale_factor = random.uniform(0.1, 0.5)

        # 缩放overlay图片
        new_size = (int(overlay.width * scale_factor), int(overlay.height * scale_factor))
        scaled_overlay = overlay.resize(new_size, Image.Resampling.LANCZOS)

        # 随机选择overlay的位置
        max_x = background.width - scaled_overlay.width
        max_y = background.height - scaled_overlay.height
        if max_x < 0 or max_y < 0:
            print(f"Scaled overlay image is larger than background {background_image_name}. Skipping.")
            continue
        rand_x = random.randint(0, max_x)
        rand_y = random.randint(0, max_y)

        # 将叠加图片放在背景图片上的随机位置
        background.paste(scaled_overlay, (rand_x, rand_y), scaled_overlay)

        # 将结果保存为新图片
        output_path = os.path.join(output_folder, background_image_name)
        # 如果输出文件格式要求为JPEG，确保转换为RGB模式
        if output_path.lower().endswith('.jpg') or output_path.lower().endswith('.jpeg'):
            background = background.convert('RGB')
        background.save(output_path)

print("所有图片处理完成。")