add GPU device count check

jveitchmichaelis · Sep 24, 2021 · d8f131e · d8f131e
1 parent c7634ac
commit d8f131e
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 6 deletions.
diff --git a/DeepLabel.pro b/DeepLabel.pro
@@ -23,7 +23,6 @@ DEFINES += QT_DEPRECATED_WARNINGS
 #DEFINES += QT_DISABLE_DEPRECATED_BEFORE=0x060000 # disables all the APIs deprecated before Qt 6.0.0
 
 CONFIG += WITH_CUDA
-DEFINES += WITH_CUDA
 
 macx{
 message("Mac")
@@ -53,13 +52,15 @@ win32{
 message("Windows")
 DEFINES += PROTOBUF_USE_DLLS
 
-defined(WITH_CUDA){
+CONFIG(WITH_CUDA){
+DEFINES += WITH_CUDA
 INCLUDEPATH += "$$_PRO_FILE_PWD_/opencv/build_cuda/include"
 LIBS += -L"$$_PRO_FILE_PWD_/opencv/build_cuda/x64/vc15/lib"
 }else{
 INCLUDEPATH += "$$_PRO_FILE_PWD_/opencv/build/include"
 LIBS += -L"$$_PRO_FILE_PWD_/opencv/build/x64/vc15/lib"
 }
+
 INCLUDEPATH += "$$_PRO_FILE_PWD_/protobuf/include"
 LIBS += -L"$$_PRO_FILE_PWD_/protobuf/lib"
 

diff --git a/src/detection/detectoropencv.cpp b/src/detection/detectoropencv.cpp
@@ -93,8 +93,16 @@ void DetectorOpenCV::setTarget(int target){
  }
 #ifdef WITH_CUDA
  else if(preferable_target == cv::dnn::DNN_TARGET_CUDA || preferable_target == cv::dnn::DNN_TARGET_CUDA_FP16){
- // Check for GPU
  try{
+ // Check for GPU
+ if(cv::cuda::getCudaEnabledDeviceCount() == 0){
+ qCritical() << "No CUDA enabled devices found";
+ preferable_target = cv::dnn::DNN_TARGET_CPU;
+ return;
+ }else{
+ cv::cuda::setDevice(0);
+ }
+
  auto devinfo = cv::cuda::DeviceInfo();
  if (!devinfo.isCompatible()){
  qWarning() << "Device is not CUDA compatible. Falling back to CPU";
@@ -103,7 +111,8 @@ void DetectorOpenCV::setTarget(int target){
  qInfo() << "NVIDIA GPU detected: " << devinfo.name();
  }
  }catch(...){
- qCritical() << "Problem checking for GPU, defaulting to CPU inference mode. Check that OpenCV was compiled with CUDA?";
+ qCritical() << "Problem calling GPU functions, defaulting to CPU inference mode.";
+ qCritical() << "Check that OpenCV was compiled with CUDA or try rebooting.";
  preferable_target = cv::dnn::DNN_TARGET_CPU;
  }
  }
@@ -306,6 +315,7 @@ std::vector<BoundingBox> DetectorOpenCV::infer(cv::Mat image){
  cv::cvtColor(image, image, cv::COLOR_BGRA2BGR);
  }
 
+ // 16-bit conversion, min-max scaling
  if(convert_depth && image.elemSize() == 2){
  double minval, maxval;
  cv::minMaxIdx(image, &minval, &maxval);
@@ -384,6 +394,11 @@ std::vector<BoundingBox> DetectorOpenCV::inferTensorflow(cv::Mat image){
  return results;
 }
 
+void DetectorOpenCV::setNormalisation(double scale_factor, cv::Scalar mean){
+ this->scale_factor = scale_factor;
+ this->mean = mean;
+}
+
 std::vector<BoundingBox> DetectorOpenCV::inferDarknet(cv::Mat image){
 
  std::vector<BoundingBox> results;
@@ -393,10 +408,12 @@ std::vector<BoundingBox> DetectorOpenCV::inferDarknet(cv::Mat image){
  mean = cv::Scalar(0);
  }
 
- // Check for 16-bit
- double scale_factor = 1/255.0;
+ // Normalisation, check for 16-bit
+ // TODO: have this as a user option
  if(image.elemSize() == 2){
  scale_factor = 1/65535.0;
+ }else{
+ scale_factor= 1/255.0;
  }
 
  auto input_size = cv::Size(input_width, input_height);

diff --git a/src/detection/detectoropencv.h b/src/detection/detectoropencv.h
@@ -52,12 +52,14 @@ class DetectorOpenCV
  void setTarget(int target);
  void setTarget(QString target);
  void setChannels(int channels);
+ void setNormalisation(double scale_factor, cv::Scalar mean);
  int getChannels(void){return input_channels;}
 
  static model_framework frameworkFromString(QString framework_string);
  static int targetFromString(QString target_string);
 
  void runOnProject(LabelProject *project);
+
 private:
 
  void postProcess(cv::Mat& frame, const std::vector<cv::Mat>& outs, std::vector<BoundingBox> &filtered_outputs);
@@ -74,6 +76,8 @@ class DetectorOpenCV
  int input_channels = 3;
  int preferable_target = cv::dnn::DNN_TARGET_OPENCL;
  model_framework framework = FRAMEWORK_DARKNET;
+ cv::Scalar mean = {0,0,0};
+ double scale_factor = 1/255.;
 
  std::vector<std::string> class_names;
  std::vector<std::string> output_names;