nvidia-greenscreen-filter.c 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739
  1. #include <obs-module.h>
  2. #include <util/threading.h>
  3. #include <dxgi.h>
  4. #include <d3d11.h>
  5. #include <d3d11_1.h>
  6. #include "nvvfx-load.h"
  7. /* -------------------------------------------------------- */
  8. #define do_log(level, format, ...) \
  9. blog(level, \
  10. "[NVIDIA RTX AI Greenscreen (Background removal): '%s'] " format, \
  11. obs_source_get_name(filter->context), ##__VA_ARGS__)
  12. #define warn(format, ...) do_log(LOG_WARNING, format, ##__VA_ARGS__)
  13. #define info(format, ...) do_log(LOG_INFO, format, ##__VA_ARGS__)
  14. #define error(format, ...) do_log(LOG_ERROR, format, ##__VA_ARGS__)
  15. #ifdef _DEBUG
  16. #define debug(format, ...) do_log(LOG_DEBUG, format, ##__VA_ARGS__)
  17. #else
  18. #define debug(format, ...)
  19. #endif
  20. /* -------------------------------------------------------- */
  21. #define S_MODE "mode"
  22. #define S_MODE_QUALITY 0
  23. #define S_MODE_PERF 1
  24. #define S_THRESHOLDFX "threshold"
  25. #define S_THRESHOLDFX_DEFAULT 1.0
  26. #define MT_ obs_module_text
  27. #define TEXT_MODE MT_("Greenscreen.Mode")
  28. #define TEXT_MODE_QUALITY MT_("Greenscreen.Quality")
  29. #define TEXT_MODE_PERF MT_("Greenscreen.Performance")
  30. #define TEXT_MODE_THRESHOLD MT_("Greenscreen.Threshold")
  31. bool nvvfx_loaded = false;
  32. struct nv_greenscreen_data {
  33. obs_source_t *context;
  34. bool images_allocated;
  35. bool initial_render;
  36. bool processing_stop;
  37. bool processed_frame;
  38. bool target_valid;
  39. volatile bool got_new_frame;
  40. /* RTX SDK vars */
  41. NvVFX_Handle handle;
  42. CUstream stream; // CUDA stream
  43. int mode; // 0 = quality, 1 = performance
  44. NvCVImage *src_img; // src img in obs format (RGBA ?) on GPU
  45. NvCVImage *BGR_src_img; // src img in BGR on GPU
  46. NvCVImage *A_dst_img; // mask img on GPU
  47. NvCVImage *dst_img; // mask texture
  48. NvCVImage *stage; // planar stage img used for transfer to texture
  49. /* alpha mask effect */
  50. gs_effect_t *effect;
  51. gs_texrender_t *render;
  52. gs_texture_t *alpha_texture;
  53. uint32_t width; // width of texture
  54. uint32_t height; // height of texture
  55. gs_eparam_t *mask_param;
  56. gs_texture_t *src_texture;
  57. gs_eparam_t *src_param;
  58. gs_eparam_t *threshold_param;
  59. double threshold;
  60. };
  61. static const char *nv_greenscreen_filter_name(void *unused)
  62. {
  63. UNUSED_PARAMETER(unused);
  64. return obs_module_text("NvidiaGreenscreenFilter");
  65. }
  66. static void nv_greenscreen_filter_update(void *data, obs_data_t *settings)
  67. {
  68. struct nv_greenscreen_data *filter = (struct nv_greenscreen_data *)data;
  69. NvCV_Status vfxErr;
  70. int mode = (int)obs_data_get_int(settings, S_MODE);
  71. if (filter->mode != mode) {
  72. filter->mode = mode;
  73. vfxErr = NvVFX_SetU32(filter->handle, NVVFX_MODE, mode);
  74. vfxErr = NvVFX_Load(filter->handle);
  75. if (NVCV_SUCCESS != vfxErr)
  76. error("Error loading AI Greenscreen FX %i", vfxErr);
  77. }
  78. filter->threshold =
  79. (double)obs_data_get_double(settings, S_THRESHOLDFX);
  80. }
  81. static void nv_greenscreen_filter_actual_destroy(void *data)
  82. {
  83. struct nv_greenscreen_data *filter = (struct nv_greenscreen_data *)data;
  84. if (!nvvfx_loaded) {
  85. bfree(filter);
  86. return;
  87. }
  88. filter->processing_stop = true;
  89. if (filter->images_allocated) {
  90. obs_enter_graphics();
  91. gs_texture_destroy(filter->alpha_texture);
  92. gs_texrender_destroy(filter->render);
  93. obs_leave_graphics();
  94. NvCVImage_Destroy(filter->src_img);
  95. NvCVImage_Destroy(filter->BGR_src_img);
  96. NvCVImage_Destroy(filter->A_dst_img);
  97. NvCVImage_Destroy(filter->dst_img);
  98. NvCVImage_Destroy(filter->stage);
  99. }
  100. if (filter->stream) {
  101. NvVFX_CudaStreamDestroy(filter->stream);
  102. }
  103. if (filter->handle) {
  104. NvVFX_DestroyEffect(filter->handle);
  105. }
  106. if (filter->effect) {
  107. obs_enter_graphics();
  108. gs_effect_destroy(filter->effect);
  109. obs_leave_graphics();
  110. }
  111. bfree(filter);
  112. }
  113. static void nv_greenscreen_filter_destroy(void *data)
  114. {
  115. obs_queue_task(OBS_TASK_GRAPHICS, nv_greenscreen_filter_actual_destroy,
  116. data, false);
  117. }
  118. static void init_images_greenscreen(struct nv_greenscreen_data *filter)
  119. {
  120. NvCV_Status vfxErr;
  121. uint32_t width = filter->width;
  122. uint32_t height = filter->height;
  123. /* 1. create alpha texture */
  124. obs_enter_graphics();
  125. if (filter->alpha_texture) {
  126. gs_texture_destroy(filter->alpha_texture);
  127. }
  128. filter->alpha_texture =
  129. gs_texture_create(width, height, GS_A8, 1, NULL, 0);
  130. if (filter->alpha_texture == NULL) {
  131. error("Alpha texture couldn't be created");
  132. goto fail;
  133. }
  134. struct ID3D11Texture2D *d11texture =
  135. (struct ID3D11Texture2D *)gs_texture_get_obj(
  136. filter->alpha_texture);
  137. obs_leave_graphics();
  138. /* 2. Create NvCVImage which will hold final alpha texture. */
  139. if (NvCVImage_Create(width, height, NVCV_A, NVCV_U8, NVCV_CHUNKY,
  140. NVCV_GPU, 1, &filter->dst_img) != NVCV_SUCCESS) {
  141. goto fail;
  142. }
  143. vfxErr = NvCVImage_InitFromD3D11Texture(filter->dst_img, d11texture);
  144. if (vfxErr != NVCV_SUCCESS) {
  145. const char *errString = NvCV_GetErrorStringFromCode(vfxErr);
  146. error("Error passing dst ID3D11Texture to img; error %i: %s",
  147. vfxErr, errString);
  148. goto fail;
  149. }
  150. /* 3. create texrender */
  151. obs_enter_graphics();
  152. if (filter->render)
  153. gs_texrender_destroy(filter->render);
  154. filter->render = gs_texrender_create(GS_BGRA_UNORM, GS_ZS_NONE);
  155. obs_leave_graphics();
  156. if (!filter->render) {
  157. error("Failed to create a texture renderer", vfxErr);
  158. goto fail;
  159. }
  160. /* 4. Create and allocate BGR NvCVImage (fx src). */
  161. if (NvCVImage_Create(width, height, NVCV_BGR, NVCV_U8, NVCV_CHUNKY,
  162. NVCV_GPU, 1,
  163. &filter->BGR_src_img) != NVCV_SUCCESS) {
  164. goto fail;
  165. }
  166. if (NvCVImage_Alloc(filter->BGR_src_img, width, height, NVCV_BGR,
  167. NVCV_U8, NVCV_CHUNKY, NVCV_GPU,
  168. 1) != NVCV_SUCCESS) {
  169. goto fail;
  170. }
  171. /* 5. Create and allocate Alpha NvCVimage (fx dst). */
  172. if (NvCVImage_Create(width, height, NVCV_A, NVCV_U8, NVCV_CHUNKY,
  173. NVCV_GPU, 1, &filter->A_dst_img) != NVCV_SUCCESS) {
  174. goto fail;
  175. }
  176. if (NvCVImage_Alloc(filter->A_dst_img, width, height, NVCV_A, NVCV_U8,
  177. NVCV_CHUNKY, NVCV_GPU, 1) != NVCV_SUCCESS) {
  178. goto fail;
  179. }
  180. /* 6. Create stage NvCVImage which will be used as buffer for transfer */
  181. if (NvCVImage_Create(width, height, NVCV_RGBA, NVCV_U8, NVCV_PLANAR,
  182. NVCV_GPU, 1, &filter->stage) != NVCV_SUCCESS) {
  183. goto fail;
  184. }
  185. if (NvCVImage_Alloc(filter->stage, width, height, NVCV_RGBA, NVCV_U8,
  186. NVCV_PLANAR, NVCV_GPU, 1) != NVCV_SUCCESS) {
  187. goto fail;
  188. }
  189. /* 7. Set input & output images for nv FX. */
  190. if (NvVFX_SetImage(filter->handle, NVVFX_INPUT_IMAGE,
  191. filter->BGR_src_img) != NVCV_SUCCESS) {
  192. goto fail;
  193. }
  194. if (NvVFX_SetImage(filter->handle, NVVFX_OUTPUT_IMAGE,
  195. filter->A_dst_img) != NVCV_SUCCESS) {
  196. goto fail;
  197. }
  198. filter->images_allocated = true;
  199. return;
  200. fail:
  201. error("Error during allocation of images");
  202. filter->processing_stop = true;
  203. return;
  204. }
  205. static bool process_texture_greenscreen(struct nv_greenscreen_data *filter)
  206. {
  207. gs_texrender_t *render = filter->render;
  208. NvCV_Status vfxErr;
  209. /* 1. Map src img holding texture. */
  210. vfxErr = NvCVImage_MapResource(filter->src_img, filter->stream);
  211. if (vfxErr != NVCV_SUCCESS) {
  212. const char *errString = NvCV_GetErrorStringFromCode(vfxErr);
  213. error("Error mapping resource for source texture; error %i : %s",
  214. vfxErr, errString);
  215. goto fail;
  216. }
  217. /* 2. Convert to BGR. */
  218. vfxErr = NvCVImage_Transfer(filter->src_img, filter->BGR_src_img, 1.0f,
  219. filter->stream, filter->stage);
  220. if (vfxErr != NVCV_SUCCESS) {
  221. const char *errString = NvCV_GetErrorStringFromCode(vfxErr);
  222. error("Error converting src to BGR img; error %i: %s", vfxErr,
  223. errString);
  224. goto fail;
  225. }
  226. vfxErr = NvCVImage_UnmapResource(filter->src_img, filter->stream);
  227. if (vfxErr != NVCV_SUCCESS) {
  228. const char *errString = NvCV_GetErrorStringFromCode(vfxErr);
  229. error("Error unmapping resource for src texture; error %i: %s",
  230. vfxErr, errString);
  231. goto fail;
  232. }
  233. /* 3. run RTX fx */
  234. vfxErr = NvVFX_Run(filter->handle, 1);
  235. if (vfxErr != NVCV_SUCCESS) {
  236. const char *errString = NvCV_GetErrorStringFromCode(vfxErr);
  237. error("Error running the FX; error %i: %s", vfxErr, errString);
  238. goto fail;
  239. }
  240. /* 4. Map dst texture before transfer from dst img provided by FX */
  241. vfxErr = NvCVImage_MapResource(filter->dst_img, filter->stream);
  242. if (vfxErr != NVCV_SUCCESS) {
  243. const char *errString = NvCV_GetErrorStringFromCode(vfxErr);
  244. error("Error mapping resource for dst texture; error %i: %s",
  245. vfxErr, errString);
  246. goto fail;
  247. }
  248. vfxErr = NvCVImage_Transfer(filter->A_dst_img, filter->dst_img, 1.0f,
  249. filter->stream, filter->stage);
  250. if (vfxErr != NVCV_SUCCESS) {
  251. const char *errString = NvCV_GetErrorStringFromCode(vfxErr);
  252. error("Error transferring mask to alpha texture; error %i: %s ",
  253. vfxErr, errString);
  254. goto fail;
  255. }
  256. vfxErr = NvCVImage_UnmapResource(filter->dst_img, filter->stream);
  257. if (vfxErr != NVCV_SUCCESS) {
  258. const char *errString = NvCV_GetErrorStringFromCode(vfxErr);
  259. error("Error unmapping resource for dst texture; error %i: %s",
  260. vfxErr, errString);
  261. goto fail;
  262. }
  263. return true;
  264. fail:
  265. filter->processing_stop = true;
  266. return false;
  267. }
  268. static void *nv_greenscreen_filter_create(obs_data_t *settings,
  269. obs_source_t *context)
  270. {
  271. struct nv_greenscreen_data *filter =
  272. (struct nv_greenscreen_data *)bzalloc(sizeof(*filter));
  273. if (!nvvfx_loaded) {
  274. nv_greenscreen_filter_destroy(filter);
  275. return NULL;
  276. }
  277. NvCV_Status vfxErr;
  278. filter->context = context;
  279. filter->mode = -1; // should be 0 or 1; -1 triggers an update
  280. filter->images_allocated = false;
  281. filter->processed_frame = true; // start processing when false
  282. filter->width = 0;
  283. filter->height = 0;
  284. filter->initial_render = false;
  285. filter->processing_stop = false;
  286. /* 1. Create FX */
  287. vfxErr = NvVFX_CreateEffect(NVVFX_FX_GREEN_SCREEN, &filter->handle);
  288. if (NVCV_SUCCESS != vfxErr) {
  289. const char *errString = NvCV_GetErrorStringFromCode(vfxErr);
  290. error("Error creating AI Greenscreen FX; error %i: %s", vfxErr,
  291. errString);
  292. nv_greenscreen_filter_destroy(filter);
  293. return NULL;
  294. }
  295. /* 2. Set models path & initialize CudaStream */
  296. char buffer[MAX_PATH];
  297. char modelDir[MAX_PATH];
  298. nvvfx_get_sdk_path(buffer, MAX_PATH);
  299. size_t max_len = sizeof(buffer) / sizeof(char);
  300. snprintf(modelDir, max_len, "%s\\models", buffer);
  301. vfxErr = NvVFX_SetString(filter->handle, NVVFX_MODEL_DIRECTORY,
  302. modelDir);
  303. vfxErr = NvVFX_CudaStreamCreate(&filter->stream);
  304. if (NVCV_SUCCESS != vfxErr) {
  305. const char *errString = NvCV_GetErrorStringFromCode(vfxErr);
  306. error("Error creating CUDA Stream; error %i: %s", vfxErr,
  307. errString);
  308. nv_greenscreen_filter_destroy(filter);
  309. return NULL;
  310. }
  311. vfxErr = NvVFX_SetCudaStream(filter->handle, NVVFX_CUDA_STREAM,
  312. filter->stream);
  313. if (NVCV_SUCCESS != vfxErr) {
  314. const char *errString = NvCV_GetErrorStringFromCode(vfxErr);
  315. error("Error setting CUDA Stream %i", vfxErr);
  316. nv_greenscreen_filter_destroy(filter);
  317. return NULL;
  318. }
  319. /* log sdk version */
  320. unsigned int version;
  321. if (NvVFX_GetVersion(&version) == NVCV_SUCCESS) {
  322. uint8_t major = (version >> 24) & 0xff;
  323. uint8_t minor = (version >> 16) & 0x00ff;
  324. uint8_t build = (version >> 8) & 0x0000ff;
  325. info("RTX VIDEO FX version: %i.%i.%i", major, minor, build);
  326. }
  327. /* 3. Load alpha mask effect. */
  328. char *effect_path = obs_module_file("rtx_greenscreen.effect");
  329. obs_enter_graphics();
  330. filter->effect = gs_effect_create_from_file(effect_path, NULL);
  331. bfree(effect_path);
  332. if (filter->effect) {
  333. filter->mask_param =
  334. gs_effect_get_param_by_name(filter->effect, "mask");
  335. filter->src_param =
  336. gs_effect_get_param_by_name(filter->effect, "image");
  337. filter->threshold_param = gs_effect_get_param_by_name(
  338. filter->effect, "threshold");
  339. }
  340. obs_leave_graphics();
  341. if (!filter->effect) {
  342. nv_greenscreen_filter_destroy(filter);
  343. return NULL;
  344. }
  345. /*---------------------------------------- */
  346. nv_greenscreen_filter_update(filter, settings);
  347. return filter;
  348. }
  349. static obs_properties_t *nv_greenscreen_filter_properties(void *data)
  350. {
  351. obs_properties_t *props = obs_properties_create();
  352. obs_property_t *mode = obs_properties_add_list(props, S_MODE, TEXT_MODE,
  353. OBS_COMBO_TYPE_LIST,
  354. OBS_COMBO_FORMAT_INT);
  355. obs_property_list_add_int(mode, TEXT_MODE_QUALITY, S_MODE_QUALITY);
  356. obs_property_list_add_int(mode, TEXT_MODE_PERF, S_MODE_PERF);
  357. obs_property_t *threshold = obs_properties_add_float_slider(
  358. props, S_THRESHOLDFX, TEXT_MODE_THRESHOLD, 0, 1, 0.05);
  359. UNUSED_PARAMETER(data);
  360. return props;
  361. }
  362. static void nv_greenscreen_filter_defaults(obs_data_t *settings)
  363. {
  364. obs_data_set_default_int(settings, S_MODE, S_MODE_QUALITY);
  365. obs_data_set_default_double(settings, S_THRESHOLDFX,
  366. S_THRESHOLDFX_DEFAULT);
  367. }
  368. static struct obs_source_frame *
  369. nv_greenscreen_filter_video(void *data, struct obs_source_frame *frame)
  370. {
  371. struct nv_greenscreen_data *filter = (struct nv_greenscreen_data *)data;
  372. os_atomic_set_bool(&filter->got_new_frame, true);
  373. return frame;
  374. }
  375. static void nv_greenscreen_filter_tick(void *data, float t)
  376. {
  377. UNUSED_PARAMETER(t);
  378. struct nv_greenscreen_data *filter = (struct nv_greenscreen_data *)data;
  379. if (filter->processing_stop) {
  380. return;
  381. }
  382. if (!obs_filter_get_target(filter->context)) {
  383. return;
  384. }
  385. obs_source_t *target = obs_filter_get_target(filter->context);
  386. uint32_t cx;
  387. uint32_t cy;
  388. filter->target_valid = true;
  389. cx = obs_source_get_base_width(target);
  390. cy = obs_source_get_base_height(target);
  391. // initially the sizes are 0
  392. if (!cx && !cy) {
  393. filter->target_valid = false;
  394. return;
  395. }
  396. /* minimum size supported by SDK is (512,288) */
  397. filter->target_valid = cx >= 512 && cy >= 288;
  398. if (!filter->target_valid) {
  399. error("Size must be larger than (512,288)");
  400. return;
  401. }
  402. if (cx != filter->width && cy != filter->height) {
  403. filter->images_allocated = false;
  404. filter->width = cx;
  405. filter->height = cy;
  406. }
  407. if (!filter->images_allocated) {
  408. init_images_greenscreen(filter);
  409. filter->initial_render = 0;
  410. }
  411. filter->processed_frame = false;
  412. }
  413. static void draw_greenscreen(struct nv_greenscreen_data *filter)
  414. {
  415. /* Render alpha mask */
  416. if (!obs_source_process_filter_begin(filter->context, GS_BGRA_UNORM,
  417. OBS_ALLOW_DIRECT_RENDERING)) {
  418. return;
  419. }
  420. gs_effect_set_texture(filter->mask_param, filter->alpha_texture);
  421. gs_effect_set_texture(filter->src_param, filter->src_texture);
  422. gs_effect_set_float(filter->threshold_param, (float)filter->threshold);
  423. while (gs_effect_loop(filter->effect, "Draw")) {
  424. gs_draw_sprite(NULL, 0, filter->width, filter->height);
  425. }
  426. obs_source_process_filter_end(filter->context, filter->effect, 0, 0);
  427. }
  428. static void draw_greenscreen_srgb(struct nv_greenscreen_data *filter)
  429. {
  430. /* Render alpha mask */
  431. if (!obs_source_process_filter_begin(filter->context, GS_BGRA_UNORM,
  432. OBS_ALLOW_DIRECT_RENDERING)) {
  433. return;
  434. }
  435. const bool previous = gs_framebuffer_srgb_enabled();
  436. gs_enable_framebuffer_srgb(true);
  437. gs_effect_set_texture_srgb(filter->mask_param, filter->alpha_texture);
  438. gs_effect_set_texture_srgb(filter->src_param, filter->src_texture);
  439. gs_effect_set_float(filter->threshold_param, (float)filter->threshold);
  440. while (gs_effect_loop(filter->effect, "Draw")) {
  441. gs_draw_sprite(NULL, 0, filter->width, filter->height);
  442. }
  443. gs_enable_framebuffer_srgb(previous);
  444. obs_source_process_filter_end(filter->context, filter->effect, 0, 0);
  445. }
  446. static void nv_greenscreen_filter_render(void *data, gs_effect_t *effect)
  447. {
  448. NvCV_Status vfxErr;
  449. bool ret;
  450. struct nv_greenscreen_data *filter = (struct nv_greenscreen_data *)data;
  451. if (filter->processing_stop)
  452. return;
  453. obs_source_t *target = obs_filter_get_target(filter->context);
  454. obs_source_t *parent = obs_filter_get_parent(filter->context);
  455. gs_texrender_t *render;
  456. /* Skip if processing of a frame hasn't yet started */
  457. if (!filter->target_valid || !target || !parent ||
  458. filter->processed_frame) {
  459. obs_source_skip_video_filter(filter->context);
  460. return;
  461. }
  462. /* 1. Render to retrieve texture. */
  463. render = filter->render;
  464. if (!render) {
  465. obs_source_skip_video_filter(filter->context);
  466. return;
  467. }
  468. uint32_t target_flags = obs_source_get_output_flags(target);
  469. uint32_t parent_flags = obs_source_get_output_flags(parent);
  470. bool custom_draw = (target_flags & OBS_SOURCE_CUSTOM_DRAW) != 0;
  471. bool async = (target_flags & OBS_SOURCE_ASYNC) != 0;
  472. bool srgb_draw = (parent_flags & OBS_SOURCE_SRGB) != 0;
  473. gs_texrender_reset(render);
  474. gs_blend_state_push();
  475. gs_blend_function(GS_BLEND_ONE, GS_BLEND_ZERO);
  476. if (gs_texrender_begin(render, filter->width, filter->height)) {
  477. struct vec4 clear_color;
  478. vec4_zero(&clear_color);
  479. gs_clear(GS_CLEAR_COLOR, &clear_color, 0.0f, 0);
  480. gs_ortho(0.0f, (float)filter->width, 0.0f,
  481. (float)filter->height, -100.0f, 100.0f);
  482. if (target == parent && !custom_draw && !async)
  483. obs_source_default_render(target);
  484. else
  485. obs_source_video_render(target);
  486. gs_texrender_end(render);
  487. }
  488. gs_blend_state_pop();
  489. /* 2. Initialize src_texture (only at startup or reset) */
  490. if (!filter->initial_render) {
  491. obs_enter_graphics();
  492. filter->src_texture = gs_texrender_get_texture(filter->render);
  493. struct ID3D11Texture2D *d11texture2 =
  494. (struct ID3D11Texture2D *)gs_texture_get_obj(
  495. filter->src_texture);
  496. obs_leave_graphics();
  497. if (!d11texture2) {
  498. error("Couldn't retrieve d3d11texture2d.");
  499. return;
  500. }
  501. vfxErr = NvCVImage_Create(filter->width, filter->height,
  502. NVCV_BGRA, NVCV_U8, NVCV_CHUNKY,
  503. NVCV_GPU, 1, &filter->src_img);
  504. if (vfxErr != NVCV_SUCCESS) {
  505. const char *errString =
  506. NvCV_GetErrorStringFromCode(vfxErr);
  507. error("Error creating src img; error %i: %s", vfxErr,
  508. errString);
  509. filter->processing_stop = true;
  510. return;
  511. }
  512. vfxErr = NvCVImage_InitFromD3D11Texture(filter->src_img,
  513. d11texture2);
  514. if (vfxErr != NVCV_SUCCESS) {
  515. const char *errString =
  516. NvCV_GetErrorStringFromCode(vfxErr);
  517. error("Error passing src ID3D11Texture to img; error %i: %s",
  518. vfxErr, errString);
  519. filter->processing_stop = true;
  520. return;
  521. }
  522. filter->initial_render = true;
  523. }
  524. /* 3. Process FX (outputs a mask) & draw. */
  525. if (filter->initial_render && filter->images_allocated) {
  526. ret = true;
  527. if (filter->got_new_frame) {
  528. ret = process_texture_greenscreen(filter);
  529. os_atomic_set_bool(&filter->got_new_frame, false);
  530. }
  531. if (ret) {
  532. if (!srgb_draw)
  533. draw_greenscreen(filter);
  534. else
  535. draw_greenscreen_srgb(filter);
  536. filter->processed_frame = true;
  537. }
  538. } else {
  539. obs_source_skip_video_filter(filter->context);
  540. return;
  541. }
  542. UNUSED_PARAMETER(effect);
  543. }
  544. bool load_nvvfx(void)
  545. {
  546. if (!load_nv_vfx_libs()) {
  547. blog(LOG_INFO,
  548. "[NVIDIA RTX VIDEO FX]: FX disabled, redistributable not found.");
  549. return false;
  550. }
  551. #define LOAD_SYM_FROM_LIB(sym, lib, dll) \
  552. if (!(sym = (sym##_t)GetProcAddress(lib, #sym))) { \
  553. DWORD err = GetLastError(); \
  554. printf("[NVIDIA RTX VIDEO FX]: Couldn't load " #sym \
  555. " from " dll ": %lu (0x%lx)", \
  556. err, err); \
  557. release_nv_vfx(); \
  558. goto unload_everything; \
  559. }
  560. #define LOAD_SYM(sym) LOAD_SYM_FROM_LIB(sym, nv_videofx, "NVVideoEffects.dll")
  561. LOAD_SYM(NvVFX_GetVersion);
  562. LOAD_SYM(NvVFX_CreateEffect);
  563. LOAD_SYM(NvVFX_DestroyEffect);
  564. LOAD_SYM(NvVFX_SetU32);
  565. LOAD_SYM(NvVFX_SetS32);
  566. LOAD_SYM(NvVFX_SetF32);
  567. LOAD_SYM(NvVFX_SetF64);
  568. LOAD_SYM(NvVFX_SetU64);
  569. LOAD_SYM(NvVFX_SetObject);
  570. LOAD_SYM(NvVFX_SetCudaStream);
  571. LOAD_SYM(NvVFX_SetImage);
  572. LOAD_SYM(NvVFX_SetString);
  573. LOAD_SYM(NvVFX_GetU32);
  574. LOAD_SYM(NvVFX_GetS32);
  575. LOAD_SYM(NvVFX_GetF32);
  576. LOAD_SYM(NvVFX_GetF64);
  577. LOAD_SYM(NvVFX_GetU64);
  578. LOAD_SYM(NvVFX_GetObject);
  579. LOAD_SYM(NvVFX_GetCudaStream);
  580. LOAD_SYM(NvVFX_GetImage);
  581. LOAD_SYM(NvVFX_GetString);
  582. LOAD_SYM(NvVFX_Run);
  583. LOAD_SYM(NvVFX_Load);
  584. LOAD_SYM(NvVFX_CudaStreamCreate);
  585. LOAD_SYM(NvVFX_CudaStreamDestroy);
  586. #undef LOAD_SYM
  587. #define LOAD_SYM(sym) LOAD_SYM_FROM_LIB(sym, nv_cvimage, "NVCVImage.dll")
  588. LOAD_SYM(NvCV_GetErrorStringFromCode);
  589. LOAD_SYM(NvCVImage_Init);
  590. LOAD_SYM(NvCVImage_InitView);
  591. LOAD_SYM(NvCVImage_Alloc);
  592. LOAD_SYM(NvCVImage_Realloc);
  593. LOAD_SYM(NvCVImage_Dealloc);
  594. LOAD_SYM(NvCVImage_Create);
  595. LOAD_SYM(NvCVImage_Destroy);
  596. LOAD_SYM(NvCVImage_ComponentOffsets);
  597. LOAD_SYM(NvCVImage_Transfer);
  598. LOAD_SYM(NvCVImage_TransferRect);
  599. LOAD_SYM(NvCVImage_TransferFromYUV);
  600. LOAD_SYM(NvCVImage_TransferToYUV);
  601. LOAD_SYM(NvCVImage_MapResource);
  602. LOAD_SYM(NvCVImage_UnmapResource);
  603. LOAD_SYM(NvCVImage_Composite);
  604. LOAD_SYM(NvCVImage_CompositeRect);
  605. LOAD_SYM(NvCVImage_CompositeOverConstant);
  606. LOAD_SYM(NvCVImage_FlipY);
  607. LOAD_SYM(NvCVImage_GetYUVPointers);
  608. LOAD_SYM(NvCVImage_InitFromD3D11Texture);
  609. LOAD_SYM(NvCVImage_ToD3DFormat);
  610. LOAD_SYM(NvCVImage_FromD3DFormat);
  611. LOAD_SYM(NvCVImage_ToD3DColorSpace);
  612. LOAD_SYM(NvCVImage_FromD3DColorSpace);
  613. #undef LOAD_SYM
  614. #define LOAD_SYM(sym) LOAD_SYM_FROM_LIB(sym, nv_cudart, "cudart64_110.dll")
  615. LOAD_SYM(cudaMalloc);
  616. LOAD_SYM(cudaStreamSynchronize);
  617. LOAD_SYM(cudaFree);
  618. LOAD_SYM(cudaMemcpy);
  619. LOAD_SYM(cudaMemsetAsync);
  620. #undef LOAD_SYM
  621. int err;
  622. NvVFX_Handle h = NULL;
  623. /* load the effect to check if the GPU is supported */
  624. err = NvVFX_CreateEffect(NVVFX_FX_GREEN_SCREEN, &h);
  625. if (err != NVCV_SUCCESS) {
  626. if (err == NVCV_ERR_UNSUPPORTEDGPU) {
  627. blog(LOG_INFO,
  628. "[NVIDIA RTX VIDEO FX]: disabled, unsupported GPU");
  629. } else {
  630. blog(LOG_ERROR,
  631. "[NVIDIA RTX VIDEO FX]: disabled, error %i", err);
  632. }
  633. goto unload_everything;
  634. }
  635. NvVFX_DestroyEffect(h);
  636. nvvfx_loaded = true;
  637. blog(LOG_INFO, "[NVIDIA RTX VIDEO FX]: enabled, redistributable found");
  638. return true;
  639. unload_everything:
  640. nvvfx_loaded = false;
  641. release_nv_vfx();
  642. return false;
  643. }
  644. #ifdef LIBNVVFX_ENABLED
  645. void unload_nvvfx(void)
  646. {
  647. release_nv_vfx();
  648. }
  649. #endif
  650. struct obs_source_info nvidia_greenscreen_filter_info = {
  651. .id = "nv_greenscreen_filter",
  652. .type = OBS_SOURCE_TYPE_FILTER,
  653. .output_flags = OBS_SOURCE_VIDEO,
  654. .get_name = nv_greenscreen_filter_name,
  655. .create = nv_greenscreen_filter_create,
  656. .destroy = nv_greenscreen_filter_destroy,
  657. .get_defaults = nv_greenscreen_filter_defaults,
  658. .get_properties = nv_greenscreen_filter_properties,
  659. .update = nv_greenscreen_filter_update,
  660. .filter_video = nv_greenscreen_filter_video,
  661. .video_render = nv_greenscreen_filter_render,
  662. .video_tick = nv_greenscreen_filter_tick,
  663. };