import cv2
import numpy as np
import matplotlib.pyplot as plt


def add_watermark(image_path, watermark_path):
    # 读取原始图像和水印图像
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    watermark = cv2.imread(watermark_path, cv2.IMREAD_GRAYSCALE)
    watermark = cv2.resize(watermark, (image.shape[1], image.shape[0]))

    # 将原始图像和水印图像转换到频域
    f_image = np.fft.fft2(image)
    f_image_shifted = np.fft.fftshift(f_image)
    f_watermark = np.fft.fft2(watermark)

    # 嵌入水印
    alpha = 0.1  # 水印强度因子
    f_image_shifted += alpha * f_watermark

    # 将嵌入水印后的频域表示转换回空间域
    f_ishift = np.fft.ifftshift(f_image_shifted)
    img_back = np.fft.ifft2(f_ishift)
    img_back = np.abs(img_back)

    # 提取水印
    f_extracted = f_image_shifted - np.fft.fftshift(f_image)
    f_extracted_shifted = np.fft.ifftshift(f_extracted / alpha)
    watermark_extracted = np.fft.ifft2(f_extracted_shifted)
    watermark_extracted = np.abs(watermark_extracted)

    # 显示图像
    plt.figure(figsize=(10, 8))
    plt.subplot(221), plt.imshow(image, cmap='gray'), plt.title('Original Image')
    plt.subplot(222), plt.imshow(watermark, cmap='gray'), plt.title('Watermark')
    plt.subplot(223), plt.imshow(img_back, cmap='gray'), plt.title('Image with Watermark')
    plt.subplot(224), plt.imshow(watermark_extracted, cmap='gray'), plt.title('Extracted Watermark')
    plt.show()


# 使用示例
add_watermark(r'c:lena.BMP', r'c:encrypted_image1.png')

                                                            实验结果