wehub-resource-sync
552 lines
24 KiB
Plaintext
552 lines
24 KiB
Plaintext
/* ******************************************************************************
|
|
*
|
|
*
|
|
* 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 <array/NDArrayFactory.h>
|
|
#include <exceptions/cuda_exception.h>
|
|
#include <legacy/NativeOps.h>
|
|
#include <ops/declarable/helpers/image_suppression.h>
|
|
|
|
#include <queue>
|
|
|
|
#include "execution/cuda/LaunchDims.h"
|
|
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
|
// needToSuppressWithThreshold - predicate for suppression
|
|
// boxes - boxes tensor buffer
|
|
// boxesShape boxes tensor shape
|
|
// previousIndex - index for current pos value
|
|
// nextIndex - index for neighbor pos value
|
|
// threshold - threashold value to suppress
|
|
//
|
|
// return value: true, if threshold is overcome, false otherwise
|
|
//
|
|
template <typename T>
|
|
static SD_DEVICE bool needToSuppressWithThreshold(T* boxes, LongType const* boxesShape, int previousIndex,
|
|
int nextIndex, T threshold) {
|
|
LongType previous0[] = {previousIndex, 0};
|
|
LongType previous1[] = {previousIndex, 1};
|
|
LongType previous2[] = {previousIndex, 2};
|
|
LongType previous3[] = {previousIndex, 3};
|
|
LongType next0[] = {nextIndex, 0};
|
|
LongType next1[] = {nextIndex, 1};
|
|
LongType next2[] = {nextIndex, 2};
|
|
LongType next3[] = {nextIndex, 3};
|
|
|
|
LongType prevOffset0, prevOffset1, prevOffset2, prevOffset3;
|
|
LongType nextOffset0, nextOffset1, nextOffset2, nextOffset3;
|
|
|
|
COORDS2INDEX(2, shape::stride(boxesShape), previous0, prevOffset0);
|
|
COORDS2INDEX(2, shape::stride(boxesShape), previous1, prevOffset1);
|
|
COORDS2INDEX(2, shape::stride(boxesShape), previous2, prevOffset2);
|
|
COORDS2INDEX(2, shape::stride(boxesShape), previous3, prevOffset3);
|
|
COORDS2INDEX(2, shape::stride(boxesShape), next0, nextOffset0);
|
|
COORDS2INDEX(2, shape::stride(boxesShape), next1, nextOffset1);
|
|
COORDS2INDEX(2, shape::stride(boxesShape), next2, nextOffset2);
|
|
COORDS2INDEX(2, shape::stride(boxesShape), next3, nextOffset3);
|
|
|
|
// we have rectangle with given max values. Compute vexes of rectangle first
|
|
|
|
T minYPrev = math::sd_min(boxes[prevOffset0], boxes[prevOffset2]);
|
|
T minXPrev = math::sd_min(boxes[prevOffset1], boxes[prevOffset3]);
|
|
T maxYPrev = math::sd_max(boxes[prevOffset0], boxes[prevOffset2]);
|
|
T maxXPrev = math::sd_max(boxes[prevOffset1], boxes[prevOffset3]);
|
|
T minYNext = math::sd_min(boxes[nextOffset0], boxes[nextOffset2]);
|
|
T minXNext = math::sd_min(boxes[nextOffset1], boxes[nextOffset3]);
|
|
T maxYNext = math::sd_max(boxes[nextOffset0], boxes[nextOffset2]);
|
|
T maxXNext = math::sd_max(boxes[nextOffset1], boxes[nextOffset3]);
|
|
|
|
// compute areas for comparation
|
|
T areaPrev = (maxYPrev - minYPrev) * (maxXPrev - minXPrev);
|
|
T areaNext = (maxYNext - minYNext) * (maxXNext - minXNext);
|
|
|
|
// of course, areas should be positive
|
|
if (areaNext <= T(0.f) || areaPrev <= T(0.f)) return false;
|
|
|
|
// compute intersection of rectangles
|
|
T minIntersectionY = math::sd_max(minYPrev, minYNext);
|
|
T minIntersectionX = math::sd_max(minXPrev, minXNext);
|
|
T maxIntersectionY = math::sd_min(maxYPrev, maxYNext);
|
|
T maxIntersectionX = math::sd_min(maxXPrev, maxXNext);
|
|
T intersectionArea = math::sd_max(T(maxIntersectionY - minIntersectionY), T(0.0f)) *
|
|
math::sd_max(T(maxIntersectionX - minIntersectionX), T(0.0f));
|
|
T intersectionValue = intersectionArea / (areaPrev + areaNext - intersectionArea);
|
|
// final check
|
|
return intersectionValue > threshold;
|
|
}
|
|
|
|
template <typename T>
|
|
static inline T similirityV3_(NDArray& boxes, LongType i, LongType j) {
|
|
const T zero = static_cast<T>(0.f);
|
|
const T yminI = math::sd_min(boxes.t<T>(i, 0), boxes.t<T>(i, 2));
|
|
const T xminI = math::sd_min(boxes.t<T>(i, 1), boxes.t<T>(i, 3));
|
|
const T ymaxI = math::sd_max(boxes.t<T>(i, 0), boxes.t<T>(i, 2));
|
|
const T xmaxI = math::sd_max(boxes.t<T>(i, 1), boxes.t<T>(i, 3));
|
|
const T yminJ = math::sd_min(boxes.t<T>(j, 0), boxes.t<T>(j, 2));
|
|
const T xminJ = math::sd_min(boxes.t<T>(j, 1), boxes.t<T>(j, 3));
|
|
const T ymaxJ = math::sd_max(boxes.t<T>(j, 0), boxes.t<T>(j, 2));
|
|
const T xmaxJ = math::sd_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::sd_max(yminI, yminJ);
|
|
const T intersectionXmin = math::sd_max(xminI, xminJ);
|
|
const T intersectionYmax = math::sd_min(ymaxI, ymaxJ);
|
|
const T intersectionXmax = math::sd_min(xmaxI, xmaxJ);
|
|
const T intersectionY = intersectionYmax - intersectionYmin;
|
|
const T intersectionX = intersectionXmax - intersectionXmin;
|
|
const T intersectionArea = math::sd_max(intersectionY, zero) * math::sd_max(intersectionX, zero);
|
|
return intersectionArea / (areaI + areaJ - intersectionArea);
|
|
}
|
|
|
|
|
|
template <typename T>
|
|
static SD_DEVICE T similirityV3(T* boxes, LongType const* boxesShape, int previousIndex, int nextIndex) {
|
|
LongType previous0[] = {previousIndex, 0};
|
|
LongType previous1[] = {previousIndex, 1};
|
|
LongType previous2[] = {previousIndex, 2};
|
|
LongType previous3[] = {previousIndex, 3};
|
|
LongType next0[] = {nextIndex, 0};
|
|
LongType next1[] = {nextIndex, 1};
|
|
LongType next2[] = {nextIndex, 2};
|
|
LongType next3[] = {nextIndex, 3};
|
|
|
|
LongType prevOffset0, prevOffset1, prevOffset2, prevOffset3;
|
|
LongType nextOffset0, nextOffset1, nextOffset2, nextOffset3;
|
|
|
|
COORDS2INDEX(2, shape::stride(boxesShape), previous0, prevOffset0);
|
|
COORDS2INDEX(2, shape::stride(boxesShape), previous1, prevOffset1);
|
|
COORDS2INDEX(2, shape::stride(boxesShape), previous2, prevOffset2);
|
|
COORDS2INDEX(2, shape::stride(boxesShape), previous3, prevOffset3);
|
|
COORDS2INDEX(2, shape::stride(boxesShape), next0, nextOffset0);
|
|
COORDS2INDEX(2, shape::stride(boxesShape), next1, nextOffset1);
|
|
COORDS2INDEX(2, shape::stride(boxesShape), next2, nextOffset2);
|
|
COORDS2INDEX(2, shape::stride(boxesShape), next3, nextOffset3);
|
|
|
|
// we have rectangle with given max values. Compute vexes of rectangle first
|
|
|
|
T minYPrev = math::sd_min(boxes[prevOffset0], boxes[prevOffset2]);
|
|
T minXPrev = math::sd_min(boxes[prevOffset1], boxes[prevOffset3]);
|
|
T maxYPrev = math::sd_max(boxes[prevOffset0], boxes[prevOffset2]);
|
|
T maxXPrev = math::sd_max(boxes[prevOffset1], boxes[prevOffset3]);
|
|
T minYNext = math::sd_min(boxes[nextOffset0], boxes[nextOffset2]);
|
|
T minXNext = math::sd_min(boxes[nextOffset1], boxes[nextOffset3]);
|
|
T maxYNext = math::sd_max(boxes[nextOffset0], boxes[nextOffset2]);
|
|
T maxXNext = math::sd_max(boxes[nextOffset1], boxes[nextOffset3]);
|
|
|
|
// compute areas for comparator
|
|
T areaPrev = (maxYPrev - minYPrev) * (maxXPrev - minXPrev);
|
|
T areaNext = (maxYNext - minYNext) * (maxXNext - minXNext);
|
|
|
|
// of course, areas should be positive
|
|
if (areaNext <= T(0.f) || areaPrev <= T(0.f)) return false;
|
|
|
|
// compute intersection of rectangles
|
|
T minIntersectionY = math::sd_max(minYPrev, minYNext);
|
|
T minIntersectionX = math::sd_max(minXPrev, minXNext);
|
|
T maxIntersectionY = math::sd_min(maxYPrev, maxYNext);
|
|
T maxIntersectionX = math::sd_min(maxXPrev, maxXNext);
|
|
T intersectionArea = math::sd_max(T(maxIntersectionY - minIntersectionY), T(0.0f)) *
|
|
math::sd_max(T(maxIntersectionX - minIntersectionX), T(0.0f));
|
|
T intersectionValue = intersectionArea / (areaPrev + areaNext - intersectionArea);
|
|
// final check
|
|
return intersectionValue;
|
|
}
|
|
|
|
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
|
// shouldSelectKernel - compute status for all selected rectangles (boxes)
|
|
//
|
|
// we compute boolean flag as shared uint32 and return it on final only for the first thread
|
|
//
|
|
template <typename T, typename I>
|
|
static SD_KERNEL void shouldSelectKernel(T* boxesBuf, LongType const* boxesShape, I* indexBuf,
|
|
I* selectedIndicesData, double threshold, int numSelected, int i,
|
|
bool* shouldSelect) {
|
|
auto tid = blockIdx.x * blockDim.x + threadIdx.x;
|
|
auto step = gridDim.x * blockDim.x;
|
|
__shared__ unsigned int shouldSelectShared;
|
|
if (threadIdx.x == 0) {
|
|
shouldSelectShared = (unsigned int)shouldSelect[0];
|
|
}
|
|
__syncthreads();
|
|
for (int j = numSelected - 1 - tid; j >= 0; j -= step) {
|
|
if (shouldSelectShared) {
|
|
if (needToSuppressWithThreshold(boxesBuf, boxesShape, indexBuf[i], indexBuf[selectedIndicesData[j]],
|
|
T(threshold)))
|
|
atomicCAS(&shouldSelectShared, 1, 0); // exchange only when need to suppress
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
// final move: collect result
|
|
if (threadIdx.x == 0) {
|
|
*shouldSelect = shouldSelectShared > 0;
|
|
}
|
|
}
|
|
|
|
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
|
// indices - type depended, indicesLong - type defined (only 64bit integers)
|
|
//
|
|
template <typename I>
|
|
static SD_KERNEL void copyIndices(void* indices, void* indicesLong, LongType len) {
|
|
I* indexBuf = reinterpret_cast<I*>(indices);
|
|
LongType* srcBuf = reinterpret_cast<LongType*>(indicesLong);
|
|
;
|
|
|
|
auto tid = threadIdx.x + blockIdx.x * blockDim.x;
|
|
auto step = blockDim.x * gridDim.x;
|
|
|
|
for (auto i = tid; i < len; i += step) indexBuf[i] = (I)srcBuf[i];
|
|
}
|
|
|
|
template <typename T, typename I>
|
|
static SD_KERNEL void suppressScores(T* scores, I* indices, LongType length, T scoreThreshold) {
|
|
auto start = blockIdx.x * blockDim.x;
|
|
auto step = gridDim.x * blockDim.x;
|
|
|
|
for (auto e = start + threadIdx.x; e < (int)length; e += step) {
|
|
if (scores[e] < scoreThreshold) {
|
|
scores[e] = scoreThreshold;
|
|
indices[e] = -1;
|
|
} else {
|
|
indices[e] = I(e);
|
|
}
|
|
}
|
|
}
|
|
|
|
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
|
// nonMaxSuppressionV2 algorithm - given from TF NonMaxSuppressionV2 implementation
|
|
//
|
|
template <typename T, typename I>
|
|
static void nonMaxSuppressionV2_(LaunchContext* context, NDArray* boxes, NDArray* scales, int maxSize,
|
|
double threshold, double scoreThreshold, NDArray* output) {
|
|
auto stream = context->getCudaStream();
|
|
NDArray::prepareSpecialUse({output}, {boxes, scales});
|
|
std::vector<sd::LongType> shape = {scales->lengthOf()};
|
|
NDArray indices (NDArrayFactory::create_<I>(
|
|
'c', shape, context)); // - 1, scales->lengthOf()); //, scales->getContext());
|
|
|
|
NDArray scores(*scales);
|
|
Pointer extras[2] = {nullptr, stream};
|
|
auto indexBuf = indices.dataBuffer()->specialAsT<I>();
|
|
auto scoreBuf = scores.dataBuffer()->specialAsT<T>();
|
|
dim3 launchDims = getLaunchDims("image_suppress_scores");
|
|
suppressScores<T, I><<<launchDims.x, launchDims.y,launchDims.z, *stream>>>(scoreBuf, indexBuf, scores.lengthOf(), T(scoreThreshold));
|
|
indices.tickWriteDevice();
|
|
sortByValue(extras, &indices,
|
|
&scores,true);
|
|
indices.tickWriteDevice();
|
|
NDArray selectedIndices = NDArrayFactory::create<I>('c', {output->lengthOf()}, context);
|
|
int numSelected = 0;
|
|
int numBoxes = boxes->sizeAt(0), tt(0);
|
|
auto boxesBuf = reinterpret_cast<T*>(boxes->specialBuffer());
|
|
|
|
auto selectedIndicesData = reinterpret_cast<I*>(selectedIndices.specialBuffer());
|
|
auto outputBuf = reinterpret_cast<I*>(output->specialBuffer());
|
|
|
|
bool* shouldSelectD;
|
|
auto err = cudaMalloc(&shouldSelectD, sizeof(bool));
|
|
if (err) {
|
|
throw cuda_exception::build("helpers::nonMaxSuppressionV2: Cannot allocate memory for bool flag", err);
|
|
}
|
|
for (I i = 0; i < boxes->sizeAt(0); ++i) {
|
|
bool shouldSelect = numSelected < output->lengthOf();
|
|
if (shouldSelect) {
|
|
err = cudaMemcpy(shouldSelectD, &shouldSelect, sizeof(bool), cudaMemcpyHostToDevice);
|
|
if (err) {
|
|
throw cuda_exception::build("helpers::nonMaxSuppressionV2: Cannot set up bool flag to device", err);
|
|
}
|
|
|
|
dim3 selectDims = getLaunchDims("image_suppress_select");
|
|
shouldSelectKernel<T, I><<<selectDims.y,selectDims.x,selectDims.z, *stream>>>(
|
|
boxesBuf, boxes->specialShapeInfo(), indexBuf, selectedIndicesData, threshold, numSelected, i, shouldSelectD);
|
|
err = cudaMemcpy(&shouldSelect, shouldSelectD, sizeof(bool), cudaMemcpyDeviceToHost);
|
|
if (err) {
|
|
throw cuda_exception::build("helpers::nonMaxSuppressionV2: Cannot set up bool flag to host", err);
|
|
}
|
|
}
|
|
|
|
if (shouldSelect) {
|
|
cudaMemcpy(reinterpret_cast<I*>(output->specialBuffer()) + numSelected, indexBuf + i, sizeof(I),
|
|
cudaMemcpyDeviceToDevice);
|
|
cudaMemcpy(selectedIndicesData + numSelected, &i, sizeof(I), cudaMemcpyHostToDevice);
|
|
numSelected++;
|
|
}
|
|
}
|
|
|
|
err = cudaFree(shouldSelectD);
|
|
if (err) {
|
|
throw cuda_exception::build("helpers::nonMaxSuppressionV2: Cannot deallocate memory for bool flag", err);
|
|
}
|
|
}
|
|
|
|
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
|
template <typename T, typename I>
|
|
static SD_DEVICE bool checkOverlapBoxes(T* boxes, LongType const* shape, T* scores, I* indices, I* selectedIndices,
|
|
I* startIndices, I selectedSize, I nextCandidateIndex, T overlapThreshold,
|
|
T scoreThreshold, bool simple) {
|
|
bool shouldHardSuppress = false;
|
|
T& nextCandidateScore = scores[nextCandidateIndex];
|
|
I selectedIndex = indices[nextCandidateIndex];
|
|
I finish = startIndices[nextCandidateIndex];
|
|
|
|
for (int j = selectedSize; j > finish; --j) {
|
|
T boxVal;
|
|
if (simple) {
|
|
LongType xPos[] = {selectedIndex, selectedIndices[j - 1]};
|
|
LongType xShift;
|
|
COORDS2INDEX(shape::rank(shape), shape::stride(shape), xPos, xShift);
|
|
boxVal = boxes[xShift];
|
|
} else {
|
|
boxVal = similirityV3(boxes, shape, selectedIndex, selectedIndices[j - 1]);
|
|
}
|
|
if (boxVal > static_cast<T>(overlapThreshold)) nextCandidateScore = static_cast<T>(0.f);
|
|
|
|
// First decide whether to perform hard suppression
|
|
if (boxVal >= overlapThreshold) {
|
|
shouldHardSuppress = true;
|
|
break;
|
|
}
|
|
|
|
// If nextCandidate survives hard suppression, apply soft suppression
|
|
if (nextCandidateScore <= static_cast<T>(scoreThreshold)) break;
|
|
}
|
|
|
|
return shouldHardSuppress;
|
|
}
|
|
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
|
template <typename T, typename I>
|
|
static SD_KERNEL void suppressNonMaxOverlapKernel(T* boxes, LongType const* boxesShape, T* scoresData, I* indices,
|
|
I* startIndices, LongType length, I maxOutputLen,
|
|
T overlapThreshold, T scoreThreshold, I* output, LongType const* outputShape, I* outputLength, bool simple) {
|
|
__shared__ I selectedSize;
|
|
__shared__ I* tempOutput;
|
|
|
|
if (threadIdx.x == 0) {
|
|
selectedSize = outputLength ? *outputLength : maxOutputLen;
|
|
extern __shared__ unsigned char shmem[];
|
|
tempOutput = (I*)shmem;
|
|
}
|
|
__syncthreads();
|
|
|
|
auto start = blockIdx.x * blockDim.x;
|
|
auto step = blockDim.x * gridDim.x;
|
|
|
|
for (I nextCandidateIndex = start + threadIdx.x; selectedSize < maxOutputLen && nextCandidateIndex < (I)length;) {
|
|
auto originalScore = scoresData[nextCandidateIndex];
|
|
I nextCandidateBoxIndex = indices[nextCandidateIndex];
|
|
auto selectedSizeMark = selectedSize;
|
|
|
|
// skip for cases when index is less than 0 (under score threshold)
|
|
if (nextCandidateBoxIndex < 0) {
|
|
nextCandidateIndex += step;
|
|
continue;
|
|
}
|
|
// check for overlaps
|
|
bool shouldHardSuppress =
|
|
checkOverlapBoxes(boxes, boxesShape, scoresData, indices, tempOutput, startIndices, selectedSize,
|
|
nextCandidateIndex, overlapThreshold, scoreThreshold, simple); // false;
|
|
T nextCandidateScore = scoresData[nextCandidateIndex];
|
|
|
|
startIndices[nextCandidateIndex] = selectedSize;
|
|
if (!shouldHardSuppress) {
|
|
if (nextCandidateScore == originalScore) {
|
|
// Suppression has not occurred, so select nextCandidate
|
|
I currSize = math::atomics::sd_atomicAdd(&selectedSize, (I)1);
|
|
if (output) {
|
|
output[currSize] = nextCandidateBoxIndex;
|
|
}
|
|
tempOutput[currSize] = nextCandidateBoxIndex;
|
|
|
|
}
|
|
|
|
if ((float) nextCandidateScore > (float) scoreThreshold) {
|
|
// Soft suppression has occurred and current score is still greater than
|
|
// scoreThreshold; add nextCandidate back onto priority queue.
|
|
continue; // in some cases, this index not 0
|
|
}
|
|
}
|
|
nextCandidateIndex += step;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
|
|
if (threadIdx.x == 0) {
|
|
printf("selectedSize: %i\n", selectedSize);
|
|
if (outputLength) *outputLength = selectedSize;
|
|
}
|
|
}
|
|
|
|
|
|
typedef NDArray (*SimilarityFunc)(NDArray& boxes, LongType i, LongType j);
|
|
template <typename T>
|
|
static inline T similarityOverlaps_(NDArray& boxes, LongType i, LongType j) {
|
|
return boxes.t<T>(i, j);
|
|
}
|
|
|
|
static NDArray similiratyOverlaps(NDArray& boxes, LongType i, LongType j) {
|
|
NDArray res(boxes.dataType(), boxes.getContext()); // = NDArrayFactory::create(0.);
|
|
BUILD_SINGLE_SELECTOR(boxes.dataType(), res = similarityOverlaps_, (boxes, i, j), SD_FLOAT_TYPES);
|
|
return res;
|
|
}
|
|
|
|
static NDArray similarityV3(NDArray& boxes, LongType i, LongType j) {
|
|
NDArray res(boxes.dataType(), boxes.getContext()); // = NDArrayFactory::create(0.);
|
|
BUILD_SINGLE_SELECTOR(boxes.dataType(), res = similirityV3_, (boxes, i, j), SD_FLOAT_TYPES);
|
|
return res;
|
|
}
|
|
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
|
template <typename T, typename I>
|
|
static LongType nonMaxSuppressionGeneric_(LaunchContext* context, NDArray* boxes, NDArray* scores,
|
|
int outputSize, float overlapThreshold, float scoreThreshold,
|
|
NDArray* output, SimilarityFunc f) {
|
|
auto stream = context->getCudaStream();
|
|
if (output)
|
|
NDArray::preparePrimaryUse({output}, {boxes, scores});
|
|
else {
|
|
if (!boxes->isActualOnHostSide()) boxes->syncToHost();
|
|
if (!scores->isActualOnHostSide()) scores->syncToHost();
|
|
}
|
|
|
|
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);
|
|
}
|
|
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 (LongType)selected.size();
|
|
}
|
|
|
|
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
|
void nonMaxSuppression(LaunchContext* context, NDArray* boxes, NDArray* scales, int maxSize, double threshold,
|
|
double scoreThreshold, NDArray* output) {
|
|
BUILD_DOUBLE_SELECTOR(boxes->dataType(), output->dataType(), nonMaxSuppressionV2_,
|
|
(context, boxes, scales, maxSize, threshold, scoreThreshold, output), SD_FLOAT_TYPES,
|
|
SD_INDEXING_TYPES);
|
|
}
|
|
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
LongType nonMaxSuppressionGeneric(LaunchContext* context, NDArray* boxes, NDArray* scales, int maxSize,
|
|
double threshold, double scoreThreshold, NDArray* output) {
|
|
BUILD_DOUBLE_SELECTOR(boxes->dataType(), output ? output->dataType() : DataType::INT32,
|
|
return nonMaxSuppressionGeneric_,
|
|
(context, boxes, scales, maxSize, threshold, scoreThreshold, output, similiratyOverlaps),
|
|
SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
|
return boxes->sizeAt(0);
|
|
}
|
|
|
|
LongType nonMaxSuppressionV3(LaunchContext* context, NDArray* boxes, NDArray* scores, int maxSize,
|
|
double overlapThreshold, double scoreThreshold, NDArray* output) {
|
|
BUILD_DOUBLE_SELECTOR(boxes->dataType(), output ? output->dataType() : DataType::INT32,
|
|
return nonMaxSuppressionGeneric_,
|
|
(context, boxes, scores, maxSize, overlapThreshold, scoreThreshold, output, similarityV3),
|
|
SD_FLOAT_TYPES, SD_INDEXING_TYPES);
|
|
return boxes->sizeAt(0);
|
|
}
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|