import numpy as np
from scipy.io import wavfile
import matplotlib.pyplot as plt

def embed_watermark(audio, watermark, alpha=0.1):
    """
    在音頻信號中嵌入水印
    :param audio: 原始音頻信號
    :param watermark: 水印信號
    :param alpha: 水印強度
    :return: 嵌入水印后的音頻信號
    """
    return audio + alpha * watermark

def extract_watermark(watermarked_audio, original_audio, alpha=0.1):
    """
    從音頻信號中提取水印
    :param watermarked_audio: 嵌入水印后的音頻信號
    :param original_audio: 原始音頻信號
    :param alpha: 水印強度
    :return: 提取的水印信號
    """
    return (watermarked_audio - original_audio) / alpha

def main():
    # 讀取原始音頻文件
    sample_rate, audio = wavfile.read('original_audio.wav')

    # 生成水印信號
    watermark = np.random.randn(len(audio))

    # 嵌入水印
    watermarked_audio = embed_watermark(audio, watermark)

    # 保存嵌入水印后的音頻文件
    wavfile.write('watermarked_audio.wav', sample_rate, watermarked_audio.astype(np.int16))

    # 提取水印
    extracted_watermark = extract_watermark(watermarked_audio, audio)

    # 繪制原始音頻和提取的水印
    plt.figure(figsize=(12, 6))
    plt.subplot(2, 1, 1)
    plt.plot(audio)
    plt.title('Original Audio')
    plt.subplot(2, 1, 2)
    plt.plot(extracted_watermark)
    plt.title('Extracted Watermark')
    plt.show()

if __name__ == "__main__":
    main()

import logging

logging.basicConfig(level=logging.INFO)

def read_audio(file_path):
    try:
        sample_rate, audio = wavfile.read(file_path)
        return sample_rate, audio
    except Exception as e:
        logging.error(f"Error reading audio file: {e}")
        return None, None

def write_audio(file_path, sample_rate, audio):
    try:
        wavfile.write(file_path, sample_rate, audio.astype(np.int16))
    except Exception as e:
        logging.error(f"Error writing audio file: {e}")

def main():
    # 讀取原始音頻文件
    sample_rate, audio = read_audio('original_audio.wav')
    if audio is None:
        return

    # 生成水印信號
    watermark = np.random.randn(len(audio))

    # 嵌入水印
    watermarked_audio = embed_watermark(audio, watermark)

    # 保存嵌入水印后的音頻文件
    write_audio('watermarked_audio.wav', sample_rate, watermarked_audio)

    # 提取水印
    extracted_watermark = extract_watermark(watermarked_audio, audio)

    # 繪制原始音頻和提取的水印
    plt.figure(figsize=(12, 6))
    plt.subplot(2, 1, 1)
    plt.plot(audio)
    plt.title('Original Audio')
    plt.subplot(2, 1, 2)
    plt.plot(extracted_watermark)
    plt.title('Extracted Watermark')
    plt.show()

if __name__ == "__main__":
    main()