The model only work on messages taken from the pool.

I found that if the message was taken from the pool file (samples/message_pool.npy), the model performed very well (BER=0). But if the message doesn't come from the pool, it behaves like a random guess. Here's my test code:

import pdb
from evaluation import restore_model
import librosa
from main import embed
from main import detect, random_message
import numpy as np
import config
import tensorflow as tf
from config import MESSAGE_POOL

embedder_model = restore_model("embedder_model")
detector_model = restore_model("detector_model")

def expand_message(a):
    # (1,512) ===> (1, 16, 2, 512)
    repeats = [16, 1]
    a_tiled = tf.tile(a, repeats)
    a_expanded = tf.expand_dims(a_tiled, axis=1)
    a_final = tf.tile(a_expanded, [1, 2, 1])
    a_expanded2 = tf.expand_dims(a_final, axis=0)
    return a_expanded2

def fun1(watermark):
    # load signal
    config.HOP_LENGTH = 490
    signal, sr = librosa.load('samples/example1.wav', sr=16000)
    signal = np.expand_dims(signal, axis=0)

    # encode:
    watermark_expanded = expand_message(watermark)
    signal_wmd = embed(embedder_model, signal, message=watermark_expanded)

    # decode:
    detector_output = detect(detector_model, signal_wmd).numpy().squeeze()
    watermark_detected_bit = tf.where(tf.greater_equal(detector_output, 0.5), 1, 0).numpy().squeeze()

    # calculate BER:
    watermark_npy = watermark.numpy().squeeze()
    ber = (watermark_npy != watermark_detected_bit).sum() / len(watermark_detected_bit)
    return ber

if __name__ == "__main__":
    watermark_from_pool = random_message()
    print("pool-message :", fun1(watermark_from_pool))
    print("1 - pool-message :", fun1(1 - watermark_from_pool))

    # random message
    watermark_from_np = np.random.randint(2, size=512)
    watermark_from_np_tensor = tf.expand_dims(tf.constant(watermark_from_np), axis=0)
    print("random message:", fun1(watermark_from_np_tensor))

and the BER output is:

pool-message : 0.0
1 - pool-message : 0.47265625
random message: 0.47265625

Also, there is a typo for the restore_model function in the evaluation.py, with an 's' missing from the statement keras.model:

restored_model = keras.model.load_model(model_path)

===>

restored_model = keras.models.load_model(model_path)

kosta-pmf / dnn-audio-watermarking

The model only work on messages taken from the pool. #6