mirror of https://github.com/alibaba/MNN.git
				
				
				
			
		
			
				
	
	
		
			270 lines
		
	
	
		
			11 KiB
		
	
	
	
		
			C++
		
	
	
	
			
		
		
	
	
			270 lines
		
	
	
		
			11 KiB
		
	
	
	
		
			C++
		
	
	
	
//
 | 
						|
//  CPUDetectionOutput.cpp
 | 
						|
//  MNN
 | 
						|
//
 | 
						|
//  Created by MNN on 2018/07/17.
 | 
						|
//  Copyright © 2018, Alibaba Group Holding Limited
 | 
						|
//
 | 
						|
/* When use MSVC compile the file on x86 Release, a compiler internal error will be report because of MSVC's bug.
 | 
						|
   reference link: https://developercommunity.visualstudio.com/comments/535612/view.html */
 | 
						|
#if defined(_MSC_VER) && defined(_M_IX86) && !defined(_DEBUG)
 | 
						|
#pragma optimize("", off)
 | 
						|
#endif
 | 
						|
 | 
						|
#include "backend/cpu/CPUDetectionOutput.hpp"
 | 
						|
#include <math.h>
 | 
						|
#include <vector>
 | 
						|
//#define MNN_OPEN_TIME_TRACE
 | 
						|
#include <MNN/AutoTime.hpp>
 | 
						|
#include "backend/cpu/CPUBackend.hpp"
 | 
						|
#include "backend/cpu/compute/CommonOptFunction.h"
 | 
						|
#include "core/TensorUtils.hpp"
 | 
						|
 | 
						|
