Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ImageSegmentation: add an option to choose between cpu and gpu #1618

Merged
merged 1 commit into from
Dec 13, 2023
Merged

ImageSegmentation: add an option to choose between cpu and gpu #1618

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 37 additions & 24 deletions src/aliceVision/segmentation/segmentation.cpp
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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