# -*- coding:utf-8 -*-
# @Author len
# @Create 2023/12/25 14:55

import cv2
import numpy as np
import random
import os

def process_image(bg_path, logo_path, output_path):
    background = cv2.imread(bg_path)
    logo = cv2.imread(logo_path, -1)  # -1 表示尝试读取透明度通道

    if logo is None or background is None:
        print(f"无法读取图像，跳过 {bg_path} 和 {logo_path}")
        return

    # 检查标志图像是否有透明度通道
    if logo.shape[2] == 4:
        alpha_channel = logo[:, :, 3]
    else:
        alpha_channel = np.ones(logo.shape[:2], dtype=logo.dtype) * 255

    # 确保标志图像小于背景图像
    if logo.shape[1] > background.shape[1] or logo.shape[0] > background.shape[0]:
        scale_factor = min(background.shape[1] / logo.shape[1], background.shape[0] / logo.shape[0]) * 0.9
    else:
        scale_factor = random.uniform(0.5, 1.5)

    width = int(logo.shape[1] * scale_factor)
    height = int(logo.shape[0] * scale_factor)
    logo = cv2.resize(logo, (width, height))
    alpha_channel = cv2.resize(alpha_channel, (width, height))

    # 确保随机位置的计算是有效的
    x_offset = random.randint(0, max(0, background.shape[1] - width))
    y_offset = random.randint(0, max(0, background.shape[0] - height))

    # 合成图像
    for c in range(0, 3):
        background[y_offset:y_offset+height, x_offset:x_offset+width, c] = \
            logo[:, :, c] * (alpha_channel / 255.0) + \
            background[y_offset:y_offset+height, x_offset:x_offset+width, c] * (1.0 - alpha_channel / 255.0)

    # 保存结果到文件
    cv2.imwrite(output_path, background)


# 文件夹路径
backgrounds_path = r'D:\Waste\嵌入式\数据集\交通标志\1225交通标志\背景'  # 替换为你的背景图片文件夹路径
logos_path = r'D:\Waste\嵌入式\数据集\交通标志\1225交通标志\sign'  # 替换为你的标志图片文件夹路径
output_path = r'D:\Waste\嵌入式\数据集\交通标志\1225交通标志\生成'  # 替换为你的输出文件夹路径


# 确保输出路径存在
if not os.path.exists(output_path):
    os.makedirs(output_path)

backgrounds = [f for f in os.listdir(backgrounds_path) if f.endswith(('.png', '.jpg', '.jpeg'))]
logos = [f for f in os.listdir(logos_path) if f.endswith(('.png', '.jpg', '.jpeg'))]

# 为每个背景图片和每个标志图片组合
for bg_file in backgrounds:
    for logo_file in logos:
        bg_full_path = os.path.join(backgrounds_path, bg_file)
        logo_full_path = os.path.join(logos_path, logo_file)
        output_full_path = os.path.join(output_path, f'combined_{bg_file[:-4]}_{logo_file}')
        process_image(bg_full_path, logo_full_path, output_full_path)

print("批量处理完成！")