namespace MNN {
 | 
						|
 | 
						|
CPUDetectionOutput::CPUDetectionOutput(Backend *backend, int classCount, float nmsThreshold, int keepTopK,
 | 
						|
                                       float confidenceThreshold, float objectnessScore)
 | 
						|
    : Execution(backend),
 | 
						|
      mClassCount(classCount),
 | 
						|
      mNMSThreshold(nmsThreshold),
 | 
						|
      mKeepTopK(keepTopK),
 | 
						|
      mConfidenceThreshold(confidenceThreshold),
 | 
						|
      mObjectnessScoreThreshold(objectnessScore) {
 | 
						|
    TensorUtils::getDescribe(&mLocation)->dimensionFormat      = MNN_DATA_FORMAT_NCHW;
 | 
						|
    TensorUtils::getDescribe(&mConfidence)->dimensionFormat    = MNN_DATA_FORMAT_NCHW;
 | 
						|
    TensorUtils::getDescribe(&mPriorbox)->dimensionFormat      = MNN_DATA_FORMAT_NCHW;
 | 
						|
    TensorUtils::getDescribe(&mArmLocation)->dimensionFormat   = MNN_DATA_FORMAT_NCHW;
 | 
						|
    TensorUtils::getDescribe(&mArmConfidence)->dimensionFormat = MNN_DATA_FORMAT_NCHW;
 | 
						|
}
 | 
						|
 | 
						|
using score_box_t = std::tuple<float, float, float, float, int, float>;
 | 
						|
#define box_rect(xmin, ymin, xmax, ymax, label, score) std::make_tuple((xmin), (ymin), (xmax), (ymax), (label), (score))
 | 
						|
#define box_rect_xmin(rect) (std::get<0>(rect))
 | 
						|
#define box_rect_ymin(rect) (std::get<1>(rect))
 | 
						|
#define box_rect_xmax(rect) (std::get<2>(rect))
 | 
						|
#define box_rect_ymax(rect) (std::get<3>(rect))
 | 
						|
#define box_label(rect) (std::get<4>(rect))
 | 
						|
#define box_score(rect) (std::get<5>(rect))
 | 
						|
 | 
						|
static inline float intersectionArea(const score_box_t& a, const score_box_t& b) {
 | 
						|
    float axmin = box_rect_xmin(a), bxmin = box_rect_xmin(b);
 | 
						|
    float axmax = box_rect_xmax(a), bxmax = box_rect_xmax(b);
 | 
						|
    float aymin = box_rect_ymin(a), bymin = box_rect_ymin(b);
 | 
						|
    float aymax = box_rect_ymax(a), bymax = box_rect_ymax(b);
 | 
						|
    if (axmin > bxmax || axmax < bxmin || aymin > bymax || aymax < bymin)
 | 
						|
        return 0.f;
 | 
						|
 | 
						|
    float interWidth  = fmin(axmax, bxmax) - fmax(axmin, bxmin);
 | 
						|
    float interHeight = fmin(aymax, bymax) - fmax(aymin, bymin);
 | 
						|
    return interWidth * interHeight;
 | 
						|
}
 | 
						|
 | 
						|
static void pickBoxes(const std::vector<score_box_t> &boxes, std::vector<int> &picked, float nmsThreshold, int topK) {
 | 
						|
    long n = boxes.size();
 | 
						|
    std::vector<float> areas;
 | 
						|
    areas.resize(n);
 | 
						|
    for (int i = 0; i < n; i++) {
 | 
						|
        auto& box     = boxes[i];
 | 
						|
        float width  = box_rect_xmax(box) - box_rect_xmin(box);
 | 
						|
        float height = box_rect_ymax(box) - box_rect_ymin(box);
 | 
						|
        areas[i]     = width * height;
 | 
						|
    }
 | 
						|
 | 
						|
    for (int i = 0; i < n; i++) {
 | 
						|
        auto& a = boxes[i];
 | 
						|
 | 
						|
        bool keep = true;
 | 
						|
        for (auto pick : picked) {
 | 
						|
            auto& b = boxes[pick];
 | 
						|
 | 
						|
            // intersection over union
 | 
						|
            float interArea = intersectionArea(a, b);
 | 
						|
            float unionArea = areas[i] + areas[pick] - interArea;
 | 
						|
            if (interArea / unionArea > nmsThreshold) {
 | 
						|
                keep = false;
 | 
						|
                break;
 | 
						|
            }
 | 
						|
        }
 | 
						|
        if (keep) {
 | 
						|
            picked.push_back(i);
 | 
						|
            if (picked.size() >= topK) {
 | 
						|
                break;
 | 
						|
            }
 | 
						|
        }
 | 
						|
    }
 | 
						|
}
 | 
						|
 | 
						|
ErrorCode CPUDetectionOutput::onResize(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
 | 
						|
    auto &location = inputs[0];
 | 
						|
    auto &priorbox = inputs[2];
 | 
						|
    if (location->channel() != priorbox->height()) {
 | 
						|
        MNN_ERROR("Error for CPUDetection output, location and pribox not match\n");
 | 
						|
        return NOT_SUPPORT;
 | 
						|
    }
 | 
						|
    // location transform space
 | 
						|
    TensorUtils::copyShape(inputs[0], &mLocation, false);
 | 
						|
    backend()->onAcquireBuffer(&mLocation, Backend::DYNAMIC);
 | 
						|
 | 
						|
    // confidence transform space
 | 
						|
    TensorUtils::copyShape(inputs[1], &mConfidence, false);
 | 
						|
    backend()->onAcquireBuffer(&mConfidence, Backend::DYNAMIC);
 | 
						|
 | 
						|
    // priorbox transform space
 | 
						|
    TensorUtils::copyShape(inputs[2], &mPriorbox, false);
 | 
						|
    backend()->onAcquireBuffer(&mPriorbox, Backend::DYNAMIC);
 | 
						|
 | 
						|
    // refine
 | 
						|
    if (inputs.size() >= 5) {
 | 
						|
        TensorUtils::copyShape(inputs[3], &mArmConfidence, false);
 | 
						|
        TensorUtils::copyShape(inputs[4], &mArmLocation, false);
 | 
						|
 | 
						|
        backend()->onAcquireBuffer(&mArmConfidence, Backend::DYNAMIC);
 | 
						|
        backend()->onAcquireBuffer(&mArmLocation, Backend::DYNAMIC);
 | 
						|
        backend()->onReleaseBuffer(&mArmConfidence, Backend::DYNAMIC);
 | 
						|
        backend()->onReleaseBuffer(&mArmLocation, Backend::DYNAMIC);
 | 
						|
    }
 | 
						|
 | 
						|
    // release temp buffer space
 | 
						|
    backend()->onReleaseBuffer(&mLocation, Backend::DYNAMIC);
 | 
						|
    backend()->onReleaseBuffer(&mConfidence, Backend::DYNAMIC);
 | 
						|
    backend()->onReleaseBuffer(&mPriorbox, Backend::DYNAMIC);
 | 
						|
    return NO_ERROR;
 | 
						|
}
 | 
						|
 | 
						|
ErrorCode CPUDetectionOutput::onExecute(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
 | 
						|
    auto &location   = inputs[0];
 | 
						|
    auto &confidence = inputs[1];
 | 
						|
    auto &priorbox   = inputs[2];
 | 
						|
    auto &output     = outputs[0];
 | 
						|
 | 
						|
    // download
 | 
						|
    MNNUnpackC4Origin(mLocation.host<float>(), location->host<float>(), location->width() * location->height(),
 | 
						|
                location->channel(), location->width() * location->height());
 | 
						|
    MNNUnpackC4Origin(mConfidence.host<float>(), confidence->host<float>(), confidence->width() * confidence->height(),
 | 
						|
                confidence->channel(), confidence->width() * confidence->height());
 | 
						|
    MNNUnpackC4Origin(mPriorbox.host<float>(), priorbox->host<float>(), priorbox->width() * priorbox->height(),
 | 
						|
                priorbox->channel(), priorbox->width() * priorbox->height());
 | 
						|
 | 
						|
    bool refineDet = inputs.size() >= 5;
 | 
						|
    if (refineDet) {
 | 
						|
        Tensor *armconfidence = inputs[3];
 | 
						|
        Tensor *armlocation   = inputs[4];
 | 
						|
        MNNUnpackC4Origin(mArmConfidence.host<float>(), armconfidence->host<float>(),
 | 
						|
                    armconfidence->width() * armconfidence->height(), armconfidence->channel(), armconfidence->width() * armconfidence->height());
 | 
						|
        MNNUnpackC4Origin(mArmLocation.host<float>(), armlocation->host<float>(),
 | 
						|
                    armlocation->width() * armlocation->height(), armlocation->channel(), armlocation->width() * armlocation->height());
 | 
						|
    }
 | 
						|
 | 
						|
    auto priorCount       = priorbox->height() / 4;
 | 
						|
    auto locationPtr      = mLocation.host<const float>();
 | 
						|
    auto confidencePtr    = mConfidence.host<const float>();
 | 
						|
    auto priorboxPtr      = mPriorbox.host<const float>();
 | 
						|
    auto variancePtr      = mPriorbox.host<const float>() + priorbox->height() * 1;
 | 
						|
    auto armlocationPtr   = refineDet ? mArmLocation.host<const float>() : NULL;
 | 
						|
    auto armconfidencePtr = refineDet ? mArmConfidence.host<const float>() : NULL;
 | 
						|
 | 
						|
    auto boxes      = std::shared_ptr<float>(new float[4 * priorCount], [](float *p) { delete[] p; });
 | 
						|
    auto decodeBoxs = [&boxes, priorCount, variancePtr](const float *priorboxPtr, const float *locationPtr) {
 | 
						|
        for (int i = 0; i < priorCount; i++) {
 | 
						|
            auto loc = locationPtr + i * 4;
 | 
						|
            auto pb  = priorboxPtr + i * 4;
 | 
						|
            auto var = variancePtr + i * 4;
 | 
						|
            auto box = boxes.get() + i * 4;
 | 
						|
 | 
						|
            float pbW  = pb[2] - pb[0];
 | 
						|
            float pbH  = pb[3] - pb[1];
 | 
						|
            float pbCX = (pb[0] + pb[2]) * 0.5f;
 | 
						|
            float pbCY = (pb[1] + pb[3]) * 0.5f;
 | 
						|
 | 
						|
            float boxCX = var[0] * loc[0] * pbW + pbCX;
 | 
						|
            float boxCY = var[1] * loc[1] * pbH + pbCY;
 | 
						|
            float boxW  = exp(var[2] * loc[2]) * pbW;
 | 
						|
            float boxH  = exp(var[3] * loc[3]) * pbH;
 | 
						|
 | 
						|
            box[0] = boxCX - boxW * 0.5f;
 | 
						|
            box[1] = boxCY - boxH * 0.5f;
 | 
						|
            box[2] = boxCX + boxW * 0.5f;
 | 
						|
            box[3] = boxCY + boxH * 0.5f;
 | 
						|
        }
 | 
						|
    };
 | 
						|
    if (refineDet) {
 | 
						|
        decodeBoxs(priorboxPtr, armlocationPtr);
 | 
						|
        decodeBoxs(boxes.get(), locationPtr);
 | 
						|
    } else {
 | 
						|
        decodeBoxs(priorboxPtr, locationPtr);
 | 
						|
    }
 | 
						|
 | 
						|
    // sort and nms for each class
 | 
						|
    std::vector<score_box_t> allClassBoxes;
 | 
						|
    auto compareFunction = [](const score_box_t &a, const score_box_t &b) { return box_score(a) > box_score(b); };
 | 
						|
    {
 | 
						|
        AUTOTIME;
 | 
						|
        for (int i = 1; i < mClassCount; i++) { // start from 1 to ignore background class
 | 
						|
            std::vector<score_box_t> classBoxes;
 | 
						|
            classBoxes.reserve(priorCount);
 | 
						|
            // filter by confidenceThreshold
 | 
						|
            for (int j = 0; j < priorCount; j++) {
 | 
						|
                float score = confidencePtr[j * mClassCount + i];
 | 
						|
                if (refineDet && (armconfidencePtr[j * 2 + 1] < mObjectnessScoreThreshold)) {
 | 
						|
                    score = 0.0;
 | 
						|
                }
 | 
						|
                if (score > mConfidenceThreshold) {
 | 
						|
                    const float *box = boxes.get() + 4 * j;
 | 
						|
                    classBoxes.push_back(box_rect(box[0], box[1], box[2], box[3], i, score));
 | 
						|
                }
 | 
						|
            }
 | 
						|
 | 
						|
            // sort inplace
 | 
						|
            std::sort(classBoxes.begin(), classBoxes.end(), compareFunction);
 | 
						|
 | 
						|
            // apply nms
 | 
						|
            std::vector<int> picked;
 | 
						|
            pickBoxes(classBoxes, picked, mNMSThreshold, mKeepTopK);
 | 
						|
 | 
						|
            // select
 | 
						|
            for (auto index : picked) {
 | 
						|
                allClassBoxes.push_back(classBoxes[index]);
 | 
						|
            }
 | 
						|
        }
 | 
						|
    }
 | 
						|
 | 
						|
    // set width
 | 
						|
    int numDetected = (int)allClassBoxes.size();
 | 
						|
    if (numDetected > mKeepTopK) {
 | 
						|
        numDetected = mKeepTopK;
 | 
						|
    }
 | 
						|
    // global sort inplace
 | 
						|
    {
 | 
						|
        AUTOTIME;
 | 
						|
        std::partial_sort(allClassBoxes.begin(), allClassBoxes.begin() + numDetected, allClassBoxes.end(), compareFunction);
 | 
						|
    }
 | 
						|
    output->buffer().dim[2].extent = numDetected;
 | 
						|
 | 
						|
    // write data
 | 
						|
    auto outPtr = output->host<float>();
 | 
						|
    for (int i = 0; i < numDetected; i++, outPtr += 6 * 4) {
 | 
						|
        auto box      = allClassBoxes[i];
 | 
						|
        outPtr[0 * 4] = box_label(box);
 | 
						|
        outPtr[1 * 4] = box_score(box);
 | 
						|
        outPtr[2 * 4] = box_rect_xmin(box);
 | 
						|
        outPtr[3 * 4] = box_rect_ymin(box);
 | 
						|
        outPtr[4 * 4] = box_rect_xmax(box);
 | 
						|
        outPtr[5 * 4] = box_rect_ymax(box);
 | 
						|
    }
 | 
						|
 | 
						|
    return NO_ERROR;
 | 
						|
}
 | 
						|
 | 
						|
class CPUDetectionOutputCreator : public CPUBackend::Creator {
 | 
						|
public:
 | 
						|
    virtual Execution *onCreate(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs,
 | 
						|
                                const MNN::Op *op, Backend *backend) const {
 | 
						|
        auto d = op->main_as_DetectionOutput();
 | 
						|
        return new CPUDetectionOutput(backend, d->classCount(), d->nmsThresholdold(), d->keepTopK(),
 | 
						|
                                      d->confidenceThreshold(), d->objectnessScore());
 | 
						|
    }
 | 
						|
};
 | 
						|
REGISTER_CPU_OP_CREATOR(CPUDetectionOutputCreator, OpType_DetectionOutput);
 | 
						|
 | 
						|
} // namespace MNN
 |