AI-based Audio Watermarking Tool
from .utils import wm_add_util, file_reader, wm_decode_util, my_parser, metric_util, path_util
from .models import my_model
import torch
import numpy as np
from huggingface_hub import hf_hub_download


def load_model(path="default"):
    if path == "default":
        resume_path = hf_hub_download(repo_id="M4869/WavMark",
                                      filename="step59000_snr39.99_pesq4.35_BERP_none0.30_mean1.81_std1.81.model.pkl",
                                      )

    model = my_model.Model(16000, num_bit=32, n_fft=1000, hop_length=400, num_layers=8)
    checkpoint = torch.load(resume_path, map_location=torch.device('cpu'))
    model_ckpt = checkpoint
    model.load_state_dict(model_ckpt, strict=True)
    model.eval()
    return model


def encode_watermark(model, signal, payload, pattern_bit_length=16, min_snr=20, max_snr=38, show_progress=False):
    device = next(model.parameters()).device

    pattern_bit = wm_add_util.fix_pattern[0:pattern_bit_length]

    watermark = np.concatenate([pattern_bit, payload])
    assert len(watermark) == 32
    signal_wmd, info = wm_add_util.add_watermark(watermark, signal, 16000, 0.1,
                                                 device, model, min_snr, max_snr,
                                                 show_progress=show_progress)
    info["snr"] = metric_util.signal_noise_ratio(signal, signal_wmd)
    return signal_wmd, info


def decode_watermark(model, signal, decode_batch_size=10, len_start_bit=16, show_progress=False):
    device = next(model.parameters()).device
    start_bit = wm_add_util.fix_pattern[0:len_start_bit]
    mean_result, info = wm_decode_util.extract_watermark_v3_batch(
        signal,
        start_bit,
        0.1,
        16000,
        model,
        device, decode_batch_size, show_progress=show_progress)

    if mean_result is None:
        return None, info

    payload = mean_result[len_start_bit:]
    return payload, info

src/wavmark/__init__.py
wavmark/wavmark
import torch
from ..models.invblock import INV_block


class Hinet(torch.nn.Module):

    def __init__(self, in_channel=2, num_layers=16):
        super(Hinet, self).__init__()
        self.inv_blocks = torch.nn.ModuleList([INV_block(in_channel) for _ in range(num_layers)])

    def forward(self, x1, x2, rev=False):
        # x1:cover
        # x2:secret
        if not rev:
            for inv_block in self.inv_blocks:
                x1, x2 = inv_block(x1, x2)
        else:
            for inv_block in reversed(self.inv_blocks):
                x1, x2 = inv_block(x1, x2, rev=True)
        return x1, x2

src/wavmark/models/hinet.py
wavmark/wavmark
import torch
import torch.nn as nn
from ..models.rrdb_denselayer import ResidualDenseBlock_out


class INV_block(nn.Module):
    def __init__(self, channel=2, subnet_constructor=ResidualDenseBlock_out, clamp=2.0):
        super().__init__()
        self.clamp = clamp

        # ρ
        self.r = subnet_constructor(channel, channel)
        # η
        self.y = subnet_constructor(channel, channel)
        # φ
        self.f = subnet_constructor(channel, channel)

    def e(self, s):
        return torch.exp(self.clamp * 2 * (torch.sigmoid(s) - 0.5))

    def forward(self, x1, x2, rev=False):
        if not rev:

            t2 = self.f(x2)
            y1 = x1 + t2
            s1, t1 = self.r(y1), self.y(y1)
            y2 = self.e(s1) * x2 + t1

        else:

            s1, t1 = self.r(x1), self.y(x1)
            y2 = (x2 - t1) / self.e(s1)
            t2 = self.f(y2)
            y1 = (x1 - t2)

        return y1, y2

src/wavmark/models/invblock.py
wavmark/wavmark
import torch.optim
import torch.nn as nn
from ..models.hinet import Hinet
import numpy as np
import random


