diff --git a/CMakeLists.txt b/CMakeLists.txt index b4fde9c..21841be 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -34,6 +34,53 @@ if(ENABLE_QT) AUTORCC ON) endif() -target_sources(${CMAKE_PROJECT_NAME} PRIVATE src/plugin-main.c) +set(USE_SYSTEM_ONNXRUNTIME + OFF + CACHE STRING "Use system ONNX Runtime") + +set(DISABLE_ONNXRUNTIME_GPU + OFF + CACHE STRING "Disables GPU support of ONNX Runtime (Only valid on Linux)") + +if(DISABLE_ONNXRUNTIME_GPU) + target_compile_definitions(${CMAKE_PROJECT_NAME} PRIVATE DISABLE_ONNXRUNTIME_GPU) +endif() + +if(USE_SYSTEM_ONNXRUNTIME) + if(OS_LINUX) + find_package(Onnxruntime 1.16.3 REQUIRED) + set(Onnxruntime_INCLUDE_PATH + ${Onnxruntime_INCLUDE_DIR} ${Onnxruntime_INCLUDE_DIR}/onnxruntime + ${Onnxruntime_INCLUDE_DIR}/onnxruntime/core/session ${Onnxruntime_INCLUDE_DIR}/onnxruntime/core/providers/cpu) + target_link_libraries(${CMAKE_PROJECT_NAME} PRIVATE "${Onnxruntime_LIBRARIES}") + target_include_directories(${CMAKE_PROJECT_NAME} SYSTEM PUBLIC "${Onnxruntime_INCLUDE_PATH}") + else() + message(FATAL_ERROR "System ONNX Runtime is only supported on Linux!") + endif() +else() + include(cmake/FetchOnnxruntime.cmake) +endif() + +set(USE_SYSTEM_OPENCV + OFF + CACHE STRING "Use system OpenCV") +if(USE_SYSTEM_OPENCV) + if(OS_LINUX) + find_package(OpenCV REQUIRED COMPONENTS core imgproc) + target_link_libraries(${CMAKE_PROJECT_NAME} PRIVATE "${OpenCV_LIBRARIES}") + target_include_directories(${CMAKE_PROJECT_NAME} SYSTEM PUBLIC "${OpenCV_INCLUDE_DIRS}") + else() + message(FATAL_ERROR "System OpenCV is only supported on Linux!") + endif() +else() + include(cmake/FetchOpenCV.cmake) + target_link_libraries(${CMAKE_PROJECT_NAME} PRIVATE OpenCV) +endif() + +target_sources(${CMAKE_PROJECT_NAME} PRIVATE src/plugin-main.c + src/detect-filter.cpp + src/detect-filter-info.c + src/obs-utils/obs-utils.cpp + src/edgeyolo/edgeyolo_onnxruntime.cpp) set_target_properties_plugin(${CMAKE_PROJECT_NAME} PROPERTIES OUTPUT_NAME ${_name}) diff --git a/buildspec.json b/buildspec.json index 6a29f1b..2f95816 100644 --- a/buildspec.json +++ b/buildspec.json @@ -33,14 +33,14 @@ }, "platformConfig": { "macos": { - "bundleId": "com.example.obs-plugintemplate" + "bundleId": "com.royshilkrot.obs-detect" } }, - "name": "obs-plugintemplate", - "displayName": "OBS Plugin Template", - "version": "1.0.0", - "author": "Your Name Here", - "website": "https://example.com", + "name": "obs-detect", + "displayName": "OBS Object Detection plugin", + "version": "0.0.1", + "author": "Roy Shilkrot", + "website": "https://github.com/occ-ai", "email": "me@example.com", "uuids": { "windowsApp": "00000000-0000-0000-0000-000000000000" diff --git a/cmake/FetchOnnxruntime.cmake b/cmake/FetchOnnxruntime.cmake new file mode 100644 index 0000000..5643d58 --- /dev/null +++ b/cmake/FetchOnnxruntime.cmake @@ -0,0 +1,113 @@ +include(FetchContent) + +set(CUSTOM_ONNXRUNTIME_URL + "" + CACHE STRING "URL of a downloaded ONNX Runtime tarball") + +set(CUSTOM_ONNXRUNTIME_HASH + "" + CACHE STRING "Hash of a downloaded ONNX Runtime tarball") + +set(Onnxruntime_VERSION "1.17.1") + +if(CUSTOM_ONNXRUNTIME_URL STREQUAL "") + set(USE_PREDEFINED_ONNXRUNTIME ON) +else() + if(CUSTOM_ONNXRUNTIME_HASH STREQUAL "") + message(FATAL_ERROR "Both of CUSTOM_ONNXRUNTIME_URL and CUSTOM_ONNXRUNTIME_HASH must be present!") + else() + set(USE_PREDEFINED_ONNXRUNTIME OFF) + endif() +endif() + +if(USE_PREDEFINED_ONNXRUNTIME) + set(Onnxruntime_BASEURL "https://github.com/microsoft/onnxruntime/releases/download/v${Onnxruntime_VERSION}") + set(Onnxruntime_WINDOWS_VERSION "v${Onnxruntime_VERSION}-1") + set(Onnxruntime_WINDOWS_BASEURL + "https://github.com/occ-ai/occ-ai-dep-onnxruntime-static-win/releases/download/${Onnxruntime_WINDOWS_VERSION}") + + if(APPLE) + set(Onnxruntime_URL "${Onnxruntime_BASEURL}/onnxruntime-osx-universal2-${Onnxruntime_VERSION}.tgz") + set(Onnxruntime_HASH SHA256=9FA57FA6F202A373599377EF75064AE568FDA8DA838632B26A86024C7378D306) + elseif(MSVC) + set(Onnxruntime_URL "${Onnxruntime_WINDOWS_BASEURL}/onnxruntime-windows-${Onnxruntime_WINDOWS_VERSION}-Release.zip") + set(OOnnxruntime_HASH SHA256=39E63850D9762810161AE1B4DEAE5E3C02363521273E4B894A9D9707AB626C38) + else() + if(CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64") + set(Onnxruntime_URL "${Onnxruntime_BASEURL}/onnxruntime-linux-aarch64-${Onnxruntime_VERSION}.tgz") + set(Onnxruntime_HASH SHA256=70B6F536BB7AB5961D128E9DBD192368AC1513BFFB74FE92F97AAC342FBD0AC1) + else() + set(Onnxruntime_URL "${Onnxruntime_BASEURL}/onnxruntime-linux-x64-gpu-${Onnxruntime_VERSION}.tgz") + set(Onnxruntime_HASH SHA256=613C53745EA4960ED368F6B3AB673558BB8561C84A8FA781B4EA7FB4A4340BE4) + endif() + endif() +else() + set(Onnxruntime_URL "${CUSTOM_ONNXRUNTIME_URL}") + set(Onnxruntime_HASH "${CUSTOM_ONNXRUNTIME_HASH}") +endif() + +FetchContent_Declare( + onnxruntime + URL ${Onnxruntime_URL} + URL_HASH ${Onnxruntime_HASH}) +FetchContent_MakeAvailable(onnxruntime) + +if(APPLE) + set(Onnxruntime_LIB "${onnxruntime_SOURCE_DIR}/lib/libonnxruntime.${Onnxruntime_VERSION}.dylib") + target_link_libraries(${CMAKE_PROJECT_NAME} PRIVATE "${Onnxruntime_LIB}") + target_include_directories(${CMAKE_PROJECT_NAME} SYSTEM PUBLIC "${onnxruntime_SOURCE_DIR}/include") + target_sources(${CMAKE_PROJECT_NAME} PRIVATE "${Onnxruntime_LIB}") + set_property(SOURCE "${Onnxruntime_LIB}" PROPERTY MACOSX_PACKAGE_LOCATION Frameworks) + source_group("Frameworks" FILES "${Onnxruntime_LIB}") + add_custom_command( + TARGET "${CMAKE_PROJECT_NAME}" + POST_BUILD + COMMAND + ${CMAKE_INSTALL_NAME_TOOL} -change "@rpath/libonnxruntime.${Onnxruntime_VERSION}.dylib" + "@loader_path/../Frameworks/libonnxruntime.${Onnxruntime_VERSION}.dylib" $) +elseif(MSVC) + add_library(Ort INTERFACE) + set(Onnxruntime_LIB_NAMES + session;providers_shared;providers_dml;optimizer;providers;framework;graph;util;mlas;common;flatbuffers) + foreach(lib_name IN LISTS Onnxruntime_LIB_NAMES) + add_library(Ort::${lib_name} STATIC IMPORTED) + set_target_properties(Ort::${lib_name} PROPERTIES IMPORTED_LOCATION + ${onnxruntime_SOURCE_DIR}/lib/onnxruntime_${lib_name}.lib) + set_target_properties(Ort::${lib_name} PROPERTIES INTERFACE_INCLUDE_DIRECTORIES ${onnxruntime_SOURCE_DIR}/include) + target_link_libraries(Ort INTERFACE Ort::${lib_name}) + endforeach() + + set(Onnxruntime_EXTERNAL_LIB_NAMES + onnx;onnx_proto;libprotobuf-lite;re2;absl_throw_delegate;absl_hash;absl_city;absl_low_level_hash;absl_raw_hash_set + ) + foreach(lib_name IN LISTS Onnxruntime_EXTERNAL_LIB_NAMES) + add_library(Ort::${lib_name} STATIC IMPORTED) + set_target_properties(Ort::${lib_name} PROPERTIES IMPORTED_LOCATION ${onnxruntime_SOURCE_DIR}/lib/${lib_name}.lib) + target_link_libraries(Ort INTERFACE Ort::${lib_name}) + endforeach() + + add_library(Ort::DirectML SHARED IMPORTED) + set_target_properties(Ort::DirectML PROPERTIES IMPORTED_LOCATION ${onnxruntime_SOURCE_DIR}/bin/DirectML.dll) + set_target_properties(Ort::DirectML PROPERTIES IMPORTED_IMPLIB ${onnxruntime_SOURCE_DIR}/bin/DirectML.lib) + + target_link_libraries(Ort INTERFACE Ort::DirectML d3d12.lib dxgi.lib dxguid.lib Dxcore.lib) + + target_link_libraries(${CMAKE_PROJECT_NAME} PRIVATE Ort) + + install(IMPORTED_RUNTIME_ARTIFACTS Ort::DirectML DESTINATION "obs-plugins/64bit") +else() + if(CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64") + set(Onnxruntime_LINK_LIBS "${onnxruntime_SOURCE_DIR}/lib/libonnxruntime.so.${Onnxruntime_VERSION}") + set(Onnxruntime_INSTALL_LIBS ${Onnxruntime_LINK_LIBS}) + else() + set(Onnxruntime_LINK_LIBS "${onnxruntime_SOURCE_DIR}/lib/libonnxruntime.so.${Onnxruntime_VERSION}") + set(Onnxruntime_INSTALL_LIBS + ${Onnxruntime_LINK_LIBS} "${onnxruntime_SOURCE_DIR}/lib/libonnxruntime_providers_shared.so" + "${onnxruntime_SOURCE_DIR}/lib/libonnxruntime_providers_cuda.so" + "${onnxruntime_SOURCE_DIR}/lib/libonnxruntime_providers_tensorrt.so") + endif() + target_link_libraries(${CMAKE_PROJECT_NAME} PRIVATE ${Onnxruntime_LINK_LIBS}) + target_include_directories(${CMAKE_PROJECT_NAME} SYSTEM PUBLIC "${onnxruntime_SOURCE_DIR}/include") + install(FILES ${Onnxruntime_INSTALL_LIBS} DESTINATION "${CMAKE_INSTALL_LIBDIR}/obs-plugins/${CMAKE_PROJECT_NAME}") + set_target_properties(${CMAKE_PROJECT_NAME} PROPERTIES INSTALL_RPATH "$ORIGIN/${CMAKE_PROJECT_NAME}") +endif() diff --git a/cmake/FetchOpenCV.cmake b/cmake/FetchOpenCV.cmake new file mode 100644 index 0000000..cccdc25 --- /dev/null +++ b/cmake/FetchOpenCV.cmake @@ -0,0 +1,80 @@ +include(FetchContent) + +set(CUSTOM_OPENCV_URL + "" + CACHE STRING "URL of a downloaded OpenCV static library tarball") + +set(CUSTOM_OPENCV_HASH + "" + CACHE STRING "Hash of a downloaded OpenCV staitc library tarball") + +if(CUSTOM_OPENCV_URL STREQUAL "") + set(USE_PREDEFINED_OPENCV ON) +else() + if(CUSTOM_OPENCV_HASH STREQUAL "") + message(FATAL_ERROR "Both of CUSTOM_OPENCV_URL and CUSTOM_OPENCV_HASH must be present!") + else() + set(USE_PREDEFINED_OPENCV OFF) + endif() +endif() + +if(USE_PREDEFINED_OPENCV) + set(OpenCV_VERSION "v4.8.1-1") + set(OpenCV_BASEURL "https://github.com/obs-ai/obs-backgroundremoval-dep-opencv/releases/download/${OpenCV_VERSION}") + + if(${CMAKE_BUILD_TYPE} STREQUAL Release OR ${CMAKE_BUILD_TYPE} STREQUAL RelWithDebInfo) + set(OpenCV_BUILD_TYPE Release) + else() + set(OpenCV_BUILD_TYPE Debug) + endif() + + if(APPLE) + if(OpenCV_BUILD_TYPE STREQUAL Debug) + set(OpenCV_URL "${OpenCV_BASEURL}/opencv-macos-${OpenCV_VERSION}-Debug.tar.gz") + set(OpenCV_HASH SHA256=2930e335a19cc03a3d825e2b76eadd0d5cf08d8baf6537747d43f503dff32454) + else() + set(OpenCV_URL "${OpenCV_BASEURL}/opencv-macos-${OpenCV_VERSION}-Release.tar.gz") + set(OpenCV_HASH SHA256=b0c4fe2370b0bd5aa65c408e875b1ab18508ba31b93083805d7e398a3ecafdac) + endif() + elseif(MSVC) + if(OpenCV_BUILD_TYPE STREQUAL Debug) + set(OpenCV_URL "${OpenCV_BASEURL}/opencv-windows-${OpenCV_VERSION}-Debug.zip") + set(OpenCV_HASH SHA256=0c5ef12cf4b4e4db7ea17a24db156165b6f01759f3f1660b069d0722e5d5dc37) + else() + set(OpenCV_URL "${OpenCV_BASEURL}/opencv-windows-${OpenCV_VERSION}-Release.zip") + set(OpenCV_HASH SHA256=5e468f71d41d3a3ea46cc4f247475877f65d3655a2764a2c01074bda3b3e6864) + endif() + else() + if(OpenCV_BUILD_TYPE STREQUAL Debug) + set(OpenCV_URL "${OpenCV_BASEURL}/opencv-linux-${OpenCV_VERSION}-Debug.tar.gz") + set(OpenCV_HASH SHA256=e2e246d6b4f279be80e7fd0a78cba8a0eeee7b53ae807f2f57428d6876306422) + else() + set(OpenCV_URL "${OpenCV_BASEURL}/opencv-linux-${OpenCV_VERSION}-Release.tar.gz") + set(OpenCV_HASH SHA256=809922a7cc9f344a2d82a232ed7b02e122c82e77cba94b4047e666a0527cc00e) + endif() + endif() +else() + set(OpenCV_URL "${CUSTOM_OPENCV_URL}") + set(OpenCV_HASH "${CUSTOM_OPENCV_HASH}") +endif() + +FetchContent_Declare( + opencv + URL ${OpenCV_URL} + URL_HASH ${OpenCV_HASH}) +FetchContent_MakeAvailable(opencv) + +add_library(OpenCV INTERFACE) +if(MSVC) + target_link_libraries( + OpenCV + INTERFACE ${opencv_SOURCE_DIR}/x64/vc17/staticlib/opencv_imgproc481.lib + ${opencv_SOURCE_DIR}/x64/vc17/staticlib/opencv_core481.lib + ${opencv_SOURCE_DIR}/x64/vc17/staticlib/zlib.lib) + target_include_directories(OpenCV SYSTEM INTERFACE ${opencv_SOURCE_DIR}/include) +else() + target_link_libraries( + OpenCV INTERFACE ${opencv_SOURCE_DIR}/lib/libopencv_imgproc.a ${opencv_SOURCE_DIR}/lib/libopencv_core.a + ${opencv_SOURCE_DIR}/lib/opencv4/3rdparty/libzlib.a) + target_include_directories(OpenCV SYSTEM INTERFACE ${opencv_SOURCE_DIR}/include/opencv4) +endif() diff --git a/cmake/windows/compilerconfig.cmake b/cmake/windows/compilerconfig.cmake index 55984f4..37a108e 100644 --- a/cmake/windows/compilerconfig.cmake +++ b/cmake/windows/compilerconfig.cmake @@ -35,7 +35,8 @@ add_compile_options( "$<$:/MP>" "$<$:${_obs_clang_c_options}>" "$<$:${_obs_clang_cxx_options}>" - $<$>:/Gy>) + $<$>:/Gy> + /IGNORE:4099) add_compile_definitions(UNICODE _UNICODE _CRT_SECURE_NO_WARNINGS _CRT_NONSTDC_NO_WARNINGS $<$:DEBUG> $<$:_DEBUG>) @@ -45,7 +46,8 @@ add_link_options($<$>:/OPT:REF> $<$>:/OPT:ICF> $<$>:/INCREMENTAL:NO> /DEBUG - /Brepro) + /Brepro + /IGNORE:4099) # cmake-format: on if(CMAKE_COMPILE_WARNING_AS_ERROR) diff --git a/data/models/edgeyolo_tiny_lrelu_coco_256x416.onnx b/data/models/edgeyolo_tiny_lrelu_coco_256x416.onnx new file mode 100644 index 0000000..cebc452 Binary files /dev/null and b/data/models/edgeyolo_tiny_lrelu_coco_256x416.onnx differ diff --git a/src/FilterData.h b/src/FilterData.h new file mode 100644 index 0000000..5b0dab4 --- /dev/null +++ b/src/FilterData.h @@ -0,0 +1,41 @@ +#ifndef FILTERDATA_H +#define FILTERDATA_H + +#include +#include "edgeyolo/edgeyolo_onnxruntime.hpp" + +/** + * @brief The filter_data struct + * + * This struct is used to store the base data needed for ORT filters. + * +*/ +struct filter_data { + std::string useGPU; + uint32_t numThreads; + + obs_source_t *source; + gs_texrender_t *texrender; + gs_stagesurf_t *stagesurface; + gs_effect_t *effect; + + cv::Mat inputBGRA; + cv::Mat outputPreviewBGRA; + + bool isDisabled; + bool preview; + + std::mutex inputBGRALock; + std::mutex outputLock; + std::mutex modelMutex; + + std::unique_ptr edgeyolo; + +#if _WIN32 + std::wstring modelFilepath; +#else + std::string modelFilepath; +#endif +}; + +#endif /* FILTERDATA_H */ diff --git a/src/consts.h b/src/consts.h new file mode 100644 index 0000000..bc3b0e6 --- /dev/null +++ b/src/consts.h @@ -0,0 +1,21 @@ +#ifndef CONSTS_H +#define CONSTS_H + +const char *const USEGPU_CPU = "cpu"; +const char *const USEGPU_DML = "dml"; +const char *const USEGPU_CUDA = "cuda"; +const char *const USEGPU_TENSORRT = "tensorrt"; +const char *const USEGPU_COREML = "coreml"; + +const char *const EFFECT_PATH = "effects/mask_alpha_filter.effect"; +const char *const KAWASE_BLUR_EFFECT_PATH = "effects/kawase_blur.effect"; +const char *const BLEND_EFFECT_PATH = "effects/blend_images.effect"; + +const char *const PLUGIN_INFO_TEMPLATE = + "Detect Plugin (%1) by " + "OCC AI ❤️ " + "Support & Follow"; +const char *const PLUGIN_INFO_TEMPLATE_UPDATE_AVAILABLE = + "
🚀 Update available! (%1)
"; + +#endif /* CONSTS_H */ diff --git a/src/detect-filter-info.c b/src/detect-filter-info.c new file mode 100644 index 0000000..c6080be --- /dev/null +++ b/src/detect-filter-info.c @@ -0,0 +1,17 @@ +#include "detect-filter.h" + +struct obs_source_info detect_filter_info = { + .id = "detect-filter", + .type = OBS_SOURCE_TYPE_FILTER, + .output_flags = OBS_SOURCE_VIDEO, + .get_name = detect_filter_getname, + .create = detect_filter_create, + .destroy = detect_filter_destroy, + .get_defaults = detect_filter_defaults, + .get_properties = detect_filter_properties, + .update = detect_filter_update, + .activate = detect_filter_activate, + .deactivate = detect_filter_deactivate, + .video_tick = detect_filter_video_tick, + .video_render = detect_filter_video_render, +}; diff --git a/src/detect-filter.cpp b/src/detect-filter.cpp new file mode 100644 index 0000000..f2fd1b3 --- /dev/null +++ b/src/detect-filter.cpp @@ -0,0 +1,389 @@ +#include "detect-filter.h" + +#include + +#ifdef _WIN32 +#include +#include +#endif // _WIN32 + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include "FilterData.h" +#include "consts.h" +#include "obs-utils/obs-utils.h" +#include "edgeyolo/utils.hpp" + +struct detect_filter : public filter_data {}; + +const char *detect_filter_getname(void *unused) +{ + UNUSED_PARAMETER(unused); + return obs_module_text("Detect"); +} + +/** PROPERTIES */ + +static bool visible_on_bool(obs_properties_t *ppts, obs_data_t *settings, + const char *bool_prop, const char *prop_name) +{ + const bool enabled = obs_data_get_bool(settings, bool_prop); + obs_property_t *p = obs_properties_get(ppts, prop_name); + obs_property_set_visible(p, enabled); + return true; +} + +static bool enable_advanced_settings(obs_properties_t *ppts, obs_property_t *p, + obs_data_t *settings) +{ + const bool enabled = obs_data_get_bool(settings, "advanced"); + + for (const char *prop_name : + {"threshold", "useGPU", "preview", "numThreads"}) { + p = obs_properties_get(ppts, prop_name); + obs_property_set_visible(p, enabled); + } + + return true; +} + +obs_properties_t *detect_filter_properties(void *data) +{ + obs_properties_t *props = obs_properties_create(); + + obs_property_t *advanced = obs_properties_add_bool( + props, "advanced", obs_module_text("Advanced")); + + // If advanced is selected show the advanced settings, otherwise hide them + obs_property_set_modified_callback(advanced, enable_advanced_settings); + + obs_properties_add_float_slider(props, "threshold", + obs_module_text("Threshold"), 0.0, 1.0, + 0.025); + + /* GPU, CPU and performance Props */ + obs_property_t *p_use_gpu = obs_properties_add_list( + props, "useGPU", obs_module_text("InferenceDevice"), + OBS_COMBO_TYPE_LIST, OBS_COMBO_FORMAT_STRING); + + obs_property_list_add_string(p_use_gpu, obs_module_text("CPU"), + USEGPU_CPU); +#if defined(__linux__) && defined(__x86_64__) + obs_property_list_add_string(p_use_gpu, obs_module_text("GPUTensorRT"), + USEGPU_TENSORRT); +#endif +#if _WIN32 + obs_property_list_add_string(p_use_gpu, obs_module_text("GPUDirectML"), + USEGPU_DML); +#endif +#if defined(__APPLE__) + obs_property_list_add_string(p_use_gpu, obs_module_text("CoreML"), + USEGPU_COREML); +#endif + + obs_properties_add_int_slider(props, "numThreads", + obs_module_text("NumThreads"), 0, 8, 1); + + obs_properties_add_bool(props, "preview", obs_module_text("Preview")); + + // Add a informative text about the plugin + // replace the placeholder with the current version + // use std::regex_replace instead of QString::arg because the latter doesn't work on Linux + std::string basic_info = std::regex_replace( + PLUGIN_INFO_TEMPLATE, std::regex("%1"), PLUGIN_VERSION); + // Check for update + // if (get_latest_version() != nullptr) { + // basic_info += std::regex_replace( + // PLUGIN_INFO_TEMPLATE_UPDATE_AVAILABLE, std::regex("%1"), + // get_latest_version()); + // } + obs_properties_add_text(props, "info", basic_info.c_str(), + OBS_TEXT_INFO); + + UNUSED_PARAMETER(data); + return props; +} + +void detect_filter_defaults(obs_data_t *settings) +{ + obs_data_set_default_bool(settings, "advanced", false); +#if _WIN32 + obs_data_set_default_string(settings, "useGPU", USEGPU_DML); +#elif defined(__APPLE__) + obs_data_set_default_string(settings, "useGPU", USEGPU_CPU); +#else + // Linux + obs_data_set_default_string(settings, "useGPU", USEGPU_CPU); +#endif + obs_data_set_default_int(settings, "numThreads", 1); + obs_data_set_default_bool(settings, "preview", false); +} + +void detect_filter_update(void *data, obs_data_t *settings) +{ + obs_log(LOG_INFO, "Detect filter updated"); + struct detect_filter *tf = reinterpret_cast(data); + + tf->isDisabled = true; + + tf->preview = obs_data_get_bool(settings, "preview"); + + const std::string newUseGpu = obs_data_get_string(settings, "useGPU"); + const uint32_t newNumThreads = + (uint32_t)obs_data_get_int(settings, "numThreads"); + + if (tf->useGPU != newUseGpu || tf->numThreads != newNumThreads) { + // lock modelMutex + std::unique_lock lock(tf->modelMutex); + + char *modelFilepath_rawPtr = obs_module_file( + "models/edgeyolo_tiny_lrelu_coco_256x416.onnx"); + + if (modelFilepath_rawPtr == nullptr) { + obs_log(LOG_ERROR, + "Unable to get model filename from plugin."); + return; + } + + std::string modelFilepath_s(modelFilepath_rawPtr); + +#if _WIN32 + int outLength = MultiByteToWideChar(CP_ACP, MB_PRECOMPOSED, + modelFilepath_rawPtr, -1, + nullptr, 0); + tf->modelFilepath = std::wstring(outLength, L'\0'); + MultiByteToWideChar(CP_ACP, MB_PRECOMPOSED, + modelFilepath_rawPtr, -1, + tf->modelFilepath.data(), outLength); +#else + tf->modelFilepath = std::string(modelFilepath_rawPtr); +#endif + + // Re-initialize model if it's not already the selected one or switching inference device + tf->useGPU = newUseGpu; + tf->numThreads = newNumThreads; + + // parameters + int onnxruntime_device_id_ = 0; + bool onnxruntime_use_parallel_ = false; + float nms_th_ = 0.45f; + float conf_th_ = 0.3f; + int num_classes_ = (int)edgeyolo_cpp::COCO_CLASSES.size(); + + // Load model + try { + tf->edgeyolo = std::make_unique< + edgeyolo_cpp::EdgeYOLOONNXRuntime>( + tf->modelFilepath, tf->numThreads, + tf->numThreads, tf->useGPU, onnxruntime_device_id_, + onnxruntime_use_parallel_, nms_th_, conf_th_, + num_classes_); + } catch (const std::exception &e) { + obs_log(LOG_ERROR, "Failed to load model: %s", + e.what()); + // disable filter + tf->isDisabled = true; + tf->edgeyolo.reset(); + return; + } + } + + // Log the currently selected options + obs_log(LOG_INFO, "Detect Filter Options:"); + // name of the source that the filter is attached to + obs_log(LOG_INFO, " Source: %s", obs_source_get_name(tf->source)); + obs_log(LOG_INFO, " Inference Device: %s", tf->useGPU.c_str()); + obs_log(LOG_INFO, " Num Threads: %d", tf->numThreads); + obs_log(LOG_INFO, " Disabled: %s", tf->isDisabled ? "true" : "false"); +#ifdef _WIN32 + obs_log(LOG_INFO, " Model file path: %S", tf->modelFilepath.c_str()); +#else + obs_log(LOG_INFO, " Model file path: %s", tf->modelFilepath.c_str()); +#endif + + // enable + tf->isDisabled = false; +} + +void detect_filter_activate(void *data) +{ + obs_log(LOG_INFO, "Detect filter activated"); + struct detect_filter *tf = reinterpret_cast(data); + tf->isDisabled = false; +} + +void detect_filter_deactivate(void *data) +{ + obs_log(LOG_INFO, "Detect filter deactivated"); + struct detect_filter *tf = reinterpret_cast(data); + tf->isDisabled = true; +} + +/** FILTER CORE */ + +void *detect_filter_create(obs_data_t *settings, obs_source_t *source) +{ + obs_log(LOG_INFO, "Detect filter created"); + void *data = bmalloc(sizeof(struct detect_filter)); + struct detect_filter *tf = new (data) detect_filter(); + + tf->source = source; + tf->texrender = gs_texrender_create(GS_BGRA, GS_ZS_NONE); + tf->effect = obs_get_base_effect(OBS_EFFECT_OPAQUE); + + detect_filter_update(tf, settings); + + return tf; +} + +void detect_filter_destroy(void *data) +{ + obs_log(LOG_INFO, "Detect filter destroyed"); + + struct detect_filter *tf = reinterpret_cast(data); + + if (tf) { + tf->isDisabled = true; + + obs_enter_graphics(); + gs_texrender_destroy(tf->texrender); + if (tf->stagesurface) { + gs_stagesurface_destroy(tf->stagesurface); + } + obs_leave_graphics(); + tf->~detect_filter(); + bfree(tf); + } +} + +void detect_filter_video_tick(void *data, float seconds) +{ + UNUSED_PARAMETER(seconds); + + struct detect_filter *tf = reinterpret_cast(data); + + if (tf->isDisabled) { + return; + } + + if (!obs_source_enabled(tf->source)) { + return; + } + + if (!tf->edgeyolo) { + obs_log(LOG_ERROR, "Model is not initialized"); + return; + } + + cv::Mat imageBGRA; + { + std::unique_lock lock(tf->inputBGRALock, + std::try_to_lock); + if (!lock.owns_lock()) { + // No data to process + return; + } + if (tf->inputBGRA.empty()) { + // No data to process + return; + } + imageBGRA = tf->inputBGRA.clone(); + } + + cv::Mat frame; + cv::cvtColor(imageBGRA, frame, cv::COLOR_BGRA2BGR); + std::vector objects; + + try { + std::unique_lock lock(tf->modelMutex); + objects = tf->edgeyolo->inference(frame); + } catch (const Ort::Exception &e) { + obs_log(LOG_ERROR, "ONNXRuntime Exception: %s", e.what()); + // TODO: Fall back to CPU if it makes sense + } catch (const std::exception &e) { + obs_log(LOG_ERROR, "%s", e.what()); + } + + if (objects.size() == 0) { + return; + } + + cv::Mat out_frame = frame.clone(); + edgeyolo_cpp::utils::draw_objects(out_frame, objects); + + std::lock_guard lock(tf->outputLock); + cv::cvtColor(out_frame, tf->outputPreviewBGRA, cv::COLOR_BGR2BGRA); +} + +void detect_filter_video_render(void *data, gs_effect_t *_effect) +{ + UNUSED_PARAMETER(_effect); + + struct detect_filter *tf = reinterpret_cast(data); + + if (tf->isDisabled) { + if (tf->source) { + obs_source_skip_video_filter(tf->source); + } + return; + } + + uint32_t width, height; + if (!getRGBAFromStageSurface(tf, width, height)) { + if (tf->source) { + obs_source_skip_video_filter(tf->source); + } + return; + } + + // if preview is enabled, render the image + if (tf->preview) { + gs_texture_t *tex = nullptr; + { + // lock the outputLock mutex + std::lock_guard lock(tf->outputLock); + if (tf->outputPreviewBGRA.empty()) { + obs_log(LOG_ERROR, "Preview image is empty"); + obs_source_skip_video_filter(tf->source); + return; + } + if ((uint32_t)tf->outputPreviewBGRA.cols != width || + (uint32_t)tf->outputPreviewBGRA.rows != height) { + obs_source_skip_video_filter(tf->source); + return; + } + + tex = gs_texture_create( + width, height, GS_BGRA, 1, + (const uint8_t **)&tf->outputPreviewBGRA.data, + 0); + } + + gs_eparam_t *imageParam = + gs_effect_get_param_by_name(tf->effect, "image"); + gs_effect_set_texture(imageParam, tex); + + gs_blend_state_push(); + gs_blend_function(GS_BLEND_ONE, GS_BLEND_ZERO); + + while (gs_effect_loop(tf->effect, "Draw")) { + gs_draw_sprite(tex, 0, 0, 0); + } + + gs_blend_state_pop(); + gs_texture_destroy(tex); + } else { + obs_source_skip_video_filter(tf->source); + } + return; +} diff --git a/src/detect-filter.h b/src/detect-filter.h new file mode 100644 index 0000000..64e10bc --- /dev/null +++ b/src/detect-filter.h @@ -0,0 +1,20 @@ +#include + +#ifdef __cplusplus +extern "C" { +#endif + +const char *detect_filter_getname(void *unused); +void *detect_filter_create(obs_data_t *settings, obs_source_t *source); +void detect_filter_destroy(void *data); +void detect_filter_defaults(obs_data_t *settings); +obs_properties_t *detect_filter_properties(void *data); +void detect_filter_update(void *data, obs_data_t *settings); +void detect_filter_activate(void *data); +void detect_filter_deactivate(void *data); +void detect_filter_video_tick(void *data, float seconds); +void detect_filter_video_render(void *data, gs_effect_t *_effect); + +#ifdef __cplusplus +} +#endif diff --git a/src/edgeyolo/coco_names.hpp b/src/edgeyolo/coco_names.hpp new file mode 100644 index 0000000..3e4c9b2 --- /dev/null +++ b/src/edgeyolo/coco_names.hpp @@ -0,0 +1,99 @@ +#ifndef _EdgeYOLO_CPP_COCO_NAMES_HPP +#define _EdgeYOLO_CPP_COCO_NAMES_HPP + +namespace edgeyolo_cpp{ + static const std::vector COCO_CLASSES = { + "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", + "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", + "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", + "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", + "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", + "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", + "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", + "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", + "hair drier", "toothbrush" + }; + const float color_list[80][3] ={ + {0.000f, 0.447f, 0.741f}, + {0.850f, 0.325f, 0.098f}, + {0.929f, 0.694f, 0.125f}, + {0.494f, 0.184f, 0.556f}, + {0.466f, 0.674f, 0.188f}, + {0.301f, 0.745f, 0.933f}, + {0.635f, 0.078f, 0.184f}, + {0.300f, 0.300f, 0.300f}, + {0.600f, 0.600f, 0.600f}, + {1.000f, 0.000f, 0.000f}, + {1.000f, 0.500f, 0.000f}, + {0.749f, 0.749f, 0.000f}, + {0.000f, 1.000f, 0.000f}, + {0.000f, 0.000f, 1.000f}, + {0.667f, 0.000f, 1.000f}, + {0.333f, 0.333f, 0.000f}, + {0.333f, 0.667f, 0.000f}, + {0.333f, 1.000f, 0.000f}, + {0.667f, 0.333f, 0.000f}, + {0.667f, 0.667f, 0.000f}, + {0.667f, 1.000f, 0.000f}, + {1.000f, 0.333f, 0.000f}, + {1.000f, 0.667f, 0.000f}, + {1.000f, 1.000f, 0.000f}, + {0.000f, 0.333f, 0.500f}, + {0.000f, 0.667f, 0.500f}, + {0.000f, 1.000f, 0.500f}, + {0.333f, 0.000f, 0.500f}, + {0.333f, 0.333f, 0.500f}, + {0.333f, 0.667f, 0.500f}, + {0.333f, 1.000f, 0.500f}, + {0.667f, 0.000f, 0.500f}, + {0.667f, 0.333f, 0.500f}, + {0.667f, 0.667f, 0.500f}, + {0.667f, 1.000f, 0.500f}, + {1.000f, 0.000f, 0.500f}, + {1.000f, 0.333f, 0.500f}, + {1.000f, 0.667f, 0.500f}, + {1.000f, 1.000f, 0.500f}, + {0.000f, 0.333f, 1.000f}, + {0.000f, 0.667f, 1.000f}, + {0.000f, 1.000f, 1.000f}, + {0.333f, 0.000f, 1.000f}, + {0.333f, 0.333f, 1.000f}, + {0.333f, 0.667f, 1.000f}, + {0.333f, 1.000f, 1.000f}, + {0.667f, 0.000f, 1.000f}, + {0.667f, 0.333f, 1.000f}, + {0.667f, 0.667f, 1.000f}, + {0.667f, 1.000f, 1.000f}, + {1.000f, 0.000f, 1.000f}, + {1.000f, 0.333f, 1.000f}, + {1.000f, 0.667f, 1.000f}, + {0.333f, 0.000f, 0.000f}, + {0.500f, 0.000f, 0.000f}, + {0.667f, 0.000f, 0.000f}, + {0.833f, 0.000f, 0.000f}, + {1.000f, 0.000f, 0.000f}, + {0.000f, 0.167f, 0.000f}, + {0.000f, 0.333f, 0.000f}, + {0.000f, 0.500f, 0.000f}, + {0.000f, 0.667f, 0.000f}, + {0.000f, 0.833f, 0.000f}, + {0.000f, 1.000f, 0.000f}, + {0.000f, 0.000f, 0.167f}, + {0.000f, 0.000f, 0.333f}, + {0.000f, 0.000f, 0.500f}, + {0.000f, 0.000f, 0.667f}, + {0.000f, 0.000f, 0.833f}, + {0.000f, 0.000f, 1.000f}, + {0.000f, 0.000f, 0.000f}, + {0.143f, 0.143f, 0.143f}, + {0.286f, 0.286f, 0.286f}, + {0.429f, 0.429f, 0.429f}, + {0.571f, 0.571f, 0.571f}, + {0.714f, 0.714f, 0.714f}, + {0.857f, 0.857f, 0.857f}, + {0.000f, 0.447f, 0.741f}, + {0.314f, 0.717f, 0.741f}, + {0.50f, 0.5f, 0.0f} + }; +} +#endif diff --git a/src/edgeyolo/core.hpp b/src/edgeyolo/core.hpp new file mode 100644 index 0000000..ffe5620 --- /dev/null +++ b/src/edgeyolo/core.hpp @@ -0,0 +1,262 @@ +#ifndef _EdgeYOLO_CPP_CORE_HPP +#define _EdgeYOLO_CPP_CORE_HPP + +#include +#include + +namespace edgeyolo_cpp +{ +/** + * @brief Define names based depends on Unicode path support + */ +#define tcout std::cout + +#ifdef _WIN32 +#define file_name_t std::wstring +#else +#define file_name_t std::string +#endif + +#define imread_t cv::imread + + struct Object + { + cv::Rect_ rect; + int label; + float prob; + }; + + struct GridAndStride + { + int grid0; + int grid1; + int stride; + }; + + class AbcEdgeYOLO + { + public: + AbcEdgeYOLO() {} + AbcEdgeYOLO(float nms_th = 0.45, float conf_th = 0.3, + int num_classes = 80) + : nms_thresh_(nms_th), bbox_conf_thresh_(conf_th), + num_classes_(num_classes) + { + } + virtual std::vector inference(const cv::Mat &frame) = 0; + + protected: + int input_w_; + int input_h_; + float nms_thresh_; + float bbox_conf_thresh_; + int num_classes_; + int num_array_; + const std::vector mean_ = {0.485f, 0.456f, 0.406f}; + const std::vector std_ = {0.229f, 0.224f, 0.225f}; + + cv::Mat static_resize(const cv::Mat &img) + { + float r = std::fminf(input_w_ / (img.cols * 1.0f), input_h_ / (img.rows * 1.0f)); + // r = std::min(r, 1.0f); + int unpad_w = (int)(r * img.cols); + int unpad_h = (int)(r * img.rows); + cv::Mat re(unpad_h, unpad_w, CV_8UC3); + cv::resize(img, re, re.size()); + cv::Mat out(input_h_, input_w_, CV_8UC3, cv::Scalar(114, 114, 114)); + re.copyTo(out(cv::Rect(0, 0, re.cols, re.rows))); + return out; + } + + // for NCHW + void blobFromImage(const cv::Mat &img, float *blob_data) + { + size_t channels = 3; + size_t img_h = img.rows; + size_t img_w = img.cols; + for (size_t c = 0; c < channels; ++c) + { + for (size_t h = 0; h < img_h; ++h) + { + for (size_t w = 0; w < img_w; ++w) + { + blob_data[(int)(c * img_w * img_h + h * img_w + w)] = (float)img.ptr((int)h)[(int)w][(int)c]; + } + } + } + } + + // for NHWC + void blobFromImage_nhwc(const cv::Mat &img, float *blob_data) + { + size_t channels = 3; + size_t img_h = img.rows; + size_t img_w = img.cols; + for (size_t i = 0; i < img_h * img_w; ++i) + { + for (size_t c = 0; c < channels; ++c) + { + blob_data[i * channels + c] = (float)img.data[i * channels + c]; + } + } + } + + void generate_edgeyolo_proposals(const int num_array, const float *feat_ptr, const float prob_threshold, std::vector &objects) + { + + for (int idx = 0; idx < num_array; ++idx) + { + const int basic_pos = idx * (num_classes_ + 5); + + float box_objectness = feat_ptr[basic_pos + 4]; + int class_id = 0; + float max_class_score = 0.0; + for (int class_idx = 0; class_idx < num_classes_; ++class_idx) + { + float box_cls_score = feat_ptr[basic_pos + 5 + class_idx]; + float box_prob = box_objectness * box_cls_score; + if (box_prob > max_class_score) + { + class_id = class_idx; + max_class_score = box_prob; + } + } + if (max_class_score > prob_threshold) + { + float x_center = feat_ptr[basic_pos + 0]; + float y_center = feat_ptr[basic_pos + 1]; + float w = feat_ptr[basic_pos + 2]; + float h = feat_ptr[basic_pos + 3]; + float x0 = x_center - w * 0.5f; + float y0 = y_center - h * 0.5f; + + Object obj; + obj.rect.x = x0; + obj.rect.y = y0; + obj.rect.width = w; + obj.rect.height = h; + obj.label = class_id; + obj.prob = max_class_score; + objects.push_back(obj); + } + } + } + + float intersection_area(const Object &a, const Object &b) + { + cv::Rect_ inter = a.rect & b.rect; + return inter.area(); + } + + void qsort_descent_inplace(std::vector &faceobjects, int left, int right) + { + int i = left; + int j = right; + float p = faceobjects[(left + right) / 2].prob; + + while (i <= j) + { + while (faceobjects[i].prob > p) + ++i; + + while (faceobjects[j].prob < p) + --j; + + if (i <= j) + { + std::swap(faceobjects[i], faceobjects[j]); + + ++i; + --j; + } + } + if (left < j) + qsort_descent_inplace(faceobjects, left, j); + if (i < right) + qsort_descent_inplace(faceobjects, i, right); + } + + void qsort_descent_inplace(std::vector &objects) + { + if (objects.empty()) + return; + + qsort_descent_inplace(objects, 0, (int)(objects.size() - 1)); + } + + void nms_sorted_bboxes(const std::vector &faceobjects, std::vector &picked, const float nms_threshold) + { + picked.clear(); + + const size_t n = faceobjects.size(); + + std::vector areas(n); + for (size_t i = 0; i < n; ++i) + { + areas[i] = faceobjects[i].rect.area(); + } + + for (size_t i = 0; i < n; ++i) + { + const Object &a = faceobjects[i]; + const size_t picked_size = picked.size(); + + int keep = 1; + for (size_t j = 0; j < picked_size; ++j) + { + const Object &b = faceobjects[picked[j]]; + + // intersection over union + float inter_area = intersection_area(a, b); + float union_area = areas[i] + areas[picked[j]] - inter_area; + // float IoU = inter_area / union_area + if (inter_area / union_area > nms_threshold) + keep = 0; + } + + if (keep) + picked.push_back((int)i); + } + } + + void decode_outputs(const float *prob, const int num_array, + std::vector &objects, const float bbox_conf_thresh, + const float scale, const int img_w, const int img_h) + { + + std::vector proposals; + generate_edgeyolo_proposals(num_array, prob, bbox_conf_thresh, proposals); + + qsort_descent_inplace(proposals); + + std::vector picked; + nms_sorted_bboxes(proposals, picked, nms_thresh_); + + int count = (int)(picked.size()); + objects.clear(); + + for (int i = 0; i < count; ++i) + { + // adjust offset to original unpadded + float x0 = (proposals[picked[i]].rect.x) / scale; + float y0 = (proposals[picked[i]].rect.y) / scale; + float x1 = (proposals[picked[i]].rect.x + proposals[picked[i]].rect.width) / scale; + float y1 = (proposals[picked[i]].rect.y + proposals[picked[i]].rect.height) / scale; + + // clip + x0 = std::max(std::min(x0, (float)(img_w - 1)), 0.f); + y0 = std::max(std::min(y0, (float)(img_h - 1)), 0.f); + x1 = std::max(std::min(x1, (float)(img_w - 1)), 0.f); + y1 = std::max(std::min(y1, (float)(img_h - 1)), 0.f); + + proposals[picked[i]].rect.x = x0; + proposals[picked[i]].rect.y = y0; + proposals[picked[i]].rect.width = x1 - x0; + proposals[picked[i]].rect.height = y1 - y0; + + objects.push_back(proposals[picked[i]]); + } + } + }; +} +#endif diff --git a/src/edgeyolo/edgeyolo_onnxruntime.cpp b/src/edgeyolo/edgeyolo_onnxruntime.cpp new file mode 100644 index 0000000..bd00f34 --- /dev/null +++ b/src/edgeyolo/edgeyolo_onnxruntime.cpp @@ -0,0 +1,162 @@ +#include "edgeyolo_onnxruntime.hpp" + +#ifdef _WIN32 +#include +#endif + +namespace edgeyolo_cpp { + +EdgeYOLOONNXRuntime::EdgeYOLOONNXRuntime(file_name_t path_to_model, + int intra_op_num_threads, + int inter_op_num_threads, + const std::string& use_gpu, int device_id, + bool use_parallel, float nms_th, + float conf_th, int num_classes) + : AbcEdgeYOLO(nms_th, conf_th, num_classes), + intra_op_num_threads_(intra_op_num_threads), + inter_op_num_threads_(inter_op_num_threads), + use_gpu(use_gpu), + device_id_(device_id), + use_parallel_(use_parallel) +{ + try { + Ort::SessionOptions session_options; + + session_options.SetGraphOptimizationLevel( + GraphOptimizationLevel::ORT_ENABLE_ALL); + if (this->use_parallel_) { + session_options.SetExecutionMode( + ExecutionMode::ORT_PARALLEL); + session_options.SetInterOpNumThreads( + this->inter_op_num_threads_); + } else { + session_options.SetExecutionMode( + ExecutionMode::ORT_SEQUENTIAL); + } + session_options.SetIntraOpNumThreads( + this->intra_op_num_threads_); + + if (this->use_gpu == "cuda") { + OrtCUDAProviderOptions cuda_option; + cuda_option.device_id = this->device_id_; + session_options.AppendExecutionProvider_CUDA( + cuda_option); + } +#ifdef _WIN32 + if (this->use_gpu == "dml") { + auto &api = Ort::GetApi(); + OrtDmlApi *dmlApi = nullptr; + Ort::ThrowOnError(api.GetExecutionProviderApi( + "DML", ORT_API_VERSION, + (const void **)&dmlApi)); + Ort::ThrowOnError( + dmlApi->SessionOptionsAppendExecutionProvider_DML( + session_options, 0)); + } +#endif + + this->session_ = Ort::Session::Session( + this->env_, path_to_model.c_str(), session_options); + } catch (std::exception &e) { + std::cerr << e.what() << std::endl; + throw e; + } + + Ort::AllocatorWithDefaultOptions ort_alloc; + + // Allocate input memory buffer + std::cout << "input:" << std::endl; + this->input_name_ = std::string( + this->session_.GetInputNameAllocated(0, ort_alloc).get()); + // this->input_name_ = this->session_.GetInputName(0, ort_alloc); + std::cout << " name: " << this->input_name_ << std::endl; + auto input_info = this->session_.GetInputTypeInfo(0); + auto input_shape_info = input_info.GetTensorTypeAndShapeInfo(); + std::vector input_shape = input_shape_info.GetShape(); + ONNXTensorElementDataType input_tensor_type = + input_shape_info.GetElementType(); + this->input_h_ = (int)(input_shape[2]); + this->input_w_ = (int)(input_shape[3]); + + std::cout << " shape:" << std::endl; + for (size_t i = 0; i < input_shape.size(); i++) { + std::cout << " - " << input_shape[i] << std::endl; + } + std::cout << " tensor_type: " << input_tensor_type << std::endl; + + size_t input_byte_count = + sizeof(float) * input_shape_info.GetElementCount(); + std::unique_ptr input_buffer = + std::make_unique(input_byte_count); + // auto input_memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + auto input_memory_info = + Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); + + this->input_tensor_ = Ort::Value::CreateTensor( + input_memory_info, input_buffer.get(), input_byte_count, + input_shape.data(), input_shape.size(), input_tensor_type); + this->input_buffer_.emplace_back(std::move(input_buffer)); + + // Allocate output memory buffer + std::cout << "outputs" << std::endl; + this->output_name_ = std::string( + this->session_.GetOutputNameAllocated(0, ort_alloc).get()); + // this->output_name_ = this->session_.GetOutputName(0, ort_alloc); + std::cout << " name: " << this->output_name_ << std::endl; + + auto output_info = this->session_.GetOutputTypeInfo(0); + auto output_shape_info = output_info.GetTensorTypeAndShapeInfo(); + auto output_shape = output_shape_info.GetShape(); + auto output_tensor_type = output_shape_info.GetElementType(); + + this->num_array_ = 1; + std::cout << " shape:" << std::endl; + for (size_t i = 0; i < output_shape.size(); i++) { + std::cout << " - " << output_shape[i] << std::endl; + this->num_array_ *= (int)(output_shape[i]); + } + std::cout << " tensor_type: " << output_tensor_type << std::endl; + this->num_array_ /= (5 + this->num_classes_); + + size_t output_byte_count = + sizeof(float) * output_shape_info.GetElementCount(); + std::unique_ptr output_buffer = + std::make_unique(output_byte_count); + // auto output_memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault); + auto output_memory_info = + Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); + + this->output_tensor_ = Ort::Value::CreateTensor( + output_memory_info, output_buffer.get(), output_byte_count, + output_shape.data(), output_shape.size(), output_tensor_type); + this->output_buffer_.emplace_back(std::move(output_buffer)); +} + +std::vector EdgeYOLOONNXRuntime::inference(const cv::Mat &frame) +{ + // preprocess + cv::Mat pr_img = static_resize(frame); + + float *blob_data = (float *)(this->input_buffer_[0].get()); + blobFromImage(pr_img, blob_data); + + const char *input_names_[] = {this->input_name_.c_str()}; + const char *output_names_[] = {this->output_name_.c_str()}; + + // Inference + Ort::RunOptions run_options; + this->session_.Run(run_options, input_names_, &this->input_tensor_, 1, + output_names_, &this->output_tensor_, 1); + + float *net_pred = (float *)this->output_buffer_[0].get(); + + // post process + float scale = std::fminf(input_w_ / (frame.cols * 1.0f), + input_h_ / (frame.rows * 1.0f)); + std::vector objects; + decode_outputs(net_pred, this->num_array_, objects, + this->bbox_conf_thresh_, scale, frame.cols, frame.rows); + return objects; +} + +} // namespace edgeyolo_cpp diff --git a/src/edgeyolo/edgeyolo_onnxruntime.hpp b/src/edgeyolo/edgeyolo_onnxruntime.hpp new file mode 100644 index 0000000..8163fb2 --- /dev/null +++ b/src/edgeyolo/edgeyolo_onnxruntime.hpp @@ -0,0 +1,47 @@ +#ifndef _EdgeYOLO_CPP_EdgeYOLO_ONNX_HPP +#define _EdgeYOLO_CPP_EdgeYOLO_ONNX_HPP + +#include +#include +#include +#include +#include +#include + +#include + +#include "core.hpp" +#include "coco_names.hpp" + +namespace edgeyolo_cpp +{ + class EdgeYOLOONNXRuntime : public AbcEdgeYOLO + { + public: + EdgeYOLOONNXRuntime(file_name_t path_to_model, + int intra_op_num_threads, int inter_op_num_threads = 1, + const std::string& use_gpu = "", int device_id = 0, bool use_parallel = false, + float nms_th = 0.45, float conf_th = 0.3, + int num_classes = 80); + std::vector inference(const cv::Mat &frame) override; + + private: + int intra_op_num_threads_ = 1; + int inter_op_num_threads_ = 1; + int device_id_ = 0; + std::string use_gpu = ""; + bool use_parallel_ = false; + + Ort::Session session_{nullptr}; + Ort::Env env_{ORT_LOGGING_LEVEL_WARNING, "Default"}; + + Ort::Value input_tensor_{nullptr}; + Ort::Value output_tensor_{nullptr}; + std::string input_name_; + std::string output_name_; + std::vector> input_buffer_; + std::vector> output_buffer_; + }; +} + +#endif diff --git a/src/edgeyolo/utils.hpp b/src/edgeyolo/utils.hpp new file mode 100644 index 0000000..5fd6734 --- /dev/null +++ b/src/edgeyolo/utils.hpp @@ -0,0 +1,78 @@ +#ifndef _EdgeYOLO_CPP_UTILS_HPP +#define _EdgeYOLO_CPP_UTILS_HPP + +#include +#include +#include +#include +#include "core.hpp" +#include "coco_names.hpp" + +namespace edgeyolo_cpp +{ + namespace utils + { + + static std::vector read_class_labels_file(file_name_t file_name) + { + std::vector class_names; + std::ifstream ifs(file_name); + std::string buff; + if (ifs.fail()) + { + return class_names; + } + while (getline(ifs, buff)) + { + if (buff == "") + continue; + class_names.push_back(buff); + } + return class_names; + } + + static void draw_objects(cv::Mat bgr, const std::vector &objects, const std::vector &class_names = COCO_CLASSES) + { + + for (size_t i = 0; i < objects.size(); i++) + { + const Object &obj = objects[i]; + + int color_index = obj.label % 80; + cv::Scalar color = cv::Scalar(color_list[color_index][0], color_list[color_index][1], color_list[color_index][2]); + float c_mean = (float)(cv::mean(color)[0]); + cv::Scalar txt_color; + if (c_mean > 0.5) + { + txt_color = cv::Scalar(0, 0, 0); + } + else + { + txt_color = cv::Scalar(255, 255, 255); + } + + cv::rectangle(bgr, obj.rect, color * 255, 2); + + char text[256]; + sprintf(text, "%s %.1f%%", class_names[obj.label].c_str(), obj.prob * 100); + + int baseLine = 0; + cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.4, 1, &baseLine); + + cv::Scalar txt_bk_color = color * 0.7 * 255; + + int x = (int)(obj.rect.x); + int y = (int)(obj.rect.y + 1); + if (y > bgr.rows) + y = bgr.rows; + + cv::rectangle(bgr, cv::Rect(cv::Point(x, y), cv::Size(label_size.width, label_size.height + baseLine)), + txt_bk_color, -1); + + cv::putText(bgr, text, cv::Point(x, y + label_size.height), + cv::FONT_HERSHEY_SIMPLEX, 0.4, txt_color, 1); + } + } + } +} +#endif diff --git a/src/main.cpp b/src/main.cpp new file mode 100644 index 0000000..7614ab1 --- /dev/null +++ b/src/main.cpp @@ -0,0 +1,70 @@ +#include + +#include "core.hpp" +#include "edgeyolo_onnxruntime.hpp" +#include "utils.hpp" +#include "coco_names.hpp" + +#include + +cv::Mat colorImageCallback( + const cv::Mat &frame, + std::unique_ptr &edgeyolo_) +{ + // fps + auto now = std::chrono::system_clock::now(); + + auto objects = edgeyolo_->inference(frame); + + auto end = std::chrono::system_clock::now(); + auto elapsed = std::chrono::duration_cast( + end - now); + printf("Inference: %f FPS\n", 1000.0f / elapsed.count()); + printf("OBJECTS: %ld\n", objects.size()); + + cv::Mat out_frame = frame.clone(); + edgeyolo_cpp::utils::draw_objects(out_frame, objects); + + return out_frame; +} + +int main(int argc, char **argv) +{ + + if (argc < 2) { + std::cerr << "Usage: " << argv[0] << " \n"; + return 1; + } + + std::wstring model_path_ = + std::wstring(argv[1], argv[1] + strlen(argv[1])); + + // parameters + int onnxruntime_intra_op_num_threads_ = 4; + int onnxruntime_inter_op_num_threads_ = 4; + bool onnxruntime_use_cuda_ = false; + int onnxruntime_device_id_ = 0; + bool onnxruntime_use_parallel_ = false; + float nms_th_ = 0.45; + float conf_th_ = 0.3; + int num_classes_ = edgeyolo_cpp::COCO_CLASSES.size(); + + // Load model + auto edgeyolo_ = std::make_unique( + model_path_, onnxruntime_intra_op_num_threads_, + onnxruntime_inter_op_num_threads_, onnxruntime_use_cuda_, + onnxruntime_device_id_, onnxruntime_use_parallel_, nms_th_, + conf_th_, num_classes_); + + cv::VideoCapture cap("wqctLW0Hb_0.mp4"); + // grab frames from camera and run inference + cv::Mat frame; + while (cap.read(frame)) { + cv::Mat frame_draw = colorImageCallback(frame, edgeyolo_); + cv::imshow("frame", frame_draw); + if (cv::waitKey(1) == 27) + break; + } + + return 0; +} diff --git a/src/obs-utils/obs-config-utils.cpp b/src/obs-utils/obs-config-utils.cpp new file mode 100644 index 0000000..291dbde --- /dev/null +++ b/src/obs-utils/obs-config-utils.cpp @@ -0,0 +1,77 @@ +#include "obs-config-utils.h" +#include "plugin-support.h" + +#include +#include +#include + +void create_config_folder() +{ + char *config_folder_path = obs_module_config_path(""); + if (config_folder_path == nullptr) { + obs_log(LOG_ERROR, "Failed to get config folder path"); + return; + } + std::filesystem::path config_folder_std_path(config_folder_path); + bfree(config_folder_path); + + // create the folder if it doesn't exist + if (!std::filesystem::exists(config_folder_std_path)) { +#ifdef _WIN32 + obs_log(LOG_INFO, "Config folder does not exist, creating: %S", + config_folder_std_path.c_str()); +#else + obs_log(LOG_INFO, "Config folder does not exist, creating: %s", + config_folder_std_path.c_str()); +#endif + // Create the config folder + std::filesystem::create_directories(config_folder_std_path); + } +} + +int getConfig(config_t **config) +{ + create_config_folder(); // ensure the config folder exists + + // Get the config file + char *config_file_path = obs_module_config_path("config.ini"); + + int ret = config_open(config, config_file_path, CONFIG_OPEN_EXISTING); + if (ret != CONFIG_SUCCESS) { + obs_log(LOG_INFO, "Failed to open config file %s", + config_file_path); + return OBS_BGREMOVAL_CONFIG_FAIL; + } + + return OBS_BGREMOVAL_CONFIG_SUCCESS; +} + +int getFlagFromConfig(const char *name, bool *returnValue, bool defaultValue) +{ + // Get the config file + config_t *config; + if (getConfig(&config) != OBS_BGREMOVAL_CONFIG_SUCCESS) { + *returnValue = defaultValue; + return OBS_BGREMOVAL_CONFIG_FAIL; + } + + *returnValue = config_get_bool(config, "config", name); + config_close(config); + + return OBS_BGREMOVAL_CONFIG_SUCCESS; +} + +int setFlagInConfig(const char *name, const bool value) +{ + // Get the config file + config_t *config; + if (getConfig(&config) != OBS_BGREMOVAL_CONFIG_SUCCESS) { + return OBS_BGREMOVAL_CONFIG_FAIL; + } + + config_set_bool(config, "config", name, value); + config_save(config); + config_close(config); + + return OBS_BGREMOVAL_CONFIG_SUCCESS; +} diff --git a/src/obs-utils/obs-config-utils.h b/src/obs-utils/obs-config-utils.h new file mode 100644 index 0000000..cf79b31 --- /dev/null +++ b/src/obs-utils/obs-config-utils.h @@ -0,0 +1,30 @@ +#ifndef OBS_CONFIG_UTILS_H +#define OBS_CONFIG_UTILS_H + +enum { + OBS_BGREMOVAL_CONFIG_SUCCESS = 0, + OBS_BGREMOVAL_CONFIG_FAIL = 1, +}; + +/** + * Get a boolean flasg from the module configuration file. + * + * @param name The name of the config item. + * @param returnValue The value of the config item. + * @param defaultValue The default value of the config item. + * @return OBS_BGREMOVAL_CONFIG_SUCCESS if the config item was found, + * OBS_BGREMOVAL_CONFIG_FAIL otherwise. + */ +int getFlagFromConfig(const char *name, bool *returnValue, bool defaultValue); + +/** + * Set a boolean flag in the module configuration file. + * + * @param name The name of the config item. + * @param value The value of the config item. + * @return OBS_BGREMOVAL_CONFIG_SUCCESS if the config item was found, + * OBS_BGREMOVAL_CONFIG_FAIL otherwise. + */ +int setFlagInConfig(const char *name, const bool value); + +#endif /* OBS_CONFIG_UTILS_H */ diff --git a/src/obs-utils/obs-utils.cpp b/src/obs-utils/obs-utils.cpp new file mode 100644 index 0000000..c65e369 --- /dev/null +++ b/src/obs-utils/obs-utils.cpp @@ -0,0 +1,73 @@ +#include "obs-utils.h" + +#include + +/** + * @brief Get RGBA from the stage surface + * + * @param tf The filter data + * @param width The width of the stage surface (output) + * @param height The height of the stage surface (output) + * @return true if successful + * @return false if unsuccessful +*/ +bool getRGBAFromStageSurface(filter_data *tf, uint32_t &width, uint32_t &height) +{ + + if (!obs_source_enabled(tf->source)) { + return false; + } + + obs_source_t *target = obs_filter_get_target(tf->source); + if (!target) { + return false; + } + width = obs_source_get_base_width(target); + height = obs_source_get_base_height(target); + if (width == 0 || height == 0) { + return false; + } + gs_texrender_reset(tf->texrender); + if (!gs_texrender_begin(tf->texrender, width, height)) { + return false; + } + struct vec4 background; + vec4_zero(&background); + gs_clear(GS_CLEAR_COLOR, &background, 0.0f, 0); + gs_ortho(0.0f, static_cast(width), 0.0f, + static_cast(height), -100.0f, 100.0f); + gs_blend_state_push(); + gs_blend_function(GS_BLEND_ONE, GS_BLEND_ZERO); + obs_source_video_render(target); + gs_blend_state_pop(); + gs_texrender_end(tf->texrender); + + if (tf->stagesurface) { + uint32_t stagesurf_width = + gs_stagesurface_get_width(tf->stagesurface); + uint32_t stagesurf_height = + gs_stagesurface_get_height(tf->stagesurface); + if (stagesurf_width != width || stagesurf_height != height) { + gs_stagesurface_destroy(tf->stagesurface); + tf->stagesurface = nullptr; + } + } + if (!tf->stagesurface) { + tf->stagesurface = + gs_stagesurface_create(width, height, GS_BGRA); + } + gs_stage_texture(tf->stagesurface, + gs_texrender_get_texture(tf->texrender)); + uint8_t *video_data; + uint32_t linesize; + if (!gs_stagesurface_map(tf->stagesurface, &video_data, &linesize)) { + return false; + } + { + std::lock_guard lock(tf->inputBGRALock); + tf->inputBGRA = + cv::Mat(height, width, CV_8UC4, video_data, linesize); + } + gs_stagesurface_unmap(tf->stagesurface); + return true; +} diff --git a/src/obs-utils/obs-utils.h b/src/obs-utils/obs-utils.h new file mode 100644 index 0000000..d00076f --- /dev/null +++ b/src/obs-utils/obs-utils.h @@ -0,0 +1,9 @@ +#ifndef OBS_UTILS_H +#define OBS_UTILS_H + +#include "FilterData.h" + +bool getRGBAFromStageSurface(filter_data *tf, uint32_t &width, + uint32_t &height); + +#endif /* OBS_UTILS_H */ diff --git a/src/plugin-main.c b/src/plugin-main.c index bd4976c..d2a16e7 100644 --- a/src/plugin-main.c +++ b/src/plugin-main.c @@ -22,8 +22,16 @@ with this program. If not, see OBS_DECLARE_MODULE() OBS_MODULE_USE_DEFAULT_LOCALE(PLUGIN_NAME, "en-US") +MODULE_EXPORT const char *obs_module_description(void) +{ + return obs_module_text("DetectFilterPlugin"); +} + +extern struct obs_source_info detect_filter_info; + bool obs_module_load(void) { + obs_register_source(&detect_filter_info); obs_log(LOG_INFO, "plugin loaded successfully (version %s)", PLUGIN_VERSION); return true;