Initial commit

mask_split_files.py

@@ -0,0 +1,115 @@
+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')