import os
import numpy as np
from PIL import Image
import scipy.ndimage as ndimage
from scipy.ndimage import label, find_objects
import random


def create_segment_mask(alpha_channel):
    """
    创建分割掩码，找出被透明线条分隔的区域
    """
    binary_mask = alpha_channel < 255  # 透明区域
    labeled_array, num_features = label(~binary_mask)

    segment_masks = []
    for i in range(1, num_features + 1):
        segment_mask = labeled_array == i
        segment_masks.append(segment_mask)

    return segment_masks


def extract_segment_with_transparent_background(image, segment_mask):
    """
    提取分割区域，区域外设置为透明
    """
    slices = find_objects(segment_mask.astype(int))[0]

    y_slice = slice(max(0, slices[0].start - 5), min(image.shape[0], slices[0].stop + 5))
    x_slice = slice(max(0, slices[1].start - 5), min(image.shape[1], slices[1].stop + 5))

    segment_image = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
    segment_image[y_slice, x_slice][segment_mask[y_slice, x_slice]] = image[y_slice, x_slice][
        segment_mask[y_slice, x_slice]]

    return segment_image


def smart_crop_and_rotate(image, border_size=15):
    """
    智能裁剪和随机旋转，保留指定像素的透明边框
    """
    # 将numpy数组转换为PIL图像
    pil_image = Image.fromarray(image)

    # 找到非透明区域的边界框
    bbox = pil_image.getbbox()

    if bbox:
        # 扩大裁剪区域，保留指定像素的透明边框
        left = max(0, bbox[0] - border_size)
        upper = max(0, bbox[1] - border_size)
        right = min(image.shape[1], bbox[2] + border_size)
        lower = min(image.shape[0], bbox[3] + border_size)

        # 裁剪
        cropped_image = pil_image.crop((left, upper, right, lower))

        # # 随机旋转角度（-45到45度）
        rotation_angle = random.uniform(-180, 180)
        rotated_image = cropped_image.rotate(rotation_angle,
                                             expand=True,
                                             fillcolor=(0, 0, 0, 0))

        return np.array(rotated_image)

    return image


def main(input_path, output_folder):
    """
    主处理函数
    """
    # 创建输出文件夹
    os.makedirs(output_folder, exist_ok=True)

    # 随机数种子（保证每次运行结果可复现）
    random.seed()

    # 打开图像
    image = Image.open(input_path)

    # 转换为RGBA模式
    image = image.convert("RGBA")

    # 转换为numpy数组
    img_array = np.array(image)

    # 获取透明通道
    alpha_channel = img_array[:, :, 3]

    # 创建分割掩码
    segment_masks = create_segment_mask(alpha_channel)

    # 保存每个分割区域
    for i, segment_mask in enumerate(segment_masks):
        # 提取分割区域，区域外设置为透明
        segment_image_array = extract_segment_with_transparent_background(img_array, segment_mask)

        # 裁剪并随机旋转，保留20像素透明边框
        processed_segment = smart_crop_and_rotate(segment_image_array, border_size=20)

        # 从numpy数组创建图像
        segment_image = Image.fromarray(processed_segment)

        # 保存图像
        output_path = os.path.join(output_folder, f'segment_{i + 1}.png')
        segment_image.save(output_path)

    print(f"分割、裁剪和旋转完成，共生成 {len(segment_masks)} 个图像片段")


# 执行分割
main('result.png', 'output10')