Merge pull request #17332 from l-bat:fix_nms

Fixed NMSBoxes bug

* Added NMS for each class

* Updated cpp sample

* Fixed errors

* Refactoring

* Added NMS for IE
This commit is contained in:
Liubov Batanina 2020-05-25 15:34:11 +03:00 committed by GitHub
parent 5393185add
commit d5e8792f55
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 65 additions and 14 deletions

View File

@ -45,7 +45,7 @@ std::vector<std::string> classes;
inline void preprocess(const Mat& frame, Net& net, Size inpSize, float scale,
const Scalar& mean, bool swapRB);
void postprocess(Mat& frame, const std::vector<Mat>& out, Net& net);
void postprocess(Mat& frame, const std::vector<Mat>& out, Net& net, int backend);
void drawPred(int classId, float conf, int left, int top, int right, int bottom, Mat& frame);
@ -148,7 +148,8 @@ int main(int argc, char** argv)
// Load a model.
Net net = readNet(modelPath, configPath, parser.get<String>("framework"));
net.setPreferableBackend(parser.get<int>("backend"));
int backend = parser.get<int>("backend");
net.setPreferableBackend(backend);
net.setPreferableTarget(parser.get<int>("target"));
std::vector<String> outNames = net.getUnconnectedOutLayersNames();
@ -245,7 +246,7 @@ int main(int argc, char** argv)
std::vector<Mat> outs = predictionsQueue.get();
Mat frame = processedFramesQueue.get();
postprocess(frame, outs, net);
postprocess(frame, outs, net, backend);
if (predictionsQueue.counter > 1)
{
@ -285,7 +286,7 @@ int main(int argc, char** argv)
std::vector<Mat> outs;
net.forward(outs, outNames);
postprocess(frame, outs, net);
postprocess(frame, outs, net, backend);
// Put efficiency information.
std::vector<double> layersTimes;
@ -319,7 +320,7 @@ inline void preprocess(const Mat& frame, Net& net, Size inpSize, float scale,
}
}
void postprocess(Mat& frame, const std::vector<Mat>& outs, Net& net)
void postprocess(Mat& frame, const std::vector<Mat>& outs, Net& net, int backend)
{
static std::vector<int> outLayers = net.getUnconnectedOutLayers();
static std::string outLayerType = net.getLayer(outLayers[0])->type;
@ -396,11 +397,48 @@ void postprocess(Mat& frame, const std::vector<Mat>& outs, Net& net)
else
CV_Error(Error::StsNotImplemented, "Unknown output layer type: " + outLayerType);
std::vector<int> indices;
NMSBoxes(boxes, confidences, confThreshold, nmsThreshold, indices);
for (size_t i = 0; i < indices.size(); ++i)
// NMS is used inside Region layer only on DNN_BACKEND_OPENCV for another backends we need NMS in sample
// or NMS is required if number of outputs > 1
if (outLayers.size() > 1 || (outLayerType == "Region" && backend != DNN_BACKEND_OPENCV))
{
std::map<int, std::vector<size_t> > class2indices;
for (size_t i = 0; i < classIds.size(); i++)
{
if (confidences[i] >= confThreshold)
{
class2indices[classIds[i]].push_back(i);
}
}
std::vector<Rect> nmsBoxes;
std::vector<float> nmsConfidences;
std::vector<int> nmsClassIds;
for (std::map<int, std::vector<size_t> >::iterator it = class2indices.begin(); it != class2indices.end(); ++it)
{
std::vector<Rect> localBoxes;
std::vector<float> localConfidences;
std::vector<size_t> classIndices = it->second;
for (size_t i = 0; i < classIndices.size(); i++)
{
localBoxes.push_back(boxes[classIndices[i]]);
localConfidences.push_back(confidences[classIndices[i]]);
}
std::vector<int> nmsIndices;
NMSBoxes(localBoxes, localConfidences, confThreshold, nmsThreshold, nmsIndices);
for (size_t i = 0; i < nmsIndices.size(); i++)
{
size_t idx = nmsIndices[i];
nmsBoxes.push_back(localBoxes[idx]);
nmsConfidences.push_back(localConfidences[idx]);
nmsClassIds.push_back(it->first);
}
}
boxes = nmsBoxes;
classIds = nmsClassIds;
confidences = nmsConfidences;
}
for (size_t idx = 0; idx < boxes.size(); ++idx)
{
int idx = indices[i];
Rect box = boxes[idx];
drawPred(classIds[idx], confidences[idx], box.x, box.y,
box.x + box.width, box.y + box.height, frame);

View File

@ -141,9 +141,6 @@ def postprocess(frame, outs):
# Network produces output blob with a shape NxC where N is a number of
# detected objects and C is a number of classes + 4 where the first 4
# numbers are [center_x, center_y, width, height]
classIds = []
confidences = []
boxes = []
for out in outs:
for detection in out:
scores = detection[5:]
@ -163,9 +160,25 @@ def postprocess(frame, outs):
print('Unknown output layer type: ' + lastLayer.type)
exit()
indices = cv.dnn.NMSBoxes(boxes, confidences, confThreshold, nmsThreshold)
# NMS is used inside Region layer only on DNN_BACKEND_OPENCV for another backends we need NMS in sample
# or NMS is required if number of outputs > 1
if len(outNames) > 1 or lastLayer.type == 'Region' and args.backend != cv.dnn.DNN_BACKEND_OPENCV:
indices = []
classIds = np.array(classIds)
boxes = np.array(boxes)
confidences = np.array(confidences)
unique_classes = set(classIds)
for cl in unique_classes:
class_indices = np.where(classIds == cl)[0]
conf = confidences[class_indices]
box = boxes[class_indices].tolist()
nms_indices = cv.dnn.NMSBoxes(box, conf, confThreshold, nmsThreshold)
nms_indices = nms_indices[:, 0] if len(nms_indices) else []
indices.extend(class_indices[nms_indices])
else:
indices = np.arange(0, len(classIds))
for i in indices:
i = i[0]
box = boxes[i]
left = box[0]
top = box[1]