115 lines
3.4 KiB
Python
115 lines
3.4 KiB
Python
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') |