prepare_stim.py

r"""
--------------------------------------------------------------------------------
--  _____       ______  _____                                                 --
-- |_   _|     |  ____|/ ____|                                                --
--   | |  _ __ | |__  | (___                                                  --
--   | | | '_ \|  __|  \___ \   Institute of Embedded Systems                 --
--  _| |_| | | | |____ ____) |  Zurich University of Applied Sciences         --
-- |_____|_| |_|______|_____/   8401 Winterthur, Switzerland                  --
--------------------------------------------------------------------------------
--! @file
--! @brief Functions to create a stimuli file for usage inside the netwrok testbench
--!
--! @author vore
--! @date 18-11-2025
--
--------------------------------------------------------------------------------
"""

import json
import os

import numpy as np
import onnx
import scipy.io.wavfile as wav
from python_speech_features import mfcc
from scipy import signal
from scipy.signal.windows import hann

from helper import read_first_multithresh, threshold_indices_f
from new_config_data import ConfigReader


class StimuliPrep:
    def __init__(self, model_path, config_file, n_samples=0):
        self.tf_desired_samples = 16000
        self.tf_window_size_samples = 480
        self.tf_sample_rate = 16000
        self.tf_window_size_ms = 30.0
        self.tf_window_stride_ms = 20.0
        self.tf_dct_coefficient_count = 10
        self.model = onnx.load(model_path)
        self.N_SAMPLES = n_samples
        self.sample_classes = [
            "down",
            "go",
            "left",
            "no",
            "off",
            "on",
            "right",
            "stop",
            "up",
            "yes",
        ]
        self.config_file = config_file

    def py_speech_preprocessing(self, raw_signal, sample_rate):
        # Resample
        num_target_samples = round(self.tf_sample_rate / sample_rate * len(raw_signal))
        resampled_data = signal.resample(raw_signal, num_target_samples)
        # Rescale
        rescaled_data = resampled_data / np.max(resampled_data)
        # Pad
        padded_data = np.pad(
            rescaled_data, [[0, self.tf_desired_samples - rescaled_data.shape[-1]]], mode="constant"
        )
        # Calculate MFCC features
        nfft = int(2 ** np.ceil(np.log2(self.tf_window_size_samples)))
        mfcc_feat_py = mfcc(
            padded_data,
            self.tf_sample_rate,
            winlen=self.tf_window_size_ms / 1000.0,
            winstep=self.tf_window_stride_ms / 1000.0,
            numcep=self.tf_dct_coefficient_count,
            nfilt=40,
            nfft=nfft,
            lowfreq=20.0,
            highfreq=4000.0,
            winfunc=hann,
            appendEnergy=False,
            preemph=0.0,
            ceplifter=0.0,
        )
        # Cut and transpose MFCC features
        mfcc_feat_py = mfcc_feat_py[:-1, :].T

        return mfcc_feat_py

    def quantize_input(self, mfcc_feat_py):
        # Scaling
        quant_mfcc_feat = mfcc_feat_py / 0.8298503756523132
        # Clamping & rounding
        quant_mfcc_feat = np.where(quant_mfcc_feat > 127.0, 127.0, quant_mfcc_feat)
        quant_mfcc_feat = np.where(quant_mfcc_feat < -127.0, -127.0, quant_mfcc_feat)
        quant_mfcc_feat = np.round(quant_mfcc_feat)
        quant_mfcc_feat = quant_mfcc_feat.astype(np.int8).reshape((1, 490))

        return quant_mfcc_feat

    def read_wav_file(self, path):
        rate, raw_signal = wav.read(path)
        mfcc_feat_py = self.py_speech_preprocessing(raw_signal, rate)
        quant_mfcc_feat = self.quantize_input(mfcc_feat_py)
        return quant_mfcc_feat

    def prepare_audio_sample(self, sample):
        sample = np.asarray(sample).reshape((490,))
        path = os.path.dirname(os.path.realpath(__file__))
        data = {}
        pre_threshs = read_first_multithresh(self.model)
        threshs = np.tile(pre_threshs, 490)
        input_data = threshold_indices_f(sample, threshs)
        data["Input Data" + str(0)] = input_data.tolist()
        with open(path + "/simulation/data/" + "qonnx.json", "w") as f:
            json.dump(data, f, indent=4)

    def prepare_all_audio_samples(self):
        path = os.path.dirname(os.path.realpath(__file__)) + "/GSCV2_data/audio_samples/"
        path2 = os.path.dirname(os.path.realpath(__file__))
        data = {}
        pre_threshs = read_first_multithresh(self.model)
        threshs = np.tile(pre_threshs, 490)
        for i, sample_class in enumerate(self.sample_classes):
            rate, raw_signal = wav.read(path + f"audio_sample_{sample_class}.wav")
            mfcc_feat = self.py_speech_preprocessing(raw_signal, rate)
            quant_mfcc = self.quantize_input(mfcc_feat)
            sample = np.asarray(quant_mfcc).reshape((490,))
            input_data = threshold_indices_f(sample, threshs)
            data["Input Data" + str(i)] = input_data.tolist()
            # write_binary_file2(path + "/simulation/data/" + "qonnx" + str(i) + ".dat", input_data, bits=32)
        with open(path2 + "/simulation/data/" + "qonnx.json", "w") as f:
            json.dump(data, f, indent=4)

    def prepare_data(self, model):
        pass

    def prepare_wav(self, sample_path):
        data = self.read_wav_file(sample_path)
        self.prepare_audio_sample(data)

    def prepare_validation_data(self, data_path):
        input_npy = np.load(data_path)
        path = os.path.dirname(os.path.realpath(__file__))
        data = {}
        pre_threshs = read_first_multithresh(self.model)
        threshs = np.tile(pre_threshs, 490)
        for i in range(self.N_SAMPLES):
            input_data = input_npy[i]
            input_data = threshold_indices_f(input_data, threshs)
            data["Input Data" + str(i)] = input_data.tolist()
        with open(path + "/simulation/data/" + "qonnx.json", "w") as f:
            json.dump(data, f, indent=4)

    def create_data_file(self):
        cr = ConfigReader(self.config_file)
        layer_data = cr.get_layer_data()
        g_layer_sizes = layer_data["ls"]
        g_kernel_sizes = layer_data["ks"]

        pre_stim = 1
        after_stim = cr.data_width * sum(g_kernel_sizes) * sum(g_layer_sizes) + 2 * len(
            g_layer_sizes
        )

        n_tests = pre_stim + cr.data_width + g_layer_sizes[0] + 2

        input_data = []
        start = []
        wen = []
        waddr = []
        path = os.path.dirname(os.path.realpath(__file__))

        with open(path + "/simulation/data/" + "qonnx.json") as data_file:
            json_data = json.load(data_file)

        for i in range(self.N_SAMPLES):
            input_data.append(
                [0] + json_data["Input Data" + str(i)] + [0 for _ in range(cr.data_width + 2)]
            )

        start = (
            [0 for _ in range(g_layer_sizes[0] * cr.data_width + 1 + pre_stim)]
            + [1]
            + [0 for _ in range(1)]
        )
        wen = [[0 for _ in range(g_layer_sizes[0])]]
        # first permutation of the write_enable vector [1,0,0...,0]
        start_list = [1] + [0 for _ in range(g_layer_sizes[0] - 1)]
        # Cycle through the input size and rotate the start_list vector
        for i in range(g_layer_sizes[0]):
            wen.extend([start_list[-i:] + start_list[:-i] for _ in range(cr.data_width)])
        # wen.extend([[0 for _ in range(g_layer_sizes[0])] for _ in range(after_stim)])
        wen.extend(
            [[0 for _ in range(g_layer_sizes[0])] for _ in range(pre_stim + cr.data_width + 2)]
        )
        # Construct waddr
        # First case is all 0
        waddr = [0]
        # Cycle through the samples
        waddr.extend(list(range(cr.data_width)))
        # Zero after the first cycle
        waddr.extend([0 for _ in range(after_stim)])
        path = os.path.dirname(os.path.realpath(__file__))
        with open(path + "/simulation/scripts/network.dat", "w") as dat_file:
            for j in range(self.N_SAMPLES):
                for i in range(n_tests):
                    test_inputs = []
                    test_inputs.append(start[i])
                    test_inputs.append(input_data[j][i])
                    wen_i = wen[i]
                    wen_i = " ".join(map(str, wen_i))
                    test_inputs.extend(wen[i])
                    waddr_i = waddr[i]
                    test_inputs.append(waddr_i)
                    dat_file.write(" ".join(list(map(str, test_inputs))))
                    dat_file.write("\n")