model.py

import os
from typing import Union

from keras.src.optimizers import Adam
from numpy import set_printoptions as np_set_print_opts
from tensorflow import expand_dims as tf_expand_dims
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.math import exp as tf_exp

import feature_extraction as fe
import pre_processing as pp

""" This module contains functions to create the model and there utility functions. """

from tensorflow.keras.layers import Conv2D, Conv2DTranspose, MaxPooling2D, UpSampling2D, Flatten, Dense, BatchNormalization


def callbacks(lr_decay_rate=-0.1):
    # Callbacks
    def exponential_decay(epoch, lr):
        if epoch < 2:
            return lr
        else:
            return float(lr * tf_exp(lr_decay_rate))

    callback = keras.callbacks.LearningRateScheduler(exponential_decay)

    if not os.path.exists(os.path.join(os.getcwd(), 'model')):
        os.mkdir(os.path.join(os.getcwd(), 'model'))

    if not os.path.exists(os.path.join(os.getcwd(), 'model', 'tmp')):
        os.mkdir(os.path.join(os.getcwd(), 'model', 'tmp'))

    file_path = os.path.join(os.getcwd(), 'model', 'tmp', "{epoch:02d}-{val_accuracy:.3f}-{val_loss:.2f}.keras")
    checkpoint = keras.callbacks.ModelCheckpoint(filepath=file_path,
                                                 monitor="val_accuracy", verbose=1,
                                                 save_best_only=True,
                                                 save_weights_only=False,
                                                 mode='auto', save_freq="epoch")

    return callback, checkpoint


def cnn_autoencoder(input_shape, learning_rate=0.001):
    """
    Creates a CNN encoder-decoder model for audio signal reconstruction.

      :param learning_rate: learning rate provided to the model
      :param input_shape: A tuple representing the shape of the input STFT (128, 376, 1).

    Returns:
      A compiled Keras model.
    """

    # Encoder
    inputs = keras.Input(shape=input_shape)

    # Encoder Block
    def encoder_block(x, filters, kernel_size=(3, 3), activation=layers.LeakyReLU(negative_slope=0.1), padding="same"):
        shortcut = x
        x = Conv2D(filters=filters, kernel_size=kernel_size, activation=activation, padding=padding)(x)
        x = BatchNormalization()(x)
        # x = Conv2D(filters=filters, kernel_size=kernel_size, activation=None, padding=padding)(x)
        # x = layers.add([x, shortcut])  # Removed for this version
        # x = layers.Activation(activation)(x)
        x = layers.Dropout(0.2)(x)
        return x

    x = encoder_block(inputs, filters=16)
    x = MaxPooling2D(pool_size=(2, 2), padding="same")(x)

    x = encoder_block(x, filters=32)
    x = MaxPooling2D(pool_size=(2, 2), padding="same")(x)

    x = encoder_block(x, filters=64)
    x = MaxPooling2D(pool_size=(2, 2), padding="same")(x)

    encoded = Flatten()(x)

    # Decoder
    # x = Dense(units=640, activation=layers.LeakyReLU(negative_slope=0.1))(encoded)
    # x = Dense(units=640, activation=layers.LeakyReLU(negative_slope=0.1))(x)
    # x = layers.Reshape((10, 12, 16))(x)

    # Decoder Block
    def decoder_block(x, filters, kernel_size=(3, 3), activation=layers.LeakyReLU(negative_slope=0.1), padding="same"):
        shortcut = Conv2DTranspose(filters=filters, kernel_size=kernel_size, strides=(1, 1), activation=None,
                                   padding=padding)(x)
        x = Conv2DTranspose(filters=filters, kernel_size=kernel_size, activation=activation, padding=padding)(x)
        x = BatchNormalization()(x)
        # x = Conv2DTranspose(filters=filters, kernel_size=kernel_size, activation=None, padding=padding)(x)
        # Removed for this version: layers.add([x, shortcut])
        # x = layers.Activation(activation)(x)
        x = layers.Dropout(0.2)(x)
        return x

    x = decoder_block(x, filters=64)
    x = UpSampling2D(size=(2, 2))(x)

    x = decoder_block(x, filters=32)
    x = UpSampling2D(size=(2, 2))(x)

    x = decoder_block(x, filters=16)

    flatten = Flatten()(x)
    dense = Dense(units=64, activation="relu")(flatten)
    dense = Dense(units=32, activation="relu")(dense)

    outputs = Dense(units=2, activation="sigmoid")(dense)

    # Model
    model = keras.Model(inputs=inputs, outputs=outputs)

    optimizer = Adam(learning_rate=learning_rate)
    model.compile(optimizer=optimizer, loss='binary_crossentropy', metrics=["accuracy"])

    model.summary()

    return model, callbacks()


