jigsaw/04extract_borders.py

import os
from PIL import Image
import numpy as np


def get_edge_opaque_pixels_with_limit(image_path):
    """
    提取图片中每一行和每一列最左、最右、最上、最下的完全不透明像素（alpha == 255），
    并根据位置限制过滤。

    Args:
        image_path (str): 图片文件路径。

    Returns:
        dict: 包含四个边的像素信息列表。
    """
    with Image.open(image_path).convert("RGBA") as img:
        width, height = img.size
        pixels = img.load()
        limit_left = width / 3
        limit_right = 2 * width / 3
        limit_top = height / 3
        limit_bottom = 2 * height / 3

        left_pixels = []
        right_pixels = []
        top_pixels = []
        bottom_pixels = []

        # 左边
        for y in range(height):
            for x in range(width):
                r, g, b, a = pixels[x, y]
                if a == 255 and x <= limit_left:
                    left_pixels.append({
                        "position": (x, y),
                        "color": (r, g, b, a)
                    })
                    break

        # 右边
        for y in range(height):
            for x in range(width - 1, -1, -1):
                r, g, b, a = pixels[x, y]
                if a == 255 and x >= limit_right:
                    right_pixels.append({
                        "position": (x, y),
                        "color": (r, g, b, a)
                    })
                    break

        # 上边
        for x in range(width):
            for y in range(height):
                r, g, b, a = pixels[x, y]
                if a == 255 and y <= limit_top:
                    top_pixels.append({
                        "position": (x, y),
                        "color": (r, g, b, a)
                    })
                    break

        # 下边
        for x in range(width):
            for y in range(height - 1, -1, -1):
                r, g, b, a = pixels[x, y]
                if a == 255 and y >= limit_bottom:
                    bottom_pixels.append({
                        "position": (x, y),
                        "color": (r, g, b, a)
                    })
                    break

    return {
        "left": left_pixels,
        "right": right_pixels,
        "top": top_pixels,
        "bottom": bottom_pixels
    }


def get_edge_opaque_pixels_with_limit_numpy(image_path):
    """
    使用NumPy提取图片中每一行和每一列最左、最右、最上、最下的完全不透明像素（alpha == 255），
    并根据位置限制过滤。

    Args:
        image_path (str): 图片文件路径。

    Returns:
        dict: 包含四个边的像素信息列表。
    """
    with Image.open(image_path).convert("RGBA") as img:
        data = np.array(img)
        height, width, _ = data.shape
        limit_left = width / 3
        limit_right = 2 * width / 3
        limit_top = height / 3
        limit_bottom = 2 * height / 3

        left_pixels = []
        right_pixels = []
        top_pixels = []
        bottom_pixels = []

        # 左边
        for y in range(height):
            row = data[y, :, :]
            opaque_indices = np.where(row[:, 3] == 255)[0]
            valid_indices = opaque_indices[opaque_indices <= limit_left]
            if valid_indices.size > 0:
                x = valid_indices[0]
                r, g, b, a = row[x]
                left_pixels.append({
                    "position": (int(x), y),
                    "color": (int(r), int(g), int(b), int(a))
                })

        # 右边
        for y in range(height):
            row = data[y, :, :]
            opaque_indices = np.where(row[:, 3] == 255)[0]
            valid_indices = opaque_indices[opaque_indices >= limit_right]
            if valid_indices.size > 0:
                x = valid_indices[-1]  # 最右边的第一个满足条件的像素
                r, g, b, a = row[x]
                right_pixels.append({
                    "position": (int(x), y),
                    "color": (int(r), int(g), int(b), int(a))
                })

        # 上边
        for x in range(width):
            column = data[:, x, :]
            opaque_indices = np.where(column[:, 3] == 255)[0]
            valid_indices = opaque_indices[opaque_indices <= limit_top]
            if valid_indices.size > 0:
                y = valid_indices[0]
                r, g, b, a = column[y]
                top_pixels.append({
                    "position": (x, int(y)),
                    "color": (int(r), int(g), int(b), int(a))
                })

        # 下边
        for x in range(width):
            column = data[:, x, :]
            opaque_indices = np.where(column[:, 3] == 255)[0]
            valid_indices = opaque_indices[opaque_indices >= limit_bottom]
            if valid_indices.size > 0:
                y = valid_indices[-1]  # 最下边的第一个满足条件的像素
                r, g, b, a = column[y]
                bottom_pixels.append({
                    "position": (x, int(y)),
                    "color": (int(r), int(g), int(b), int(a))
                })

    return {
        "left": left_pixels,
        "right": right_pixels,
        "top": top_pixels,
        "bottom": bottom_pixels
    }


def process_directory(directory_path, use_numpy=False):
    """
    遍历指定文件夹下所有图片文件，提取每张图片的边缘不透明像素信息。

    Args:
        directory_path (str): 文件夹路径。
        use_numpy (bool): 是否使用NumPy加速处理。

    Returns:
        dict: 每个文件对应的四个方向的像素信息。
    """
    supported_extensions = ('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tiff')
    result_map = {}

    for filename in os.listdir(directory_path):
        if filename.lower().endswith(supported_extensions):
            file_path = os.path.join(directory_path, filename)
            if use_numpy:
                edge_pixels = get_edge_opaque_pixels_with_limit_numpy(file_path)
            else:
                edge_pixels = get_edge_opaque_pixels_with_limit(file_path)
            result_map[filename] = edge_pixels

    return result_map


if __name__ == "__main__":
    directory_path = "output4"  # 替换为你的文件夹路径
    use_numpy = True  # 设置为True使用NumPy优化，False使用PIL
    result = process_directory(directory_path, use_numpy)