Skip to content

Commit

Permalink
Merge pull request #1783 from alicevision/bugfix/imageSegmentationOnn…
Browse files Browse the repository at this point in the history 
…xOutput

[segmentation] Image segmentation CPU bugfix and ONNX output update
  • Loading branch information
@cbentejac
cbentejac authored Feb 13, 2025
2 parents 5716cee + 86d5e05 commit 31039fc
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 51 deletions.
66 changes: 24 additions & 42 deletions src/aliceVision/segmentation/segmentation.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -67,12 +67,7 @@ bool Segmentation::initialize()
_ortSession = std::make_unique<Ort::Session>(*_ortEnvironment, _parameters.modelWeights.c_str(), ortSessionOptions);
#endif

Ort::MemoryInfo memInfoCuda("Cuda", OrtAllocatorType::OrtArenaAllocator, 0, OrtMemType::OrtMemTypeDefault);
Ort::Allocator cudaAllocator(*_ortSession, memInfoCuda);

_output.resize(_parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);
_cudaInput = cudaAllocator.Alloc(_output.size() * sizeof(float));
_cudaOutput = cudaAllocator.Alloc(_output.size() * sizeof(float));
#endif
}
else
Expand All @@ -88,18 +83,6 @@ bool Segmentation::initialize()
return true;
}

bool Segmentation::terminate()
{
#if ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_ONNX_GPU)
Ort::MemoryInfo mem_info_cuda("Cuda", OrtAllocatorType::OrtArenaAllocator, 0, OrtMemType::OrtMemTypeDefault);
Ort::Allocator cudaAllocator(*_ortSession, mem_info_cuda);
cudaAllocator.Free(_cudaInput);
cudaAllocator.Free(_cudaOutput);
#endif

return true;
}

bool Segmentation::processImage(image::Image<IndexT>& labels, const image::Image<image::RGBfColor>& source)
{
// Todo : handle orientation and small images smaller than model input
Expand Down Expand Up @@ -244,7 +227,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++)
{
Expand All @@ -255,10 +238,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;
Expand All @@ -281,76 +262,77 @@ 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))
if (!labelsFromOutputTensor(labels, outTensor[0]))
{
return false;
}

std::vector<float> output(_parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);
auto *outTData = outTensor.front().GetTensorMutableData<float>();
output.assign(outTData, outTData + _parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);

return true;
}

bool Segmentation::processTileGPU(image::Image<ScoredLabel>& labels, const image::Image<image::RGBfColor>::Base& source)
{
ALICEVISION_LOG_TRACE("Process tile using gpu");
#if ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_CUDA)
Ort::MemoryInfo mem_info_cuda("Cuda", OrtAllocatorType::OrtArenaAllocator, 0, OrtMemType::OrtMemTypeDefault);
Ort::Allocator cudaAllocator(*_ortSession, mem_info_cuda);
Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);

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);

cudaMemcpy(_cudaInput, transformedInput.data(), sizeof(float) * transformedInput.size(), cudaMemcpyHostToDevice);
std::vector<Ort::Value> inputTensors;
inputTensors.emplace_back(Ort::Value::CreateTensor<float>(memInfo,
transformedInput.data(),
transformedInput.size(),
inputDimensions.data(),
inputDimensions.size()));

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.data(), 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);

if (!labelsFromModelOutput(labels, _output))
if (!labelsFromOutputTensor(labels, outTensor[0]))
{
return false;
}

auto *outTData = outTensor.front().GetTensorMutableData<float>();
_output.assign(outTData, outTData + _parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);

#endif

return true;
Expand Down
13 changes: 4 additions & 9 deletions src/aliceVision/segmentation/segmentation.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ class Segmentation
}
}

virtual ~Segmentation() { terminate(); }
virtual ~Segmentation() {}

/**
* Process an input image to estimate segmentation
Expand All @@ -72,11 +72,6 @@ class Segmentation
*/
bool initialize();

/**
* Onnx destruction code
*/
bool terminate();

/**
* Assume the source image is the correct size
* @param labels the output label image
Expand All @@ -86,10 +81,10 @@ class Segmentation

/**
* Transform model output to a label image
* @param labels the output labels imaage
* @param modeloutput the model output vector
* @param labels the output labels image
* @param modeloutput the model output tensor
*/
bool labelsFromModelOutput(image::Image<ScoredLabel>& labels, const std::vector<float>& modelOutput);
bool labelsFromOutputTensor(image::Image<ScoredLabel>& labels, Ort::Value& modelOutput);

/**
* Process effectively a buffer of the model input size
Expand Down

0 comments on commit 31039fc

Please sign in to comment.