mirror of
https://github.com/opencv/opencv.git
synced 2024-11-24 03:00:14 +08:00
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:
parent
5393185add
commit
d5e8792f55
@ -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);
|
||||
|
@ -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]
|
||||
|
Loading…
Reference in New Issue
Block a user