wehub-resource-sync
219 lines
7.5 KiB
C++
219 lines
7.5 KiB
C++
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
==============================================================================*/
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <vector>
|
|
|
|
#include "flatbuffers/flexbuffers.h" // from @flatbuffers
|
|
#include "tensorflow/lite/c/c_api_types.h"
|
|
#include "tensorflow/lite/c/common.h"
|
|
#include "tensorflow/lite/experimental/microfrontend/lib/frontend.h"
|
|
#include "tensorflow/lite/experimental/microfrontend/lib/frontend_util.h"
|
|
#include "tensorflow/lite/kernels/internal/compatibility.h"
|
|
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
|
|
#include "tensorflow/lite/kernels/kernel_util.h"
|
|
|
|
namespace tflite {
|
|
namespace ops {
|
|
namespace custom {
|
|
namespace audio_microfrontend {
|
|
|
|
constexpr int kInputTensor = 0;
|
|
constexpr int kOutputTensor = 0;
|
|
|
|
typedef struct {
|
|
int sample_rate;
|
|
FrontendState* state;
|
|
int left_context;
|
|
int right_context;
|
|
int frame_stride;
|
|
bool zero_padding;
|
|
int out_scale;
|
|
bool out_float;
|
|
} TfLiteAudioMicrofrontendParams;
|
|
|
|
void* Init(TfLiteContext* context, const char* buffer, size_t length) {
|
|
auto* data = new TfLiteAudioMicrofrontendParams;
|
|
|
|
const uint8_t* buffer_t = reinterpret_cast<const uint8_t*>(buffer);
|
|
const flexbuffers::Map& m = flexbuffers::GetRoot(buffer_t, length).AsMap();
|
|
|
|
data->sample_rate = m["sample_rate"].AsInt32();
|
|
|
|
struct FrontendConfig config;
|
|
config.window.size_ms = m["window_size"].AsInt32();
|
|
config.window.step_size_ms = m["window_step"].AsInt32();
|
|
config.filterbank.num_channels = m["num_channels"].AsInt32();
|
|
config.filterbank.upper_band_limit = m["upper_band_limit"].AsFloat();
|
|
config.filterbank.lower_band_limit = m["lower_band_limit"].AsFloat();
|
|
config.noise_reduction.smoothing_bits = m["smoothing_bits"].AsInt32();
|
|
config.noise_reduction.even_smoothing = m["even_smoothing"].AsFloat();
|
|
config.noise_reduction.odd_smoothing = m["odd_smoothing"].AsFloat();
|
|
config.noise_reduction.min_signal_remaining =
|
|
m["min_signal_remaining"].AsFloat();
|
|
config.pcan_gain_control.enable_pcan = m["enable_pcan"].AsBool();
|
|
config.pcan_gain_control.strength = m["pcan_strength"].AsFloat();
|
|
config.pcan_gain_control.offset = m["pcan_offset"].AsFloat();
|
|
config.pcan_gain_control.gain_bits = m["gain_bits"].AsInt32();
|
|
config.log_scale.enable_log = m["enable_log"].AsBool();
|
|
config.log_scale.scale_shift = m["scale_shift"].AsInt32();
|
|
|
|
data->state = new FrontendState;
|
|
FrontendPopulateState(&config, data->state, data->sample_rate);
|
|
|
|
data->left_context = m["left_context"].AsInt32();
|
|
data->right_context = m["right_context"].AsInt32();
|
|
data->frame_stride = m["frame_stride"].AsInt32();
|
|
data->zero_padding = m["zero_padding"].AsBool();
|
|
data->out_scale = m["out_scale"].AsInt32();
|
|
data->out_float = m["out_float"].AsBool();
|
|
|
|
return data;
|
|
}
|
|
|
|
void Free(TfLiteContext* context, void* buffer) {
|
|
auto* data = reinterpret_cast<TfLiteAudioMicrofrontendParams*>(buffer);
|
|
FrontendFreeStateContents(data->state);
|
|
delete data->state;
|
|
delete data;
|
|
}
|
|
|
|
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
|
|
auto* data =
|
|
reinterpret_cast<TfLiteAudioMicrofrontendParams*>(node->user_data);
|
|
|
|
TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
|
|
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
|
|
|
|
const TfLiteTensor* input;
|
|
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
|
|
TfLiteTensor* output;
|
|
TF_LITE_ENSURE_OK(context,
|
|
GetOutputSafe(context, node, kOutputTensor, &output));
|
|
|
|
TF_LITE_ENSURE_EQ(context, NumDimensions(input), 1);
|
|
|
|
TF_LITE_ENSURE_EQ(context, input->type, kTfLiteInt16);
|
|
output->type = kTfLiteInt32;
|
|
if (data->out_float) {
|
|
output->type = kTfLiteFloat32;
|
|
}
|
|
|
|
TfLiteIntArray* output_size = TfLiteIntArrayCreate(2);
|
|
int num_frames = 0;
|
|
if (input->dims->data[0] >= data->state->window.size) {
|
|
num_frames = (input->dims->data[0] - data->state->window.size) /
|
|
data->state->window.step / data->frame_stride +
|
|
1;
|
|
}
|
|
output_size->data[0] = num_frames;
|
|
output_size->data[1] = data->state->filterbank.num_channels *
|
|
(1 + data->left_context + data->right_context);
|
|
|
|
return context->ResizeTensor(context, output, output_size);
|
|
}
|
|
|
|
template <typename T>
|
|
void GenerateFeatures(TfLiteAudioMicrofrontendParams* data,
|
|
const TfLiteTensor* input, TfLiteTensor* output) {
|
|
const int16_t* audio_data = GetTensorData<int16_t>(input);
|
|
int64_t audio_size = input->dims->data[0];
|
|
|
|
T* filterbanks_flat = GetTensorData<T>(output);
|
|
|
|
int num_frames = 0;
|
|
if (audio_size >= data->state->window.size) {
|
|
num_frames = (input->dims->data[0] - data->state->window.size) /
|
|
data->state->window.step +
|
|
1;
|
|
}
|
|
std::vector<std::vector<T>> frame_buffer(num_frames);
|
|
|
|
int frame_index = 0;
|
|
while (audio_size > 0) {
|
|
size_t num_samples_read;
|
|
struct FrontendOutput output = FrontendProcessSamples(
|
|
data->state, audio_data, audio_size, &num_samples_read);
|
|
audio_data += num_samples_read;
|
|
audio_size -= num_samples_read;
|
|
|
|
if (output.values != nullptr) {
|
|
frame_buffer[frame_index].reserve(output.size);
|
|
int i;
|
|
for (i = 0; i < output.size; ++i) {
|
|
frame_buffer[frame_index].push_back(static_cast<T>(output.values[i]) /
|
|
data->out_scale);
|
|
}
|
|
++frame_index;
|
|
}
|
|
}
|
|
|
|
int index = 0;
|
|
std::vector<T> pad(data->state->filterbank.num_channels, 0);
|
|
int anchor;
|
|
for (anchor = 0; anchor < frame_buffer.size(); anchor += data->frame_stride) {
|
|
int frame;
|
|
for (frame = anchor - data->left_context;
|
|
frame <= anchor + data->right_context; ++frame) {
|
|
std::vector<T>* feature;
|
|
if (data->zero_padding && (frame < 0 || frame >= frame_buffer.size())) {
|
|
feature = &pad;
|
|
} else if (frame < 0) {
|
|
feature = &frame_buffer[0];
|
|
} else if (frame >= frame_buffer.size()) {
|
|
feature = &frame_buffer[frame_buffer.size() - 1];
|
|
} else {
|
|
feature = &frame_buffer[frame];
|
|
}
|
|
for (auto f : *feature) {
|
|
filterbanks_flat[index++] = f;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
|
|
auto* data =
|
|
reinterpret_cast<TfLiteAudioMicrofrontendParams*>(node->user_data);
|
|
FrontendReset(data->state);
|
|
|
|
const TfLiteTensor* input;
|
|
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
|
|
TfLiteTensor* output;
|
|
TF_LITE_ENSURE_OK(context,
|
|
GetOutputSafe(context, node, kOutputTensor, &output));
|
|
|
|
if (data->out_float) {
|
|
GenerateFeatures<float>(data, input, output);
|
|
} else {
|
|
GenerateFeatures<int32>(data, input, output);
|
|
}
|
|
|
|
return kTfLiteOk;
|
|
}
|
|
|
|
} // namespace audio_microfrontend
|
|
|
|
TfLiteRegistration* Register_AUDIO_MICROFRONTEND() {
|
|
static TfLiteRegistration r = {
|
|
audio_microfrontend::Init, audio_microfrontend::Free,
|
|
audio_microfrontend::Prepare, audio_microfrontend::Eval};
|
|
return &r;
|
|
}
|
|
|
|
} // namespace custom
|
|
} // namespace ops
|
|
} // namespace tflite
|