class Model(nn.Module):
    def __init__(self, num_point, num_bit, n_fft, hop_length, num_layers):
        super(Model, self).__init__()
        self.hinet = Hinet(num_layers=num_layers)
        self.watermark_fc = torch.nn.Linear(num_bit, num_point)
        self.watermark_fc_back = torch.nn.Linear(num_point, num_bit)
        self.n_fft = n_fft
        self.hop_length = hop_length

    def stft(self, data):
        window = torch.hann_window(self.n_fft).to(data.device)
        tmp = torch.stft(data, n_fft=self.n_fft, hop_length=self.hop_length, window=window, return_complex=False)
        # [1, 501, 41, 2]
        return tmp

    def istft(self, signal_wmd_fft):
        window = torch.hann_window(self.n_fft).to(signal_wmd_fft.device)
        return torch.istft(signal_wmd_fft, n_fft=self.n_fft, hop_length=self.hop_length, window=window,
                           return_complex=False)

    def encode(self, signal, message):
        signal_fft = self.stft(signal)
        # (batch,freq_bins,time_frames,2)

        message_expand = self.watermark_fc(message)
        message_fft = self.stft(message_expand)

        signal_wmd_fft, msg_remain = self.enc_dec(signal_fft, message_fft, rev=False)
        # (batch,freq_bins,time_frames,2)
        signal_wmd = self.istft(signal_wmd_fft)
        return signal_wmd

    def decode(self, signal):
        signal_fft = self.stft(signal)
        watermark_fft = signal_fft
        _, message_restored_fft = self.enc_dec(signal_fft, watermark_fft, rev=True)
        message_restored_expanded = self.istft(message_restored_fft)
        message_restored_float = self.watermark_fc_back(message_restored_expanded).clamp(-1, 1)
        return message_restored_float

    def enc_dec(self, signal, watermark, rev):
        signal = signal.permute(0, 3, 2, 1)
        # [4, 2, 41, 501]
        watermark = watermark.permute(0, 3, 2, 1)
        signal2, watermark2 = self.hinet(signal, watermark, rev)
        return signal2.permute(0, 3, 2, 1), watermark2.permute(0, 3, 2, 1)

src/wavmark/models/my_model.py
wavmark/wavmark
import torch
import torch.nn as nn
from ..models import module_util as mutil


# Dense connection
class ResidualDenseBlock_out(nn.Module):
    def __init__(self, in_channel, out_channel, bias=True):
        super(ResidualDenseBlock_out, self).__init__()
        self.conv1 = nn.Conv2d(in_channel, 32, 3, 1, 1, bias=bias)
        self.conv2 = nn.Conv2d(in_channel + 32, 32, 3, 1, 1, bias=bias)
        self.conv3 = nn.Conv2d(in_channel + 2 * 32, 32, 3, 1, 1, bias=bias)
        self.conv4 = nn.Conv2d(in_channel + 3 * 32, 32, 3, 1, 1, bias=bias)
        self.conv5 = nn.Conv2d(in_channel + 4 * 32, out_channel, 3, 1, 1, bias=bias)
        self.lrelu = nn.LeakyReLU(inplace=True)
        # initialization
        mutil.initialize_weights([self.conv5], 0.)

    def forward(self, x):
        x1 = self.lrelu(self.conv1(x))
        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
        return x5

src/wavmark/models/rrdb_denselayer.py
wavmark/wavmark
import os
import soundfile
import librosa
import resampy
import numpy as np


def is_wav_file(filename):
    file_extension = os.path.splitext(filename)[1]
    return file_extension.lower() == ".wav"


