/*
|
* Copyright (C) 2018 The Android Open Source Project
|
*
|
* 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 "utils/tflite/dist_diversification.h"
|
|
#include <algorithm>
|
#include "tensorflow/lite/context.h"
|
#include "tensorflow/lite/kernels/kernel_util.h"
|
#include "tensorflow/lite/model.h"
|
|
namespace libtextclassifier3 {
|
namespace {
|
|
// Returns a vector of row indices in a distance matrix.
|
// Indices are increasing and the distance of every selected index to others
|
// is larger than `min_distance`.
|
template <typename DistanceMatrixType>
|
std::vector<int> DiversifyByDistance(const DistanceMatrixType& distance_matrix,
|
const int matrix_size,
|
const float min_distance,
|
const int max_num_results) {
|
std::vector<int> result{0};
|
result.reserve(max_num_results);
|
int index = 1;
|
while (result.size() < max_num_results && index < matrix_size) {
|
for (; index < matrix_size; ++index) {
|
bool too_close = false;
|
for (const int selected_index : result) {
|
if (distance_matrix(index, selected_index) < min_distance) {
|
too_close = true;
|
break;
|
}
|
}
|
if (!too_close) {
|
result.push_back(index);
|
++index;
|
break;
|
}
|
}
|
}
|
return result;
|
}
|
|
// Input parameters for the op.
|
enum DistDiversificationInputs {
|
DIST_DIVERSIFICATION_INPUT_DISTANCE_MATRIX = 0,
|
DIST_DIVERSIFICATION_INPUT_MIN_DISTANCE = 1,
|
DIST_DIVERSIFICATION_INPUT_NUM_RESULTS = 2
|
};
|
|
// Output parameters for the op.
|
enum DistDiversificationOutputs {
|
DIST_DIVERSIFICATION_OUTPUT_INDICES = 0,
|
DIST_DIVERSIFICATION_OUTPUT_LENGTH = 1,
|
};
|
|
TfLiteIntArray* CreateSizeArray(const std::initializer_list<int>& sizes) {
|
TfLiteIntArray* array_size = TfLiteIntArrayCreate(sizes.size());
|
int index = 0;
|
for (const int size : sizes) {
|
array_size->data[index++] = size;
|
}
|
return array_size;
|
}
|
|
TfLiteStatus AllocateOutputIndexes(TfLiteContext* context, TfLiteNode* node) {
|
const TfLiteTensor& num_results =
|
context
|
->tensors[node->inputs->data[DIST_DIVERSIFICATION_INPUT_NUM_RESULTS]];
|
TfLiteTensor& output_indices =
|
context
|
->tensors[node->outputs->data[DIST_DIVERSIFICATION_OUTPUT_INDICES]];
|
return context->ResizeTensor(context, &output_indices,
|
CreateSizeArray({num_results.data.i32[0]}));
|
}
|
|
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
|
const TfLiteTensor& num_results =
|
context
|
->tensors[node->inputs->data[DIST_DIVERSIFICATION_INPUT_NUM_RESULTS]];
|
if (tflite::IsConstantTensor(&num_results)) {
|
TF_LITE_ENSURE_OK(context, AllocateOutputIndexes(context, node));
|
} else {
|
TfLiteTensor& output_indices =
|
context
|
->tensors[node->outputs->data[DIST_DIVERSIFICATION_OUTPUT_INDICES]];
|
tflite::SetTensorToDynamic(&output_indices);
|
}
|
TfLiteTensor& output_length =
|
context->tensors[node->outputs->data[DIST_DIVERSIFICATION_OUTPUT_LENGTH]];
|
TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, &output_length,
|
CreateSizeArray({1})));
|
return kTfLiteOk;
|
}
|
|
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
|
TfLiteTensor& output_indices =
|
context
|
->tensors[node->outputs->data[DIST_DIVERSIFICATION_OUTPUT_INDICES]];
|
if (tflite::IsDynamicTensor(&output_indices)) {
|
TF_LITE_ENSURE_OK(context, AllocateOutputIndexes(context, node));
|
}
|
const TfLiteTensor& distance_matrix =
|
context->tensors[node->inputs
|
->data[DIST_DIVERSIFICATION_INPUT_DISTANCE_MATRIX]];
|
const int distance_matrix_dim = distance_matrix.dims->data[0];
|
const float min_distance =
|
context
|
->tensors[node->inputs->data[DIST_DIVERSIFICATION_INPUT_MIN_DISTANCE]]
|
.data.f[0];
|
const int num_results =
|
context
|
->tensors[node->inputs->data[DIST_DIVERSIFICATION_INPUT_NUM_RESULTS]]
|
.data.i32[0];
|
const auto indices = DiversifyByDistance(
|
[&](int row, int col) {
|
return distance_matrix.data.f[row * distance_matrix_dim + col];
|
},
|
distance_matrix_dim, min_distance, num_results);
|
std::copy(indices.begin(), indices.end(), output_indices.data.i32);
|
std::fill_n(output_indices.data.i32 + indices.size(),
|
num_results - indices.size(), -1);
|
TfLiteTensor& output_length =
|
context->tensors[node->outputs->data[DIST_DIVERSIFICATION_OUTPUT_LENGTH]];
|
*output_length.data.i32 = indices.size();
|
return kTfLiteOk;
|
}
|
|
} // namespace
|
} // namespace libtextclassifier3
|
|
namespace tflite {
|
namespace ops {
|
namespace custom {
|
TfLiteRegistration* Register_DISTANCE_DIVERSIFICATION() {
|
static TfLiteRegistration r = {nullptr, nullptr, libtextclassifier3::Prepare,
|
libtextclassifier3::Eval};
|
return &r;
|
}
|
} // namespace custom
|
} // namespace ops
|
} // namespace tflite
|
