Fix Reduce Mean error for MobileNets DNN
authorJoe <joegeisbauer@gmail.com>
Tue, 17 Nov 2020 19:58:42 +0000 (13:58 -0600)
committerJoe <joegeisbauer@gmail.com>
Fri, 20 Nov 2020 17:17:02 +0000 (11:17 -0600)
Fix for index error for Reduce Mean

Correct Reduce Mean indexing error

modules/dnn/src/onnx/onnx_importer.cpp

index 56683f4..01d84d9 100644 (file)
@@ -494,14 +494,17 @@ void ONNXImporter::handleNode(const opencv_onnx::NodeProto& node_proto_)
                 MatShape inpShape = outShapes[node_proto.input(0)];
                 DictValue axes = layerParams.get("axes");
                 bool keepdims = layerParams.get<int>("keepdims");
-                MatShape targetShape = inpShape;
+                MatShape targetShape;
+                std::vector<bool> shouldDelete(inpShape.size(), false);
                 for (int i = 0; i < axes.size(); i++) {
                     int axis = clamp(axes.get<int>(i), inpShape.size());
-                    if (keepdims) {
-                        targetShape[axis] = 1;
-                    } else {
-                        targetShape.erase(targetShape.begin() + axis);
-                    }
+                    shouldDelete[axis] = true;
+                }
+                for (int axis = 0; axis < inpShape.size(); ++axis){
+                    if (!shouldDelete[axis])
+                        targetShape.push_back(inpShape[axis]);
+                    else if (keepdims)
+                        targetShape.push_back(1);
                 }
 
                 if (inpShape.size() == 3 && axes.size() <= 2)