def read_as_single_channel_16k(audio_file, def_sr=16000, verbose=True, aim_second=None):
    assert os.path.exists(audio_file)

    file_extension = os.path.splitext(audio_file)[1].lower()

    if file_extension == ".mp3":
        data, origin_sr = librosa.load(audio_file, sr=None)
    elif file_extension in [".wav", ".flac"]:
        data, origin_sr = soundfile.read(audio_file)
    else:
        raise Exception("unsupported file:" + file_extension)

    # channel check
    if len(data.shape) == 2:
        left_channel = data[:, 0]
        if verbose:
            print("Warning! the input audio has multiple chanel, this tool only use the first channel!")
        data = left_channel

    # sample rate check
    if origin_sr != def_sr:
        data = resampy.resample(data, origin_sr, def_sr)
        if verbose:
            print("Warning! The original samplerate is not 16Khz; the watermarked audio will be re-sampled to 16KHz")

    sr = def_sr
    audio_length_second = 1.0 * len(data) / sr
    # if verbose:
    #     print("input length :%d second" % audio_length_second)

    if aim_second is not None:
        signal = data
        assert len(signal) > 0
        current_second = len(signal) / sr
        if current_second < aim_second:
            repeat_count = int(aim_second / current_second) + 1
            signal = np.repeat(signal, repeat_count)
        data = signal[0:sr * aim_second]

    return data, sr, audio_length_second


def read_as_single_channel(file, aim_sr):
    if file.endswith(".mp3"):
        data, sr = librosa.load(file, sr=aim_sr)
    else:
        data, sr = soundfile.read(file)

    if len(data.shape) == 2:  # multi-channel
        data = data[:, 0]  # only use the first channel

    if sr != aim_sr:
        data = resampy.resample(data, sr, aim_sr)
    return data

src/wavmark/utils/file_reader.py
wavmark/wavmark
import torch
import numpy as np


def calc_ber(watermark_decoded_tensor, watermark_tensor, threshold=0.5):
    watermark_decoded_binary = watermark_decoded_tensor >= threshold
    watermark_binary = watermark_tensor >= threshold
    ber_tensor = 1 - (watermark_decoded_binary == watermark_binary).to(torch.float32).mean()
    return ber_tensor


def to_equal_length(original, signal_watermarked):
    if original.shape != signal_watermarked.shape:
        print("Warning: length not equal:", len(original), len(signal_watermarked))
        min_length = min(len(original), len(signal_watermarked))
        original = original[0:min_length]
        signal_watermarked = signal_watermarked[0:min_length]
    assert original.shape == signal_watermarked.shape
    return original, signal_watermarked


def signal_noise_ratio(original, signal_watermarked):
    original, signal_watermarked = to_equal_length(original, signal_watermarked)
    noise_strength = np.sum((original - signal_watermarked) ** 2)
    if noise_strength == 0:  #
        return np.inf
    signal_strength = np.sum(original ** 2)
    ratio = signal_strength / noise_strength
    ratio = max(1e-10, ratio)
    return 10 * np.log10(ratio)


def batch_signal_noise_ratio(original, signal_watermarked):
    signal = original.detach().cpu().numpy()
    signal_watermarked = signal_watermarked.detach().cpu().numpy()
    tmp_list = []
    for s, swm in zip(signal, signal_watermarked):
        out = signal_noise_ratio(s, swm)
        tmp_list.append(out)
    return np.mean(tmp_list)


def resample_to16k(data, old_sr):
    new_fs = 16000
    new_data = data[::int(old_sr / new_fs)]
    return new_data

src/wavmark/utils/metric_util.py
wavmark/wavmark
import argparse
import ast


