diff --git a/modules/dnn/src/net_impl.cpp b/modules/dnn/src/net_impl.cpp index dc2444e994..ba9d6bc36b 100644 --- a/modules/dnn/src/net_impl.cpp +++ b/modules/dnn/src/net_impl.cpp @@ -1631,21 +1631,20 @@ void Net::Impl::setInput(InputArray blob, const String& name, double scalefactor Mat Net::Impl::getParam(int layer, int numParam) const { - LayerData& ld = getLayerData(layer); - std::vector& layerBlobs = getLayerInstance(ld)->blobs; + std::vector& layerBlobs = getLayer(layer)->blobs; CV_Assert(numParam < (int)layerBlobs.size()); return layerBlobs[numParam]; } void Net::Impl::setParam(int layer, int numParam, const Mat& blob) { - LayerData& ld = getLayerData(layer); - // FIXIT we should not modify "execution" instance - std::vector& layerBlobs = getLayerInstance(ld)->blobs; + std::vector& layerBlobs = getLayer(layer)->blobs; CV_Assert(numParam < (int)layerBlobs.size()); // we don't make strong checks, use this function carefully layerBlobs[numParam] = blob; + if (mainGraph) + finalizeLayers = true; } void Net::Impl::setParam(const std::string& outputTensorName, int numParam, const Mat& blob) @@ -1657,9 +1656,16 @@ void Net::Impl::setParam(const std::string& outputTensorName, int numParam, cons if (excl != std::string::npos) it = argnames.find(outputTensorName.substr(excl + 1)); } - if (it == argnames.end()) + if (it == argnames.end()) { + // Not a tensor name; try it as a layer name. + int lid = getLayerId(outputTensorName); + if (lid >= 0) { + setParam(lid, numParam, blob); + return; + } CV_Error_(Error::StsObjectNotFound, ("DNN: tensor '%s' not found in the graph", outputTensorName.c_str())); + } int targetIdx = (int)it->second; const std::vector>& prog = mainGraph->prog(); diff --git a/modules/dnn/src/net_impl2.cpp b/modules/dnn/src/net_impl2.cpp index 6fbabd8da6..e32f82e9b1 100644 --- a/modules/dnn/src/net_impl2.cpp +++ b/modules/dnn/src/net_impl2.cpp @@ -689,6 +689,24 @@ void Net::Impl::forwardMainGraph(InputArrayOfArrays inputs, OutputArrayOfArrays kvCacheManager.applyRoutes(); } +// Assign a single result to an output array, including a (pre-allocated) vector +// of Mat/UMat, which _OutputArray::assign(Mat) does not handle. +static void assignSingleOutput(OutputArrayOfArrays outputBlobs, const Mat& result) +{ + _InputArray::KindFlag k = outputBlobs.kind(); + if (k == _InputArray::STD_VECTOR_MAT) { + std::vector& v = outputBlobs.getMatVecRef(); + v.resize(1); + result.copyTo(v[0]); + } else if (k == _InputArray::STD_VECTOR_UMAT) { + std::vector& v = outputBlobs.getUMatVecRef(); + v.resize(1); + result.copyTo(v[0]); + } else { + outputBlobs.assign(result); + } +} + void Net::Impl::forwardWithSingleOutput(const std::string& outname, OutputArrayOfArrays outputBlobs) { #ifdef HAVE_ONNXRUNTIME @@ -715,7 +733,7 @@ void Net::Impl::forwardWithSingleOutput(const std::string& outname, OutputArrayO std::vector outIdxs(1, outIdx); std::vector outs = runOrtSession(netInputLayer->blobs, outIdxs); CV_Assert(outs.size() == 1); - outputBlobs.assign(outs[0]); + assignSingleOutput(outputBlobs, outs[0]); return; } #endif @@ -743,7 +761,7 @@ void Net::Impl::forwardWithSingleOutput(const std::string& outname, OutputArrayO const std::vector& gr_outputs = mainGraph->outputs(); for (size_t i = 0; i < gr_outputs.size(); i++) { if (gr_outputs[i].idx == targetArg.idx) { - outputBlobs.assign(outs[i]); + assignSingleOutput(outputBlobs, outs[i]); return; } } @@ -760,9 +778,9 @@ void Net::Impl::forwardWithSingleOutput(const std::string& outname, OutputArrayO if (result.shape().layout == DATA_LAYOUT_BLOCK) { Mat converted; transformLayout(result, converted, originalLayout, originalLayout, defaultC0); - outputBlobs.assign(converted); + assignSingleOutput(outputBlobs, converted); } else { - outputBlobs.assign(result.clone()); + assignSingleOutput(outputBlobs, result.clone()); } return; } @@ -771,7 +789,7 @@ void Net::Impl::forwardWithSingleOutput(const std::string& outname, OutputArrayO std::vector inps, outs; forwardMainGraph(inps, outs); CV_Assert(!outs.empty()); - outputBlobs.assign(outs[0]); + assignSingleOutput(outputBlobs, outs[0]); } void Net::Impl::forwardWithMultipleOutputs(OutputArrayOfArrays outblobs, const std::vector& outnames)