Skip to content

Commit

Permalink
Option for gpu for image segmentation
Browse files Browse the repository at this point in the history
  • Loading branch information
@servantftechnicolor @cbentejac
servantftechnicolor authored and cbentejac committed Dec 11, 2023
1 parent 19d4fd2 commit b5cdfa2
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 31 deletions.
61 changes: 37 additions & 24 deletions src/aliceVision/segmentation/segmentation.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,33 +50,40 @@ bool Segmentation::initialize()

Ort::SessionOptions ortSessionOptions;

#if ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_ONNX_GPU)
OrtCUDAProviderOptionsV2* cuda_options = nullptr;
api.CreateCUDAProviderOptions(&cuda_options);
api.SessionOptionsAppendExecutionProvider_CUDA_V2(static_cast<OrtSessionOptions*>(ortSessionOptions), cuda_options);
api.ReleaseCUDAProviderOptions(cuda_options);

#if defined(_WIN32) || defined(_WIN64)
std::wstring modelWeights(_parameters.modelWeights.begin(), _parameters.modelWeights.end());
_ortSession = std::make_unique<Ort::Session>(*_ortEnvironment, modelWeights.c_str(), ortSessionOptions);
#else
_ortSession = std::make_unique<Ort::Session>(*_ortEnvironment, _parameters.modelWeights.c_str(), ortSessionOptions);
#endif
// this is false if ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_ONNX_GPU) is false
if (_parameters.useGpu)
{
//Disable for compilation purpose if needed
#if ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_ONNX_GPU)
OrtCUDAProviderOptionsV2* cuda_options = nullptr;
api.CreateCUDAProviderOptions(&cuda_options);
api.SessionOptionsAppendExecutionProvider_CUDA_V2(static_cast<OrtSessionOptions*>(ortSessionOptions), cuda_options);
api.ReleaseCUDAProviderOptions(cuda_options);

#if defined(_WIN32) || defined(_WIN64)
std::wstring modelWeights(_parameters.modelWeights.begin(), _parameters.modelWeights.end());
_ortSession = std::make_unique<Ort::Session>(*_ortEnvironment, modelWeights.c_str(), ortSessionOptions);
#else
_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);
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));
#else
_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
{
#if defined(_WIN32) || defined(_WIN64)
std::wstring modelWeights(_parameters.modelWeights.begin(), _parameters.modelWeights.end());
_ortSession = std::make_unique<Ort::Session>(*_ortEnvironment, modelWeights.c_str(), ortSessionOptions);
#else
_ortSession = std::make_unique<Ort::Session>(*_ortEnvironment, _parameters.modelWeights.c_str(), ortSessionOptions);
#endif
#endif
}

return true;
}
Expand Down Expand Up @@ -199,11 +206,14 @@ bool Segmentation::tiledProcess(image::Image<IndexT>& labels, const image::Image
// Compute tile
image::Image<ScoredLabel> tileLabels(_parameters.modelWidth, _parameters.modelHeight, true, {0, 0.0f});

#if ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_ONNX_GPU)
if (_parameters.useGpu)
{
processTileGPU(tileLabels, block);
#else
}
else
{
processTile(tileLabels, block);
#endif
}

// Update the global labeling
mergeLabels(scoredLabels, tileLabels, x, y);
Expand Down Expand Up @@ -265,6 +275,7 @@ bool Segmentation::labelsFromModelOutput(image::Image<ScoredLabel>& labels, cons

bool Segmentation::processTile(image::Image<ScoredLabel>& labels, const image::Image<image::RGBfColor>::Base& source)
{
ALICEVISION_LOG_TRACE("Process tile using cpu");
Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);

std::vector<const char*> inputNames{"input"};
Expand Down Expand Up @@ -300,9 +311,10 @@ bool Segmentation::processTile(image::Image<ScoredLabel>& labels, const image::I
return true;
}

#if ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_CUDA)
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);

Expand Down Expand Up @@ -339,9 +351,10 @@ bool Segmentation::processTileGPU(image::Image<ScoredLabel>& labels, const image
return false;
}

#endif

return true;
}
#endif

} // namespace segmentation
} // namespace aliceVision
18 changes: 11 additions & 7 deletions src/aliceVision/segmentation/segmentation.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,12 +39,18 @@ class Segmentation
int modelWidth;
int modelHeight;
double overlapRatio;
bool useGpu = true;
};

public:
Segmentation(const Parameters& parameters)
: _parameters(parameters)
{
//Disable gpu if disabled on compilation side
#if !ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_ONNX_GPU)
_parameters.useGpu = false;
#endif

if (!initialize())
{
throw std::runtime_error("Error on segmentation initialization");
Expand Down Expand Up @@ -92,14 +98,12 @@ class Segmentation
*/
bool processTile(image::Image<ScoredLabel>& labels, const image::Image<image::RGBfColor>::Base& source);

/**
* Process effectively a buffer of the model input size
* param labels the output labels
* @param source the source tile
*/
#if ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_CUDA)
/**
* Process effectively a buffer of the model input size
* param labels the output labels
* @param source the source tile
*/
bool processTileGPU(image::Image<ScoredLabel>& labels, const image::Image<image::RGBfColor>::Base& source);
#endif

/**
* Merge tile labels with global labels image
Expand Down
4 changes: 4 additions & 0 deletions src/software/pipeline/main_imageSegmentation.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,7 @@ int aliceVision_main(int argc, char** argv)
bool maskInvert = false;
int rangeStart = -1;
int rangeSize = 1;
bool useGpu = true;

// Description of mandatory parameters
po::options_description requiredParams("Required parameters");
Expand All @@ -104,6 +105,8 @@ int aliceVision_main(int argc, char** argv)
"Names of classes which are to be considered.")
("maskInvert", po::value<bool>(&maskInvert)->default_value(maskInvert),
"Invert mask values. If selected, the pixels corresponding to the mask will be set to 0.0 instead of 1.0.")
("useGpu", po::value<bool>(&useGpu)->default_value(useGpu),
"Use GPU if available.")
("rangeStart", po::value<int>(&rangeStart)->default_value(rangeStart),
"Range start for processing views (ordered by image filepath). Set to -1 to process all images.")
("rangeSize", po::value<int>(&rangeSize)->default_value(rangeSize),
Expand Down Expand Up @@ -178,6 +181,7 @@ int aliceVision_main(int argc, char** argv)
parameters.modelWidth = 1280;
parameters.modelHeight = 720;
parameters.overlapRatio = 0.3;
parameters.useGpu = useGpu;

aliceVision::segmentation::Segmentation seg(parameters);

Expand Down

0 comments on commit b5cdfa2

Please sign in to comment.