mirror of
https://github.com/opencv/opencv.git
synced 2025-06-11 20:09:23 +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,
|
inline void preprocess(const Mat& frame, Net& net, Size inpSize, float scale,
|
||||||
const Scalar& mean, bool swapRB);
|
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);
|
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.
|
// Load a model.
|
||||||
Net net = readNet(modelPath, configPath, parser.get<String>("framework"));
|
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"));
|
net.setPreferableTarget(parser.get<int>("target"));
|
||||||
std::vector<String> outNames = net.getUnconnectedOutLayersNames();
|
std::vector<String> outNames = net.getUnconnectedOutLayersNames();
|
||||||
|
|
||||||
@ -245,7 +246,7 @@ int main(int argc, char** argv)
|
|||||||
std::vector<Mat> outs = predictionsQueue.get();
|
std::vector<Mat> outs = predictionsQueue.get();
|
||||||
Mat frame = processedFramesQueue.get();
|
Mat frame = processedFramesQueue.get();
|
||||||
|
|
||||||
postprocess(frame, outs, net);
|
postprocess(frame, outs, net, backend);
|
||||||
|
|
||||||
if (predictionsQueue.counter > 1)
|
if (predictionsQueue.counter > 1)
|
||||||
{
|
{
|
||||||
@ -285,7 +286,7 @@ int main(int argc, char** argv)
|
|||||||
std::vector<Mat> outs;
|
std::vector<Mat> outs;
|
||||||
net.forward(outs, outNames);
|
net.forward(outs, outNames);
|
||||||
|
|
||||||
postprocess(frame, outs, net);
|
postprocess(frame, outs, net, backend);
|
||||||
|
|
||||||
// Put efficiency information.
|
// Put efficiency information.
|
||||||
std::vector<double> layersTimes;
|
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::vector<int> outLayers = net.getUnconnectedOutLayers();
|
||||||
static std::string outLayerType = net.getLayer(outLayers[0])->type;
|
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
|
else
|
||||||
CV_Error(Error::StsNotImplemented, "Unknown output layer type: " + outLayerType);
|
CV_Error(Error::StsNotImplemented, "Unknown output layer type: " + outLayerType);
|
||||||
|
|
||||||
std::vector<int> indices;
|
// NMS is used inside Region layer only on DNN_BACKEND_OPENCV for another backends we need NMS in sample
|
||||||
NMSBoxes(boxes, confidences, confThreshold, nmsThreshold, indices);
|
// or NMS is required if number of outputs > 1
|
||||||
for (size_t i = 0; i < indices.size(); ++i)
|
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];
|
Rect box = boxes[idx];
|
||||||
drawPred(classIds[idx], confidences[idx], box.x, box.y,
|
drawPred(classIds[idx], confidences[idx], box.x, box.y,
|
||||||
box.x + box.width, box.y + box.height, frame);
|
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
|
# 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
|
# detected objects and C is a number of classes + 4 where the first 4
|
||||||
# numbers are [center_x, center_y, width, height]
|
# numbers are [center_x, center_y, width, height]
|
||||||
classIds = []
|
|
||||||
confidences = []
|
|
||||||
boxes = []
|
|
||||||
for out in outs:
|
for out in outs:
|
||||||
for detection in out:
|
for detection in out:
|
||||||
scores = detection[5:]
|
scores = detection[5:]
|
||||||
@ -163,9 +160,25 @@ def postprocess(frame, outs):
|
|||||||
print('Unknown output layer type: ' + lastLayer.type)
|
print('Unknown output layer type: ' + lastLayer.type)
|
||||||
exit()
|
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:
|
for i in indices:
|
||||||
i = i[0]
|
|
||||||
box = boxes[i]
|
box = boxes[i]
|
||||||
left = box[0]
|
left = box[0]
|
||||||
top = box[1]
|
top = box[1]
|
||||||
|
Loading…
Reference in New Issue
Block a user