Initial commit

extract_mask.py

@@ -0,0 +1,92 @@
+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()