Update model output processing

Author: demoulinv

src/aliceVision/segmentation/segmentation.cpp

@@ -244,7 +244,7 @@ bool Segmentation::mergeLabels(image::Image<ScoredLabel>& labels, image::Image<S
     return true;
 }
 
-bool Segmentation::labelsFromModelOutput(image::Image<ScoredLabel>& labels, const std::vector<float>& modelOutput)
+bool Segmentation::labelsFromOutputTensor(image::Image<ScoredLabel>& labels, Ort::Value& modelOutput)
 {
     for (int outputY = 0; outputY < _parameters.modelHeight; outputY++)
     {
@@ -255,10 +255,8 @@ bool Segmentation::labelsFromModelOutput(image::Image<ScoredLabel>& labels, cons
 
             for (int classe = 0; classe < _parameters.classes.size(); classe++)
             {
-                int classPos = classe * _parameters.modelWidth * _parameters.modelHeight;
-                int pos = classPos + outputY * _parameters.modelWidth + outputX;
-
-                float val = modelOutput[pos];
+                const std::vector<int64_t> coords = {0,classe,outputY,outputX};
+                const float val = modelOutput.At<float>(coords);
                 if (val > maxVal)
                 {
                     maxVal = val;
@@ -281,29 +279,40 @@ bool Segmentation::processTile(image::Image<ScoredLabel>& labels, const image::I
     std::vector<const char*> inputNames{"input"};
     std::vector<const char*> outputNames{"output"};
     std::vector<int64_t> inputDimensions = {1, 3, _parameters.modelHeight, _parameters.modelWidth};
-    std::vector<int64_t> outputDimensions = {1, static_cast<int64_t>(_parameters.classes.size()), _parameters.modelHeight, _parameters.modelWidth};
-
-    std::vector<float> output(_parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);
-    Ort::Value outputTensors =
-      Ort::Value::CreateTensor<float>(memInfo, output.data(), output.size(), outputDimensions.data(), outputDimensions.size());
 
     std::vector<float> transformedInput;
     imageToPlanes(transformedInput, source);
 
     Ort::Value inputTensors =
       Ort::Value::CreateTensor<float>(memInfo, transformedInput.data(), transformedInput.size(), inputDimensions.data(), inputDimensions.size());
 
+    std::vector<Ort::Value> outTensor;
+
     try
     {
-        _ortSession->Run(Ort::RunOptions{nullptr}, inputNames.data(), &inputTensors, 1, outputNames.data(), &outputTensors, 1);
+        outTensor = _ortSession->Run(Ort::RunOptions{nullptr}, inputNames.data(), &inputTensors, 1, outputNames.data(), 1);
     }
     catch (const Ort::Exception& exception)
     {
         ALICEVISION_LOG_ERROR("ERROR running model inference: " << exception.what());
         return false;
     }
 
-    if (!labelsFromModelOutput(labels, output))
+    std::vector<float> output(_parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);
+    int idx = 0;
+    for (int ch = 0; ch < _parameters.classes.size(); ch++)
+    {
+        for (int i = 0; i < _parameters.modelHeight; i++)
+        {
+            for (int j = 0; j < _parameters.modelWidth; j++)
+            {
+                const std::vector<int64_t> coords = {0, ch, i, j};
+                output[idx++] = outTensor[0].At<float>(coords);
+            }
+        }
+    }
+
+    if (!labelsFromOutputTensor(labels, outTensor[0]))
     {
         return false;
     }
@@ -321,10 +330,6 @@ bool Segmentation::processTileGPU(image::Image<ScoredLabel>& labels, const image
     std::vector<const char*> inputNames{"input"};
     std::vector<const char*> outputNames{"output"};
     std::vector<int64_t> inputDimensions = {1, 3, _parameters.modelHeight, _parameters.modelWidth};
-    std::vector<int64_t> outputDimensions = {1, static_cast<int64_t>(_parameters.classes.size()), _parameters.modelHeight, _parameters.modelWidth};
-
-    Ort::Value outputTensors = Ort::Value::CreateTensor<float>(
-      mem_info_cuda, reinterpret_cast<float*>(_cudaOutput), _output.size(), outputDimensions.data(), outputDimensions.size());
 
     std::vector<float> transformedInput;
     imageToPlanes(transformedInput, source);
@@ -334,19 +339,32 @@ bool Segmentation::processTileGPU(image::Image<ScoredLabel>& labels, const image
     Ort::Value inputTensors = Ort::Value::CreateTensor<float>(
       mem_info_cuda, reinterpret_cast<float*>(_cudaInput), transformedInput.size(), inputDimensions.data(), inputDimensions.size());
 
+    std::vector<Ort::Value> outTensor;
+
     try
     {
-        _ortSession->Run(Ort::RunOptions{nullptr}, inputNames.data(), &inputTensors, 1, outputNames.data(), &outputTensors, 1);
+        outTensor = _ortSession->Run(Ort::RunOptions{nullptr}, inputNames.data(), &inputTensors, 1, outputNames.data(), 1);
     }
     catch (const Ort::Exception& exception)
     {
         ALICEVISION_LOG_ERROR("ERROR running model inference: " << exception.what());
         return false;
     }
 
-    cudaMemcpy(_output.data(), _cudaOutput, sizeof(float) * _output.size(), cudaMemcpyDeviceToHost);
+    int idx = 0;
+    for (int ch = 0; ch < _parameters.classes.size(); ch++)
+    {
+        for (int i = 0; i < _parameters.modelHeight; i++)
+        {
+            for (int j = 0; j < _parameters.modelWidth; j++)
+            {
+                const std::vector<int64_t> coords = {0, ch, i, j};
+                _output[idx++] = outTensor[0].At<float>(coords);
+            }
+        }
+    }
 
-    if (!labelsFromModelOutput(labels, _output))
+    if (!labelsFromOutputTensor(labels, outTensor[0]))
     {
         return false;
     }

src/aliceVision/segmentation/segmentation.hpp

@@ -86,11 +86,18 @@ class Segmentation
 
     /**
      * Transform model output to a label image
-     * @param labels the output labels imaage
+     * @param labels the output labels image
      * @param modeloutput the model output vector
      */
     bool labelsFromModelOutput(image::Image<ScoredLabel>& labels, const std::vector<float>& modelOutput);
 
+    /**
+     * Transform model output to a label image
+     * @param labels the output labels image
+     * @param modeloutput the model output tensor
+     */
+    bool labelsFromOutputTensor(image::Image<ScoredLabel>& labels, Ort::Value& modelOutput);
+
     /**
      * Process effectively a buffer of the model input size
      * param labels the output labels