/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* 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.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author sgazeos@gmail.com
//
#include <ops/declarable/helpers/image_suppression.h>
#include <array/NDArrayFactory.h>
#include <algorithm>
#include <numeric>
#include <queue>
namespace sd {
namespace ops {
namespace helpers {
template <typename T>
static void nonMaxSuppressionV2_(NDArray* boxes, NDArray* scales, int maxSize, double overlapThreshold,
double scoreThreshold, NDArray* output) {
std::vector<int> indices(scales->lengthOf());
std::iota(indices.begin(), indices.end(), 0);
auto actualIndicesCount = indices.size();
for (auto e = 0; e < scales->lengthOf(); e++) {
if (scales->e<float>(e) < (float)scoreThreshold) {
indices[e] = -1;
actualIndicesCount--;
}
}
std::sort(indices.begin(), indices.end(), [scales](int i, int j) {return i >= 0 && j >=0?scales->e<T>(i) > scales->e<T>(j):(i > j);});
// std::vector<int> selected(output->lengthOf());
std::vector<int> selectedIndices(output->lengthOf(), 0);
auto needToSuppressWithThreshold = [] (NDArray& boxes, int previousIndex, int nextIndex, T threshold) -> bool {
if (previousIndex < 0 || nextIndex < 0) return true;
T minYPrev = sd::math::nd4j_min(boxes.t<T>(previousIndex, 0), boxes.t<T>(previousIndex, 2));
T minXPrev = sd::math::nd4j_min(boxes.t<T>(previousIndex, 1), boxes.t<T>(previousIndex, 3));
T maxYPrev = sd::math::nd4j_max(boxes.t<T>(previousIndex, 0), boxes.t<T>(previousIndex, 2));
T maxXPrev = sd::math::nd4j_max(boxes.t<T>(previousIndex, 1), boxes.t<T>(previousIndex, 3));
T minYNext = sd::math::nd4j_min(boxes.t<T>(nextIndex, 0), boxes.t<T>(nextIndex, 2));
T minXNext = sd::math::nd4j_min(boxes.t<T>(nextIndex, 1), boxes.t<T>(nextIndex, 3));
T maxYNext = sd::math::nd4j_max(boxes.t<T>(nextIndex, 0), boxes.t<T>(nextIndex, 2));
T maxXNext = sd::math::nd4j_max(boxes.t<T>(nextIndex, 1), boxes.t<T>(nextIndex, 3));
T areaPrev = (maxYPrev - minYPrev) * (maxXPrev - minXPrev);
T areaNext = (maxYNext - minYNext) * (maxXNext - minXNext);
if (areaNext <= T(0.f) || areaPrev <= T(0.f)) return false;
T minIntersectionY = sd::math::nd4j_max(minYPrev, minYNext);
T minIntersectionX = sd::math::nd4j_max(minXPrev, minXNext);
T maxIntersectionY = sd::math::nd4j_min(maxYPrev, maxYNext);
T maxIntersectionX = sd::math::nd4j_min(maxXPrev, maxXNext);
T intersectionArea =
sd::math::nd4j_max(T(maxIntersectionY - minIntersectionY), T(0.0f)) *
sd::math::nd4j_max(T(maxIntersectionX - minIntersectionX), T(0.0f));
T intersectionValue = intersectionArea / (areaPrev + areaNext - intersectionArea);
return intersectionValue > threshold;
};
// int numSelected = 0;
int numBoxes = actualIndicesCount; //boxes->sizeAt(0);
int numSelected = 0;
for (int i = 0; i < numBoxes; ++i) {
bool shouldSelect = numSelected < output->lengthOf();
// FIXME: add parallelism here
for (int j = numSelected - 1; j >= 0; --j) {
if (shouldSelect)
if (needToSuppressWithThreshold(*boxes, indices[i], indices[selectedIndices[j]], T(overlapThreshold))) {
shouldSelect = false;
}
}
if (shouldSelect) {
output->p(numSelected, indices[i]);
selectedIndices[numSelected++] = i;
}
}
}
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Return intersection-over-union overlap between boxes i and j
template <typename T>
static inline T similirityV3_(NDArray const& boxes, Nd4jLong i, Nd4jLong j) {
const T zero = static_cast<T>(0.f);
const T yminI = math::nd4j_min(boxes.t<T>(i, 0), boxes.t<T>(i, 2));
const T xminI = math::nd4j_min(boxes.t<T>(i, 1), boxes.t<T>(i, 3));
const T ymaxI = math::nd4j_max(boxes.t<T>(i, 0), boxes.t<T>(i, 2));
const T xmaxI = math::nd4j_max(boxes.t<T>(i, 1), boxes.t<T>(i, 3));
const T yminJ = math::nd4j_min(boxes.t<T>(j, 0), boxes.t<T>(j, 2));
const T xminJ = math::nd4j_min(boxes.t<T>(j, 1), boxes.t<T>(j, 3));
const T ymaxJ = math::nd4j_max(boxes.t<T>(j, 0), boxes.t<T>(j, 2));
const T xmaxJ = math::nd4j_max(boxes.t<T>(j, 1), boxes.t<T>(j, 3));
const T areaI = (ymaxI - yminI) * (xmaxI - xminI);
const T areaJ = (ymaxJ - yminJ) * (xmaxJ - xminJ);
if (areaI <= zero || areaJ <= zero) {
return zero;
}
const T intersectionYmin = math::nd4j_max(yminI, yminJ);
const T intersectionXmin = math::nd4j_max(xminI, xminJ);
const T intersectionYmax = math::nd4j_min(ymaxI, ymaxJ);
const T intersectionXmax = math::nd4j_min(xmaxI, xmaxJ);
const T intersectionY = intersectionYmax - intersectionYmin;
const T intersectionX = intersectionXmax - intersectionXmin;
const T intersectionArea = math::nd4j_max(intersectionY, zero) * math::nd4j_max(intersectionX, zero);
return intersectionArea / (areaI + areaJ - intersectionArea);
}
template <typename T>
static inline T similiratyOverlaps_(NDArray const& boxes, Nd4jLong i, Nd4jLong j) {
return boxes.t<T>(i, j);
}
typedef NDArray (*SimiliratyFunc)(NDArray const& boxes, Nd4jLong i, Nd4jLong j);
static NDArray similiratyOverlaps(NDArray const& boxes, Nd4jLong i, Nd4jLong j) {
NDArray res(boxes.dataType(), boxes.getContext()); // = NDArrayFactory::create(0.);
BUILD_SINGLE_SELECTOR(boxes.dataType(), res = similiratyOverlaps_, (boxes, i, j) , FLOAT_TYPES);
return res;
}
static NDArray similiratyV3(NDArray const& boxes, Nd4jLong i, Nd4jLong j) {
NDArray res(boxes.dataType(), boxes.getContext()); // = NDArrayFactory::create(0.);
BUILD_SINGLE_SELECTOR(boxes.dataType(), res = similirityV3_, (boxes, i, j) , FLOAT_TYPES);
return res;
}
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T, typename I>
static Nd4jLong
nonMaxSuppressionGeneric_(sd::LaunchContext* context, NDArray* boxes, NDArray* scores, int outputSize,
float overlapThreshold, float scoreThreshold, NDArray* output, SimiliratyFunc f) {
auto numBoxes = boxes->sizeAt(0);
T* scoresData = scores->dataBuffer()->primaryAsT<T>();
// Data structure for a selection candidate in NMS.
struct Candidate {
int _boxIndex;
T _score;
int _suppressBeginIndex;
};
auto cmp = [](const Candidate& bsI, const Candidate& bsJ) -> bool{
return ((bsI._score == bsJ._score) && (bsI._boxIndex > bsJ._boxIndex)) ||
(bsI._score < bsJ._score);
};
std::priority_queue<Candidate, std::deque<Candidate>, decltype(cmp)> candidatePriorityQueue(cmp);
for (auto i = 0; i < scores->lengthOf(); ++i) {
if ((float)scoresData[i] > (float)scoreThreshold) {
candidatePriorityQueue.emplace(Candidate({i, scoresData[i], 0}));
}
}
std::vector<I> selected;
T similarity, originalScore;
Candidate nextCandidate;
while (selected.size() < outputSize && !candidatePriorityQueue.empty()) {
nextCandidate = candidatePriorityQueue.top();
originalScore = nextCandidate._score;
candidatePriorityQueue.pop();
// Overlapping boxes are likely to have similar scores, therefore we
// iterate through the previously selected boxes backwards in order to
// see if `nextCandidate` should be suppressed. We also enforce a property
// that a candidate can be suppressed by another candidate no more than
// once via `suppress_begin_index` which tracks which previously selected
// boxes have already been compared against next_candidate prior to a given
// iteration. These previous selected boxes are then skipped over in the
// following loop.
bool shouldHardSuppress = false;
for (int j = static_cast<int>(selected.size()) - 1; j >= nextCandidate._suppressBeginIndex; --j) {
auto similarityA = f(*boxes, nextCandidate._boxIndex, selected[j]); //boxes->t<T>(nextCandidate._boxIndex, selected[j]);
similarity = similarityA.template t<T>(0);
nextCandidate._score *= T(similarity <= overlapThreshold?1.0:0.); //suppressWeightFunc(similarity);
// First decide whether to perform hard suppression
if ((float)similarity >= static_cast<float>(overlapThreshold)) {
shouldHardSuppress = true;
break;
}
// If next_candidate survives hard suppression, apply soft suppression
if ((float)nextCandidate._score <= (float)scoreThreshold) break;
}
// If `nextCandidate._score` has not dropped below `scoreThreshold`
// by this point, then we know that we went through all of the previous
// selections and can safely update `suppress_begin_index` to
// `selected.size()`. If on the other hand `next_candidate.score`
// *has* dropped below the score threshold, then since `suppressWeight`
// always returns values in [0, 1], further suppression by items that were
// not covered in the above for loop would not have caused the algorithm
// to select this item. We thus do the same update to
// `suppressBeginIndex`, but really, this element will not be added back
// into the priority queue in the following.
nextCandidate._suppressBeginIndex = selected.size();
if (!shouldHardSuppress) {
if (nextCandidate._score == originalScore) {
// Suppression has not occurred, so select next_candidate
selected.push_back(nextCandidate._boxIndex);
// selected_scores.push_back(nextCandidate._score);
}
if ((float)nextCandidate._score > (float)scoreThreshold) {
// Soft suppression has occurred and current score is still greater than
// score_threshold; add next_candidate back onto priority queue.
candidatePriorityQueue.push(nextCandidate);
}
}
}
if (output) {
DataBuffer buf(selected.data(), selected.size() * sizeof(I), DataTypeUtils::fromT<I>());
output->dataBuffer()->copyBufferFrom(buf, buf.getLenInBytes());
}
return (Nd4jLong)selected.size();
}
Nd4jLong
nonMaxSuppressionGeneric(sd::LaunchContext* context, NDArray* boxes, NDArray* scores, int maxSize,
double overlapThreshold, double scoreThreshold, NDArray* output) {
BUILD_DOUBLE_SELECTOR(boxes->dataType(), output == nullptr?DataType::INT32:output->dataType(), return nonMaxSuppressionGeneric_, (context, boxes, scores, maxSize, overlapThreshold, scoreThreshold, output, similiratyOverlaps), FLOAT_TYPES, INTEGER_TYPES);
return 0;
}
Nd4jLong
nonMaxSuppressionV3(sd::LaunchContext* context, NDArray* boxes, NDArray* scores, int maxSize,
double overlapThreshold, double scoreThreshold, NDArray* output) {
BUILD_DOUBLE_SELECTOR(boxes->dataType(), output == nullptr?DataType::INT32:output->dataType(), return nonMaxSuppressionGeneric_, (context, boxes, scores, maxSize, overlapThreshold, scoreThreshold, output, similiratyV3), FLOAT_TYPES, INTEGER_TYPES);
return 0;
}
BUILD_DOUBLE_TEMPLATE(template Nd4jLong nonMaxSuppressionGeneric_, (sd::LaunchContext* context, NDArray* boxes, NDArray* scores, int maxSize,
float overlapThreshold, float scoreThreshold, NDArray* output, SimiliratyFunc similiratyFunc), FLOAT_TYPES, INTEGER_TYPES);
void
nonMaxSuppression(sd::LaunchContext * context, NDArray* boxes, NDArray* scales, int maxSize,
double overlapThreshold, double scoreThreshold, NDArray* output) {
BUILD_SINGLE_SELECTOR(boxes->dataType(), nonMaxSuppressionV2_, (boxes, scales, maxSize,
overlapThreshold, scoreThreshold, output), NUMERIC_TYPES);
}
BUILD_SINGLE_TEMPLATE(template void nonMaxSuppressionV2_, (NDArray* boxes, NDArray* scales, int maxSize,
double overlapThreshold, double scoreThreshold, NDArray* output), NUMERIC_TYPES);
}
}
}