import cv2
import numpy as np
from PIL import Image
import os

def split_image_with_template(image_path, template_path, output_dir):
    """
    根据拼图模板，将图片裁剪成拼图块。
    :param image_path: 输入图片路径
    :param template_path: 拼图模板路径
    :param output_dir: 输出拼图块保存目录
    """
    # 加载图片和模板
    img = cv2.imread(image_path)
    template = cv2.imread(template_path, cv2.IMREAD_GRAYSCALE)

    # 确保模板是二值化图像
    _, mask = cv2.threshold(template, 127, 255, cv2.THRESH_BINARY)

    # 查找轮廓
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)

    # 裁剪每块拼图
    for i, contour in enumerate(contours):
        # 创建与原始图片相同大小的掩码
        piece_mask = np.zeros_like(mask)
        cv2.drawContours(piece_mask, [contour], -1, 255, thickness=-1)

        # 裁剪区域
        x, y, w, h = cv2.boundingRect(contour)
        piece = cv2.bitwise_and(img, img, mask=piece_mask)
        piece_cropped = piece[y:y + h, x:x + w]

        # 保存拼图块
        piece_path = os.path.join(output_dir, f"piece_{i}.png")
        cv2.imwrite(piece_path, piece_cropped)

    print(f"拼图已生成，共 {len(contours)} 块，保存到：{output_dir}")


# 示例调用
image_path = "input.jpg"  # 原始图片路径
template_path = "mask.png"  # 拼图模板路径
output_dir = "output_jigsaw"  # 输出路径
split_image_with_template(image_path, template_path, output_dir)