diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 476adae..297d55d 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -22,7 +22,7 @@ jobs: uses: actions/upload-artifact@v2 with: name: windows-archive - path: build/obs-virtualbg-*.zip + path: build*/obs-virtualbg-*.zip build_mac: name: build on Mac runs-on: macos-latest diff --git a/.gitignore b/.gitignore index 85c2e84..8cc9e2c 100644 --- a/.gitignore +++ b/.gitignore @@ -44,5 +44,6 @@ CTestTestfile.cmake _deps build/ +build_gpu/ deps/ release/ diff --git a/.vscode/c_cpp_properties.json b/.vscode/c_cpp_properties.json new file mode 100644 index 0000000..1ea2cf9 --- /dev/null +++ b/.vscode/c_cpp_properties.json @@ -0,0 +1,23 @@ +{ + "configurations": [ + { + "name": "Win32", + "includePath": [ + "${workspaceFolder}/**", + "${workspaceFolder}/deps/onnxruntime-win-x64-gpu-1.10.0/include" + ], + "defines": [ + "_DEBUG", + "UNICODE", + "_UNICODE", + "USE_CUDA=1" + ], + "windowsSdkVersion": "10.0.19041.0", + "compilerPath": "C:/Program Files (x86)/Microsoft Visual Studio/2019/Community/VC/Tools/MSVC/14.29.30133/bin/Hostx64/x64/cl.exe", + "cStandard": "c17", + "cppStandard": "c++17", + "intelliSenseMode": "windows-msvc-x64" + } + ], + "version": 4 +} \ No newline at end of file diff --git a/CMakeLists.txt b/CMakeLists.txt index 2a6ee7c..86ced8d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -19,6 +19,8 @@ find_package(Halide REQUIRED COMPONENTS shared) # find_package(Halide REQUIRED) set(RELEASE_DIR "${PROJECT_SOURCE_DIR}/Release" CACHE STRING "Directory for building release packages") +set(OBS_VIRTUALBG_USE_CUDA OFF CACHE BOOL "Use CUDA or not") + set(SOURCES src/plugin.cpp @@ -65,6 +67,10 @@ if(WIN32) "${OBS_FRONTEND_LIB}" ) + if(${OBS_VIRTUALBG_USE_CUDA}) + target_compile_definitions(${CMAKE_PROJECT_NAME} PRIVATE USE_CUDA=1) + endif() + # --- Release package helper --- # The "release" folder has a structure similar OBS' one on Windows set(RELEASE_DIR "${PROJECT_SOURCE_DIR}/release") @@ -91,20 +97,6 @@ if(WIN32) "${RELEASE_DIR}/obs-plugins/${ARCH_NAME}" ) - # package the onnxruntime DLL with the plugin DLL - COMMAND if $,$>==1 ( - "${CMAKE_COMMAND}" -E copy - "${OnnxRuntimePath}/runtimes/win-x64/native/onnxruntime.dll" - "${RELEASE_DIR}/obs-plugins/${ARCH_NAME}" - ) - - # package the DirectML runtime DLL - COMMAND if $,$>==1 ( - "${CMAKE_COMMAND}" -E copy - "${PROJECT_SOURCE_DIR}/deps/directml/bin/x64-win/DirectML.dll" - "${RELEASE_DIR}/obs-plugins/${ARCH_NAME}" - ) - # If config is RelWithDebInfo, copy the pdb file COMMAND if $==1 ( "${CMAKE_COMMAND}" -E copy @@ -136,11 +128,55 @@ if(WIN32) "${LIBOBS_INCLUDE_DIR}/../${OBS_BUILDDIR_ARCH}/rundir/$/data/obs-plugins/${CMAKE_PROJECT_NAME}" ) ) + + + if(${OBS_VIRTUALBG_USE_CUDA}) + add_custom_command(TARGET ${CMAKE_PROJECT_NAME} POST_BUILD + # package the onnxruntime DLL with the plugin DLL + COMMAND if $,$>==1 ( + "${CMAKE_COMMAND}" -E copy + "${OnnxRuntimePath}/lib/onnxruntime.dll" + "${RELEASE_DIR}/obs-plugins/${ARCH_NAME}" + ) + + # package the CUDA runtime DLL + COMMAND if $,$>==1 ( + "${CMAKE_COMMAND}" -E copy + "${OnnxRuntimePath}/lib/onnxruntime_providers_cuda.dll" + "${RELEASE_DIR}/obs-plugins/${ARCH_NAME}" + ) + COMMAND if $,$>==1 ( + "${CMAKE_COMMAND}" -E copy + "${OnnxRuntimePath}/lib/onnxruntime_providers_shared.dll" + "${RELEASE_DIR}/obs-plugins/${ARCH_NAME}" + ) + ) +else() +add_custom_command(TARGET ${CMAKE_PROJECT_NAME} POST_BUILD + # package the onnxruntime DLL with the plugin DLL + COMMAND if $,$>==1 ( + "${CMAKE_COMMAND}" -E copy + "${OnnxRuntimePath}/runtimes/win-x64/native/onnxruntime.dll" + "${RELEASE_DIR}/obs-plugins/${ARCH_NAME}" + ) + + # package the DirectML runtime DLL + COMMAND if $,$>==1 ( + "${CMAKE_COMMAND}" -E copy + "${PROJECT_SOURCE_DIR}/deps/directml/bin/x64-win/DirectML.dll" + "${RELEASE_DIR}/obs-plugins/${ARCH_NAME}" + ) + ) + endif() + install(DIRECTORY "${RELEASE_DIR}/" DESTINATION .) set(CPACK_GENERATOR "ZIP") set(CPACK_INCLUDE_TOPLEVEL_DIRECTORY OFF) set(CPACK_VERBATIM_VARIABLES YES) set(CPACK_PACKAGE_VERSION v${SEMVER}) + if(${OBS_VIRTUALBG_USE_CUDA}) + set(CPACK_PACKAGE_NAME "${CMAKE_PROJECT_NAME}-cuda") + endif() include(CPack) elseif(APPLE) # set(CMAKE_OSX_DEPLOYMENT_TARGET 10.13) diff --git a/cmake/FindOnnxRuntime.cmake b/cmake/FindOnnxRuntime.cmake index 4ec663d..197a13e 100644 --- a/cmake/FindOnnxRuntime.cmake +++ b/cmake/FindOnnxRuntime.cmake @@ -5,6 +5,7 @@ if(NOT WIN32 AND EXISTS ${OnnxRuntime_DEPS_DIR}) set(OnnxRuntime_INCLUDE_DIR ${OnnxRuntime_DEPS_DIR}/include) message("OnnxRuntime_DEPS_DIR Found!!!!!! ${OnnxRuntime_DEPS_DIR} ${OnnxRuntime_LIBRARIES}") else() + message("kotti [${OnnxRuntimePath}]") find_path(OnnxRuntime_INCLUDE_DIR NAMES "onnxruntime_c_api.h" "onnxruntime/core/session/onnxruntime_c_api.h" @@ -18,7 +19,7 @@ else() ) find_library(OnnxRuntime_LIBRARIES - NAMES onnxruntime.lib onnxruntime_providers_shared.lib libonnxruntime.dylib + NAMES onnxruntime.lib onnxruntime_providers_shared.lib onnxruntime_providers_cuda.lib libonnxruntime.dylib PATHS ENV OnnxRuntimePath ${OnnxRuntimePath} diff --git a/data/model.onnx b/data/model.onnx index 8cd5bb1..801391d 100644 Binary files a/data/model.onnx and b/data/model.onnx differ diff --git a/model/.gitignore b/model/.gitignore new file mode 100644 index 0000000..c228bb2 --- /dev/null +++ b/model/.gitignore @@ -0,0 +1,3 @@ +* +!README.md +!.gitignore diff --git a/model/README.md b/model/README.md new file mode 100644 index 0000000..b9834e4 --- /dev/null +++ b/model/README.md @@ -0,0 +1,11 @@ +curl https://github.com/google/mediapipe/raw/master/mediapipe/modules/selfie_segmentation/selfie_segmentation_landscape.tflite +docker pull ghcr.io/pinto0309/tflite2tensorflow:latest + +docker run ...... + +tflite2tensorflow \ + --model_path segm_full_v679.tflite \ + --flatc_path ../flatc \ + --schema_path ../schema.fbs \ + --output_pb + diff --git a/scripts/build_obs.bat b/scripts/build_obs.bat index a03a1fe..303293a 100644 --- a/scripts/build_obs.bat +++ b/scripts/build_obs.bat @@ -1,4 +1,4 @@ -@echo off +@echo on setlocal enabledelayedexpansion cd %~dp0 diff --git a/scripts/build_win.bat b/scripts/build_win.bat index d2a14ab..931f087 100644 --- a/scripts/build_win.bat +++ b/scripts/build_win.bat @@ -13,6 +13,10 @@ set ONNXRUNTIME_URL=https://github.com/microsoft/onnxruntime/releases/download/v set ONNXRUNTIME_ZIP=%DEPS_DIR%\Microsoft.ML.OnnxRuntime.DirectML.%ONNXRUNTIME_VERSION%.zip set ONNXRUNTIME_DIR=%DEPS_DIR%\onnxruntime +set ONNXRUNTIME_GPU_URL=https://github.com/microsoft/onnxruntime/releases/download/v%ONNXRUNTIME_VERSION%/onnxruntime-win-x64-gpu-%ONNXRUNTIME_VERSION%.zip +set ONNXRUNTIME_GPU_ZIP=%DEPS_DIR%\onnxruntime-win-x64-gpu-%ONNXRUNTIME_VERSION%.zip +set ONNXRUNTIME_GPU_DIR=%DEPS_DIR%\onnxruntime-win-x64-gpu-%ONNXRUNTIME_VERSION% + set DIRECTML_URL=https://www.nuget.org/api/v2/package/Microsoft.AI.DirectML/%DIRECTML_VERSION% set DIRECTML_ZIP=%DEPS_DIR%\Microsoft.AI.DirectML-%DIRECTML_VERSION%.zip set DIRECTML_DIR=%DEPS_DIR%\directml @@ -28,6 +32,9 @@ pushd %DEPS_DIR% IF not exist %ONNXRUNTIME_ZIP% curl -L -o %ONNXRUNTIME_ZIP% %ONNXRUNTIME_URL% IF not exist %ONNXRUNTIME_DIR% 7z x -o%ONNXRUNTIME_DIR% %ONNXRUNTIME_ZIP% + IF not exist %ONNXRUNTIME_GPU_ZIP% curl -L -o %ONNXRUNTIME_GPU_ZIP% %ONNXRUNTIME_GPU_URL% + IF not exist %ONNXRUNTIME_GPU_DIR% 7z x -o%DEPS_DIR% %ONNXRUNTIME_GPU_ZIP% + IF not exist %DIRECTML_ZIP% curl -L -o %DIRECTML_ZIP% %DIRECTML_URL% IF not exist %DIRECTML_DIR% 7z x -o%DIRECTML_DIR% %DIRECTML_ZIP% @@ -56,6 +63,24 @@ pushd build IF ERRORLEVEL 1 GOTO ERR popd +IF exist build_gpu rmdir /s /q build_gpu +mkdir build_gpu +pushd build_gpu + cmake ^ + -DOBS_VIRTUALBG_USE_CUDA=ON ^ + -DobsPath=%DEPS_DIR%\obs-studio ^ + -DOnnxRuntimePath=%ONNXRUNTIME_GPU_DIR% ^ + -DCMAKE_SYSTEM_VERSION=10.0.18363.657 ^ + -DHalide_DIR=%HALIDE_DIR%\lib\cmake\Halide ^ + -DHalideHelpers_DIR=%HALIDE_DIR%\lib\cmake\HalideHelpers ^ + .. + IF ERRORLEVEL 1 GOTO ERR + cmake --build . --config RelWithDebInfo + IF ERRORLEVEL 1 GOTO ERR + cpack + IF ERRORLEVEL 1 GOTO ERR +popd + goto END :ERR echo BUILD ERROR %ERRORLEVEL% diff --git a/src/obs-virtualbg-detector.cpp b/src/obs-virtualbg-detector.cpp index a15650b..6939516 100644 --- a/src/obs-virtualbg-detector.cpp +++ b/src/obs-virtualbg-detector.cpp @@ -7,13 +7,18 @@ #include #ifdef _WIN32 +#ifdef USE_CUDA +// no include file required +#else #include #endif +#endif #ifdef __APPLE__ #include // #include #else +#include #include #endif @@ -86,6 +91,16 @@ void detector_destroy(void *data) { } void detector_setup_ort_session_gpu(Ort::SessionOptions &sessionOptions) { +#if USE_CUDA + try { + Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(sessionOptions, 0)); + } catch (const std::exception &ex) { + blog(LOG_ERROR, + "[Virtual BG detector] Can't Append Execution Provider CUDA. Will use CPU inference (its so heavy). " + "error: %s", + ex.what()); + } +#else #if _WIN32 try { sessionOptions.DisableMemPattern(); @@ -98,6 +113,7 @@ void detector_setup_ort_session_gpu(Ort::SessionOptions &sessionOptions) { ex.what()); } #endif +#endif } void detector_setup_ort_session_load_model(virtual_bg_filter_data *filter_data, @@ -203,7 +219,7 @@ void detector_update(void *data, obs_data_t *settings) { filter_data->threshold = (float)obs_data_get_double(settings, THRESHOLD_VALUE); filter_data->use_mask_blur = obs_data_get_bool(settings, USE_MASK_BLUR); -#ifdef _WIN32 +#if defined(_WIN32) || USE_CUDA filter_data->use_gpu = obs_data_get_bool(settings, USE_GPU); #endif @@ -246,7 +262,7 @@ void detector_defaults(obs_data_t *settings) { obs_data_set_default_bool(settings, USE_THRESHOLD, true); obs_data_set_default_double(settings, THRESHOLD_VALUE, 0.5); obs_data_set_default_bool(settings, USE_MASK_BLUR, true); -#if _WIN32 +#if defined(_WIN32) || USE_CUDA obs_data_set_default_bool(settings, USE_GPU, false); #endif } @@ -257,7 +273,7 @@ obs_properties_t *detector_properties(void *data) { obs_properties_add_bool(ppts, USE_THRESHOLD, obs_module_text(USE_THRESHOLD)); obs_properties_add_float_slider(ppts, THRESHOLD_VALUE, obs_module_text(THRESHOLD_VALUE), 0.0, 1.0, 0.05); obs_properties_add_bool(ppts, USE_MASK_BLUR, obs_module_text(USE_MASK_BLUR)); -#if _WIN32 +#if defined(_WIN32) || USE_CUDA obs_properties_add_bool(ppts, USE_GPU, obs_module_text(USE_GPU)); #endif