import cv2
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

# 读取图像
image_path = 'jigsaw_mask.png'  # 替换为您的图像路径
image = cv2.imread(image_path)

if image is None:
    raise ValueError("无法读取图像，请检查路径是否正确。")

# 转换为RGB
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 转换为灰度图
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

# 应用固定阈值进行二值化
_, mask = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY_INV)

# 进行形态学操作以连接断裂的线条
# 定义较大的核
kernel = np.ones((1,1), np.uint8)

# 增加膨胀次数以连接断裂
mask = cv2.dilate(mask, kernel, iterations=1)

# 闭运算进一步连接
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=5)

# 再次确保掩膜是二值的
_, mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)

# 创建 Alpha 通道
alpha = cv2.bitwise_not(mask)

# 确保 Alpha 通道只有0和255
alpha = cv2.threshold(alpha, 127, 255, cv2.THRESH_BINARY)[1]

# 创建纯白色背景
white_bg = 255 * np.ones_like(image_rgb, dtype=np.uint8)

# 合并 RGB 图像和 Alpha 通道
final_rgb = cv2.bitwise_and(white_bg, white_bg, mask=cv2.bitwise_not(mask))
image_rgba = cv2.cvtColor(final_rgb, cv2.COLOR_RGB2RGBA)
image_rgba[:, :, 3] = alpha

# 将 OpenCV 的 RGBA 图像转换为 Pillow Image 对象
image_pil = Image.fromarray(image_rgba)

# 转换为 NumPy 数组以进一步处理
data = np.array(image_pil)

# 分离颜色通道
red, green, blue, alpha = data.T

# 定义黑色的阈值范围
black_threshold = 50  # 根据需要调整

# 创建掩膜：黑色区域为True，其余为False
black_areas = (red < black_threshold) & (green < black_threshold) & (blue < black_threshold)

# 设置黑色区域的 Alpha 为0，非黑色区域的 Alpha 为255
data[..., 3][black_areas.T] = 0
data[..., 3][~black_areas.T] = 255

# 将非黑色区域设置为纯白色
data[..., 0:3][~black_areas.T] = [255, 255, 255]

# 创建新的 Image 对象
final_image = Image.fromarray(data)

# 保存结果为 PNG（支持透明度）
output_path = 'mask.png'
final_image.save(output_path)

print(f"已保存带纯白背景和透明部分的图像到 {output_path}")

# 可选：显示图像
plt.figure(figsize=(12,6))
plt.subplot(1,2,1)
plt.title('原图')
plt.imshow(image_rgb)
plt.axis('off')

plt.subplot(1,2,2)
plt.title('带纯白背景和透明部分的图像')
plt.imshow(final_image)
plt.axis('off')

plt.show()