class MyParser():

    def __init__(self, epoch=0, batch_size=0, worker=0, seed=2526,
                 max_hour=100, early_stop=5, lr=1e-3,
                 model_save_folder=None):
        super(MyParser, self).__init__()
        parser = argparse.ArgumentParser()
        parser.add_argument("--seed", default=seed, type=int)
        parser.add_argument("--worker", default=worker, type=int)
        parser.add_argument("--epoch", default=epoch, type=int)
        parser.add_argument("--batch_size", default=batch_size, type=int)
        parser.add_argument("--max_hour", default=max_hour, type=int)
        parser.add_argument("--early_stop", default=early_stop, type=int)
        parser.add_argument("--lr", default=lr, type=float)
        parser.add_argument("--model_save_folder", default=model_save_folder, type=str)
        self.core_parser = parser

    def use_wb(self, project, name, dryrun=True):
        self.project = project
        self.name = name
        self.dryrun = dryrun
        parser = self.core_parser
        parser.add_argument("--project", default=self.project, type=str)
        parser.add_argument("--name", default=self.name, type=str)
        parser.add_argument("--dryrun", default=self.dryrun, type=ast.literal_eval)

    def custom(self, the_dict):
        parser = self.core_parser
        for key in the_dict:
            value = the_dict[key]
            if type(value) == str or value is None:
                parser.add_argument("--" + key, default=value, type=str)
            elif type(value) == int:
                parser.add_argument("--" + key, default=value, type=int)
            elif type(value) == float:
                parser.add_argument("--" + key, default=value, type=float)
            elif type(value) == bool:
                parser.add_argument("--" + key, default=value, type=ast.literal_eval)
            else:
                raise Exception("unsupported type:" + type(value))

    def parse(self):
        args = parse_it(self.core_parser)
        return args

    def show(self):
        the_dic = vars(self.parse())
        keys = list(the_dic.keys())
        keys.sort()
        for key in keys:
            print(key, ":", the_dic[key])

    def parse_in_jupyter(self):
        args = self.core_parser.parse_args([])
        return args


def parse_it(parser):
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    pass

src/wavmark/utils/my_parser.py
wavmark/wavmark
import pathlib
import os


def mk_parent_dir_if_necessary(img_save_path):
    folder = pathlib.Path(img_save_path).parent
    if not os.path.exists(folder):
        os.makedirs(folder)


def mk_dir_if_necessary(folder):
    if not os.path.exists(folder):
        os.makedirs(folder)

src/wavmark/utils/path_util.py
wavmark/wavmark
import torch
import numpy as np
from ..utils import metric_util
import tqdm
import time

# The pattern bits can be any random sequence.
# But don't use all-zeros, all-ones, or any periodic sequence, which will seriously hurt decoding performance.
fix_pattern = [1, 1, 1, 1, 0, 0, 1, 0, 0, 1, 1, 1, 0, 1, 1, 0, 0, 1, 1, 0, 0, 0,
               0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 1,
               1, 1, 0, 1, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1,
               1, 1, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0,
               0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 1, 0]


def add_watermark(bit_arr, data, num_point, shift_range, device, model, min_snr, max_snr, show_progress):
    t1 = time.time()

    chunk_size = num_point + int(num_point * shift_range)
    num_segments = int(len(data) / chunk_size)
    len_remain = len(data) - num_segments * chunk_size

    output_chunks = []
    encoded_sections = 0
    skip_sections = 0

    the_iter = range(num_segments)
    if show_progress:
        the_iter = tqdm.tqdm(the_iter, desc="Processing")

    for i in the_iter:
        start_point = i * chunk_size
        current_chunk = data[start_point:start_point + chunk_size].copy()
        # [watermark_segment | shift_area ]
        current_chunk_cover_area = current_chunk[0:num_point]
        current_chunk_shift_area = current_chunk[num_point:]
        current_chunk_cover_area_wmd, state = encode_trunck_with_snr_check(i, current_chunk_cover_area,
                                                                           bit_arr,
                                                                           device, model, min_snr, max_snr)

        if state == "skip":
            skip_sections += 1
        else:
            encoded_sections += 1

        output = np.concatenate([current_chunk_cover_area_wmd, current_chunk_shift_area])
        assert output.shape == current_chunk.shape
        output_chunks.append(output)

    assert len(output_chunks) > 0
    if len_remain > 0:
        output_chunks.append(data[len(data) - len_remain:])

    reconstructed_array = np.concatenate(output_chunks)

    time_cost = time.time() - t1

    info = {
        "time_cost": time_cost,
        "encoded_sections": encoded_sections,
        "skip_sections": skip_sections,
    }
    return reconstructed_array, info


