import os
from PIL import Image
import random

def load_images_from_folder(folder_path):
    images = []
    for filename in os.listdir(folder_path):
        if filename.endswith(".jpg") or filename.endswith(".png"):
            image_path = os.path.join(folder_path, filename)
            image = Image.open(image_path).convert("RGBA")
            images.append(image)
    return images

def check_overlap(positions, new_position, image_size):
    for position in positions:
        x1, y1, x2, y2 = position
        new_x1, new_y1 = new_position
        new_x2 = new_x1 + image_size[0]
        new_y2 = new_y1 + image_size[1]
        if (new_x1 < x2 and new_x2 > x1) and (new_y1 < y2 and new_y2 > y1):
            return True
    return False

def place_images_on_background(background_image, png_images, num_images):
    bg_width, bg_height = background_image.size
    positions = []
    count = 0

    chongdie = 0

    while count < num_images:
        png_image = random.choice(png_images)
        # 旋转图像
        png_image = png_image.rotate(random.randint(1, 360), expand=True)

        png_width, png_height = png_image.size
        image_size = (png_width, png_height)

        # 尝试找到不重叠的位置
        overlapped = True
        number = 0
        while overlapped:
            number += 1
            # 随机生成位置
            x = random.randint(0, bg_width - png_width)
            y = random.randint(0, bg_height - png_height)
            new_position = (x, y)

            # 检查是否与已有图片重叠
            overlapped = check_overlap(positions, new_position, image_size)
            if number == 3:
                chongdie += 1
                overlapped = False

        # 将PNG图片粘贴到背景图片上
        background_image.paste(png_image, new_position, png_image)
        positions.append((x, y, x + png_width, y + png_height))
        count += 1

    return background_image, chongdie

background_folder_path = "data/background"
png_folder_path = "data/pro"
output_folder_path = "data/result"
num_images = 6



os.makedirs(output_folder_path, exist_ok=True)  # 创建输出文件夹

index = 99
sum = 0
while sum <= 300:
    sum += 1
    background_images = load_images_from_folder(background_folder_path)
    png_images = load_images_from_folder(png_folder_path)
    for i in range(len(background_images)):
        background_image = background_images[i]
        result_image, chongdie = place_images_on_background(background_image, png_images, num_images)

        if chongdie >= 1:
            continue

        # 将图像转换为RGB模式
        result_image = result_image.convert("RGB")

        # 保存结果图片为JPEG格式
        result_image.save(f"{output_folder_path}/result{index}.jpg", "JPEG")
        index += 1
    print(index)