diff --git a/src/yolo_v2_class.cpp b/src/yolo_v2_class.cpp index 980db93a64c..8fcf9358bb8 100644 --- a/src/yolo_v2_class.cpp +++ b/src/yolo_v2_class.cpp @@ -60,6 +60,28 @@ int dispose() { return 1; } +int get_device_count() { +#ifdef GPU + int count = 0; + cudaGetDeviceCount(&count); + return count; +#else + return -1; +#endif // GPU +} + +int get_device_name(int gpu, char* deviceName) { +#ifdef GPU + cudaDeviceProp prop; + cudaGetDeviceProperties(&prop, gpu); + std::string result = prop.name; + std::copy(result.begin(), result.end(), deviceName); + return 1; +#else + return -1; +#endif // GPU +} + #ifdef GPU void check_cuda(cudaError_t status) { if (status != cudaSuccess) { diff --git a/src/yolo_v2_class.hpp b/src/yolo_v2_class.hpp index 8626d97f0cb..6f7e47bdcb7 100644 --- a/src/yolo_v2_class.hpp +++ b/src/yolo_v2_class.hpp @@ -49,6 +49,8 @@ extern "C" YOLODLL_API int init(const char *configurationFilename, const char *w extern "C" YOLODLL_API int detect_image(const char *filename, bbox_t_container &container); extern "C" YOLODLL_API int detect_mat(const uint8_t* data, const size_t data_length, bbox_t_container &container); extern "C" YOLODLL_API int dispose(); +extern "C" YOLODLL_API int get_device_count(); +extern "C" YOLODLL_API int get_device_name(int gpu, char* deviceName); class Detector { std::shared_ptr detector_gpu_ptr;