| 103 | } |
| 104 | |
| 105 | std::vector<Express::VARP> NMSModule::onForward(const std::vector<Express::VARP>& inputs) { |
| 106 | const int maxDetections = inputs[2]->readMap<int>()[0]; |
| 107 | float iouThreshold = 0, scoreThreshold = std::numeric_limits<float>::lowest(); |
| 108 | if (inputs.size() > 3) { |
| 109 | iouThreshold = inputs[3]->readMap<float>()[0]; |
| 110 | } |
| 111 | if (inputs.size() > 4) { |
| 112 | scoreThreshold = inputs[4]->readMap<float>()[0]; |
| 113 | } |
| 114 | auto boxes = inputs[0], score = inputs[1]; |
| 115 | auto info = boxes->getInfo(), infoScore = score->getInfo(); |
| 116 | MNN_ASSERT(info->dim[info->dim.size() - 1] == 4); |
| 117 | int batch = 1, numClass = 1, numBoxes = info->dim[0]; |
| 118 | bool onnxFormat = (infoScore->dim.size() > 1); |
| 119 | if (onnxFormat) { |
| 120 | batch = infoScore->dim[0]; |
| 121 | numClass = infoScore->dim[1]; |
| 122 | numBoxes = infoScore->dim[2]; |
| 123 | } |
| 124 | INTS outputData; |
| 125 | for (int b = 0; b < batch; ++b) { |
| 126 | const auto boxesPtr = boxes->readMap<float>() + b * numBoxes * 4; |
| 127 | for (int c = 0; c < numClass; ++c) { |
| 128 | std::vector<int> selected; |
| 129 | const auto scorePtr = score->readMap<float>() + (b * numClass + c) * numBoxes; |
| 130 | NonMaxSuppressionSingleClasssImpl(boxesPtr, scorePtr, numBoxes, maxDetections, iouThreshold, scoreThreshold, &selected); |
| 131 | for (int i = 0; i < selected.size(); ++i) { |
| 132 | if (onnxFormat) { |
| 133 | outputData.push_back(b); |
| 134 | outputData.push_back(c); |
| 135 | } |
| 136 | outputData.push_back(selected[i]); |
| 137 | } |
| 138 | } |
| 139 | } |
| 140 | |
| 141 | Variable::Info outInfo; |
| 142 | outInfo.order = info->order; |
| 143 | outInfo.type = halide_type_of<int>(); |
| 144 | if (onnxFormat) { |
| 145 | outInfo.dim.assign({(int)outputData.size() / 3, 3}); |
| 146 | } else { |
| 147 | outInfo.dim.assign({(int)outputData.size()}); |
| 148 | } |
| 149 | outInfo.syncSize(); |
| 150 | VARPS outputs; |
| 151 | outputs.push_back(Variable::create(Expr::create(std::move(outInfo), outputData.data(), VARP::CONSTANT))); |
| 152 | return outputs; |
| 153 | } |
| 154 | |
| 155 | Module* NMSModule::clone(CloneContext* ctx) const { |
| 156 | NMSModule* module(new NMSModule); |