# A complex cnn model.
def complex_cnn_model(input_shape, learning_rate=0.001):
    model = keras.Sequential([
        layers.InputLayer(shape=input_shape),

        # First conv block
        layers.Conv2D(filters=32, kernel_size=5, strides=1, padding="Same", name="1b1Conv",
                      kernel_initializer='he_normal',
                      activation=layers.LeakyReLU(negative_slope=0.1), ),
        layers.BatchNormalization(),
        layers.Conv2D(filters=32, kernel_size=3, strides=1, padding="Same", name="1b2Conv",
                      kernel_initializer='he_normal',
                      activation=layers.LeakyReLU(negative_slope=0.1), ),
        layers.GaussianDropout(0.5),
        layers.BatchNormalization(),

        # Second conv block
        layers.Conv2D(filters=64, kernel_size=3, strides=2, padding="same", name="2b1Conv",
                      kernel_initializer='he_normal',
                      activation=layers.LeakyReLU(negative_slope=0.1), ),
        layers.BatchNormalization(),
        layers.Conv2D(filters=64, kernel_size=3, strides=1, padding="same", name="2b2Conv",
                      kernel_initializer='he_normal',
                      activation=layers.LeakyReLU(negative_slope=0.1), ),
        layers.GaussianDropout(0.5),
        layers.BatchNormalization(),

        # Third conv block
        layers.Conv2D(filters=128, kernel_size=3, strides=2, padding="same", name="3b1Conv",
                      kernel_initializer='he_normal',
                      activation=layers.LeakyReLU(negative_slope=0.1), ),
        layers.BatchNormalization(),
        layers.Conv2D(filters=128, kernel_size=3, strides=1, padding="same", name="3b2Conv",
                      kernel_initializer='he_normal',
                      activation=layers.LeakyReLU(negative_slope=0.1), ),
        layers.GaussianDropout(0.5),
        layers.BatchNormalization(),

        # Fourth conv block
        layers.Conv2D(filters=256, kernel_size=3, strides=2, padding="same", name="4b1Conv",
                      kernel_initializer='he_normal',
                      activation=layers.LeakyReLU(negative_slope=0.1), ),
        layers.BatchNormalization(),
        layers.Conv2D(filters=256, kernel_size=3, strides=1, padding="same", name="4b2Conv",
                      kernel_initializer='he_normal',
                      activation=layers.LeakyReLU(negative_slope=0.1), ),
        layers.BatchNormalization(),
        layers.GaussianDropout(0.5),

        # 5th
        layers.Conv2D(filters=256, kernel_size=3, strides=2, padding="same", name="5b1Conv",
                      kernel_initializer='he_normal',
                      activation=layers.LeakyReLU(negative_slope=0.1), ),
        layers.BatchNormalization(),
        layers.Conv2D(filters=256, kernel_size=3, strides=1, padding="same", name="5b2Conv",
                      kernel_initializer='he_normal',
                      activation=layers.LeakyReLU(negative_slope=0.1), ),
        layers.BatchNormalization(),
        layers.GaussianDropout(0.5),

        # Flattening and then classifying
        layers.Flatten(),
        layers.Dense(units=64, name="1FC", activation=layers.LeakyReLU(negative_slope=0.1),
                     kernel_initializer='he_uniform', ),
        layers.Dense(units=64, name="2FC", activation=layers.LeakyReLU(negative_slope=0.1),
                     kernel_initializer='he_uniform', ),
        layers.GaussianDropout(0.6),
        layers.Dense(units=64, name="3FC", activation=layers.LeakyReLU(negative_slope=0.1),
                     kernel_initializer='he_uniform', ),
        layers.Dense(units=32, name="4FC", activation=layers.LeakyReLU(negative_slope=0.1),
                     kernel_initializer='he_uniform', ),

        layers.Dense(units=2, name="Output", activation='sigmoid', kernel_initializer='glorot_uniform', ),

    ])
    # Compiling model

    optimizer = Adam(learning_rate=learning_rate)
    model.compile(optimizer=optimizer, loss='binary_crossentropy', metrics=["accuracy"])

    model.summary()

    return model, callbacks()


