Merge pull request #68 from kounoike/feature/windows-cuda

add CUDA build for windows

.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

.gitignore

@@ -44,5 +44,6 @@ CTestTestfile.cmake
 _deps
 
 build/
+build_gpu/
 deps/
 release/

.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
+}

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 $<OR:$<CONFIG:Release>,$<CONFIG:RelWithDebInfo>>==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 $<OR:$<CONFIG:Release>,$<CONFIG:RelWithDebInfo>>==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 $<CONFIG:RelWithDebInfo>==1 (
             "${CMAKE_COMMAND}" -E copy
@@ -136,11 +128,55 @@ if(WIN32)
                 "${LIBOBS_INCLUDE_DIR}/../${OBS_BUILDDIR_ARCH}/rundir/$<CONFIG>/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 $<OR:$<CONFIG:Release>,$<CONFIG:RelWithDebInfo>>==1 (
+          "${CMAKE_COMMAND}" -E copy
+              "${OnnxRuntimePath}/lib/onnxruntime.dll"
+              "${RELEASE_DIR}/obs-plugins/${ARCH_NAME}"
+      )
+
+      # package the CUDA runtime DLL
+      COMMAND if $<OR:$<CONFIG:Release>,$<CONFIG:RelWithDebInfo>>==1 (
+          "${CMAKE_COMMAND}" -E copy
+          "${OnnxRuntimePath}/lib/onnxruntime_providers_cuda.dll"
+          "${RELEASE_DIR}/obs-plugins/${ARCH_NAME}"
+      )
+      COMMAND if $<OR:$<CONFIG:Release>,$<CONFIG:RelWithDebInfo>>==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 $<OR:$<CONFIG:Release>,$<CONFIG:RelWithDebInfo>>==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 $<OR:$<CONFIG:Release>,$<CONFIG:RelWithDebInfo>>==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)

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}

model/.gitignore

@@ -0,0 +1,3 @@
+*
+!README.md
+!.gitignore

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
+

scripts/build_obs.bat

@@ -1,4 +1,4 @@
-@echo off
+@echo on
 setlocal enabledelayedexpansion
 
 cd %~dp0

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%

src/obs-virtualbg-detector.cpp

@@ -7,13 +7,18 @@
 #include <obs.h>
 
 #ifdef _WIN32
+#ifdef USE_CUDA
+// no include file required
+#else
 #include <dml_provider_factory.h>
 #endif
+#endif
 
 #ifdef __APPLE__
 #include <onnxruntime/core/session/onnxruntime_cxx_api.h>
 // #include <onnxruntime/core/providers/cpu/cpu_provider_factory.h>
 #else
+#include <onnxruntime_c_api.h>
 #include <onnxruntime_cxx_api.h>
 #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