def encode_trunck_with_snr_check(idx_trunck, signal, wm, device, model, min_snr, max_snr):
    signal_for_encode = signal
    encode_times = 0
    while True:
        encode_times += 1
        signal_wmd = encode_trunck(signal_for_encode, wm, device, model)
        snr = metric_util.signal_noise_ratio(signal, signal_wmd)
        if encode_times == 1 and snr < min_snr:
            print("skip section:%d, snr too low:%.1f" % (idx_trunck, min_snr))
            return signal, "skip"

        if snr < max_snr:
            return signal_wmd, encode_times
        # snr is too hugh
        signal_for_encode = signal_wmd

        if encode_times > 10:
            return signal_wmd, encode_times


def encode_trunck(trunck, wm, device, model):
    with torch.no_grad():
        signal = torch.FloatTensor(trunck).to(device)[None]
        message = torch.FloatTensor(np.array(wm)).to(device)[None]
        signal_wmd_tensor = model.encode(signal, message)
        signal_wmd = signal_wmd_tensor.detach().cpu().numpy().squeeze()
        return signal_wmd

src/wavmark/utils/wm_add_util.py
wavmark/wavmark
# import pdb

import torch
import numpy as np
import tqdm
import time


def decode_trunck(trunck, model, device):
    with torch.no_grad():
        signal = torch.FloatTensor(trunck).to(device).unsqueeze(0)
        message = (model.decode(signal) >= 0.5).int()
        message = message.detach().cpu().numpy().squeeze()
    return message


def extract_watermark_v3_batch(data, start_bit, shift_range, num_point, model, device, batch_size=10,
                               shift_range_p=0.5, show_progress=False):
    assert type(show_progress) == bool
    start_time = time.time()
    # 1.determine the shift step length:
    shift_step = int(shift_range * num_point * shift_range_p)

    # 2.determine where to perform detection
    # pdb.set_trace()
    total_detections = (len(data) - num_point) // shift_step
    total_detect_points = [i * shift_step for i in range(total_detections)]

    # 3.construct batch for detection
    total_batch_counts = len(total_detect_points) // batch_size + 1
    results = []

    the_iter = range(total_batch_counts)
    if show_progress:
        the_iter = tqdm.tqdm(range(total_batch_counts))

    for i in the_iter:
        detect_points = total_detect_points[i * batch_size:i * batch_size + batch_size]
        if len(detect_points) == 0:
            break
        current_batch = np.array([data[p:p + num_point] for p in detect_points])
        with torch.no_grad():
            signal = torch.FloatTensor(current_batch).to(device)
            batch_message = (model.decode(signal) >= 0.5).int().detach().cpu().numpy()
            for p, bit_array in zip(detect_points, batch_message):
                decoded_start_bit = bit_array[0:len(start_bit)]
                ber_start_bit = 1 - np.mean(start_bit == decoded_start_bit)
                num_equal_bits = np.sum(start_bit == decoded_start_bit)
                if ber_start_bit > 0:  # exact match
                    continue
                results.append({
                    "sim": 1 - ber_start_bit,
                    "num_equal_bits": num_equal_bits,
                    "msg": bit_array,
                    "start_position": p,
                    "start_time_position": p / 16000
                })

    end_time = time.time()
    time_cost = end_time - start_time

    info = {
        "time_cost": time_cost,
        "results": results,
    }

    if len(results) == 0:
        return None, info

    results_1 = [i["msg"] for i in results if np.isclose(i["sim"], 1.0)]
    mean_result = (np.array(results_1).mean(axis=0) >= 0.5).astype(int)
    return mean_result, info

src/wavmark/utils/wm_decode_util.py
wavmark/wavmark