# convert audio to wav or flac format
def verify(audio, model: keras.Model) -> Union[tuple[str, bool], str]:
    # extracting feature from audio
    feature_extractor = fe.ExtractSTFTDB(n_fft=512, hop_length=128)
    normalizer = pp.TFMinMaxNormalize(new_min=-1, new_max=1)

    feature = feature_extractor.extract(audio)
    print("\nFeatures:\n", feature, feature.shape)
    norm_feature = normalizer.normalize(feature)
    print("\nNormalized Features:\n", norm_feature, norm_feature.shape)

    reshaped_feature = tf_expand_dims(norm_feature, axis=0)
    reshaped_feature = tf_expand_dims(reshaped_feature, axis=-1)
    print("\nShape for predicting: ", reshaped_feature.shape)

    np_set_print_opts(precision=2)

    # predicting
    prediction = model.predict(reshaped_feature)

    print(prediction)
    if prediction[0, 0] > 0.5:
        return (f'This audio is {prediction[0, 0] * 100:.2f}% is likely to be fake.\n'
                f'This audio is {prediction[0, 1] * 100:.2f}% is likely to be real.'), False
    elif prediction[0, 1] > 0.5:
        return (f'This audio is {prediction[0, 1] * 100:.2f}% is likely to be real.\n'
                f'This audio is {prediction[0, 0] * 100:.2f}% is likely to be fake.'), True
    else:
        return 'Sorry! The model is not working properly.'


# deprecated
def load_model(model_path):
    try:
        model = keras.models.browse_model(model_path)
        if model:
            print(f'loaded model: {model_path}')
        return model
    except Exception as e:
        print('Tried loading the model: Failed')
        print('trying to load the model with custom function LeakyReLU(negative_slope=0.1)')

    try:
        model = keras.models.load_model(model_path,
                                        custom_objects={'LeakyReLU': keras.layers.LeakyReLU(negative_slope=0.1)})
        if model:
            print(f'loaded model: {model_path}')
        return model
    except Exception as e:
        raise f"Tried loading the model with custom function LeakyReLU(negative_slope=0.1): Failed with: {e}"


# loading video
import io
from moviepy.editor import VideoFileClip
from soundfile import write
from librosa import load


def convert_to_wav_in_memory(video_path):
    """
    Extracts audio from video, converts it to WAV in-memory, and returns it.

    Args:
      video_path: Path to the video file.

    Returns:
      A byte string containing the WAV-encoded audio data or None on error.
    """
    # Extract audio using moviepy (same as previous example)
    clip = VideoFileClip(video_path)
    audio_data = clip.audio.to_ndarray()
    clip.close()

    # Convert audio data to WAV format in-memory
    wav_data = io.BytesIO()
    write(wav_data, audio_data, samplerate=clip.audio.fps, format="wav")

    # Return the WAV-encoded audio data
    return wav_data.getvalue()


def load_video(video_path):
    # Example usage
    video_path = video_path
    wav_data = convert_to_wav_in_memory(video_path)

    if wav_data:
        # Now you can use wav_data as the source for librosa.load
        y, sr = load(io.BytesIO(wav_data), sr=None)
        # Process the loaded audio data using librosa
    else:
        print("Error converting audio to WAV format.")