History

libavfi/dnn: add LibTorch as one of DNN backend

PyTorch is an open source machine learning framework that accelerates the path from research prototyping to production deployment. Official website: https://pytorch.org/. We call the C++ library of PyTorch as LibTorch, the same below. To build FFmpeg with LibTorch, please take following steps as reference: 1. download LibTorch C++ library in https://pytorch.org/get-started/locally/, please select C++/Java for language, and other options as your need. Please download cxx11 ABI version: (libtorch-cxx11-abi-shared-with-deps-*.zip). 2. unzip the file to your own dir, with command unzip libtorch-shared-with-deps-latest.zip -d your_dir 3. export libtorch_root/libtorch/include and libtorch_root/libtorch/include/torch/csrc/api/include to $PATH export libtorch_root/libtorch/lib/ to $LD_LIBRARY_PATH 4. config FFmpeg with ../configure --enable-libtorch \ --extra-cflag=-I/libtorch_root/libtorch/include \ --extra-cflag=-I/libtorch_root/libtorch/include/torch/csrc/api/include \ --extra-ldflags=-L/libtorch_root/libtorch/lib/ 5. make To run FFmpeg DNN inference with LibTorch backend: ./ffmpeg -i input.jpg -vf \ dnn_processing=dnn_backend=torch:model=LibTorch_model.pt -y output.jpg The LibTorch_model.pt can be generated by Python with torch.jit.script() api. https://pytorch.org/tutorials/advanced/cpp_export.html. This is pytorch official guide about how to convert and load torchscript model. Please note, torch.jit.trace() is not recommanded, since it does not support ambiguous input size.

Wenbin Chen committed 2 years ago

Commit f4e06641 parent d24b1367 files · +624 −4Browse files

Unified · Split

7 files
  1. configure +4 −1
  2. libavfilter/dnn/Makefile +1 −0
  3. libavfilter/dnn/dnn_backend_torch.cpp +597 −0
  4. libavfilter/dnn/dnn_interface.c +5 −0
  5. libavfilter/dnn_filter_common.c +13 −2
  6. libavfilter/dnn_interface.h +1 −1
  7. libavfilter/vf_dnn_processing.c +3 −0
configure+4−1
@@ -281,6 +281,7 @@ External library support: · 280 unmodified lines
281 --enable-libtheora enable Theora encoding via libtheora [no]281 --enable-libtheora enable Theora encoding via libtheora [no]
282 --enable-libtls enable LibreSSL (via libtls), needed for https support282 --enable-libtls enable LibreSSL (via libtls), needed for https support
283 if openssl, gnutls or mbedtls is not used [no]283 if openssl, gnutls or mbedtls is not used [no]
284 --enable-libtorch enable Torch as one DNN backend [no]
284 --enable-libtwolame enable MP2 encoding via libtwolame [no]285 --enable-libtwolame enable MP2 encoding via libtwolame [no]
285 --enable-libuavs3d enable AVS3 decoding via libuavs3d [no]286 --enable-libuavs3d enable AVS3 decoding via libuavs3d [no]
286 --enable-libv4l2 enable libv4l2/v4l-utils [no]287 --enable-libv4l2 enable libv4l2/v4l-utils [no]
@@ -1905,6 +1906,7 @@ EXTERNAL_LIBRARY_LIST=" · 1618 unmodified lines
1905 libtensorflow1906 libtensorflow
1906 libtesseract1907 libtesseract
1907 libtheora1908 libtheora
1909 libtorch
1908 libtwolame1910 libtwolame
1909 libuavs3d1911 libuavs3d
1910 libv4l21912 libv4l2
@@ -2785,7 +2787,7 @@ cbs_vp9_select="cbs" · 874 unmodified lines
2785deflate_wrapper_deps="zlib"2787deflate_wrapper_deps="zlib"
2786dirac_parse_select="golomb"2788dirac_parse_select="golomb"
2787dovi_rpu_select="golomb"2789dovi_rpu_select="golomb"
2788dnn_suggest="libtensorflow libopenvino"2790dnn_suggest="libtensorflow libopenvino libtorch"
2789dnn_deps="avformat swscale"2791dnn_deps="avformat swscale"
2790error_resilience_select="me_cmp"2792error_resilience_select="me_cmp"
2791evcparse_select="golomb"2793evcparse_select="golomb"
@@ -6884,6 +6886,7 @@ enabled libtensorflow && require libtensorflow tensorflow/c/c_api.h TF_Versi · 4092 unmodified lines
6884enabled libtesseract && require_pkg_config libtesseract tesseract tesseract/capi.h TessBaseAPICreate6886enabled libtesseract && require_pkg_config libtesseract tesseract tesseract/capi.h TessBaseAPICreate
6885enabled libtheora && require libtheora theora/theoraenc.h th_info_init -ltheoraenc -ltheoradec -logg6887enabled libtheora && require libtheora theora/theoraenc.h th_info_init -ltheoraenc -ltheoradec -logg
6886enabled libtls && require_pkg_config libtls libtls tls.h tls_configure6888enabled libtls && require_pkg_config libtls libtls tls.h tls_configure
6889enabled libtorch && check_cxxflags -std=c++17 && require_cpp libtorch torch/torch.h "torch::Tensor" -ltorch -lc10 -ltorch_cpu -lstdc++ -lpthread
6887enabled libtwolame && require libtwolame twolame.h twolame_init -ltwolame &&6890enabled libtwolame && require libtwolame twolame.h twolame_init -ltwolame &&
6888 { check_lib libtwolame twolame.h twolame_encode_buffer_float32_interleaved -ltwolame ||6891 { check_lib libtwolame twolame.h twolame_encode_buffer_float32_interleaved -ltwolame ||
6889 die "ERROR: libtwolame must be installed and version must be >= 0.3.10"; }6892 die "ERROR: libtwolame must be installed and version must be >= 0.3.10"; }
libavfilter/dnn/Makefile+1−0
@@ -6,5 +6,6 @@ OBJS-$(CONFIG_DNN) += dnn/dnn_backend_common.o · 5 unmodified lines
6 6
7DNN-OBJS-$(CONFIG_LIBTENSORFLOW) += dnn/dnn_backend_tf.o7DNN-OBJS-$(CONFIG_LIBTENSORFLOW) += dnn/dnn_backend_tf.o
8DNN-OBJS-$(CONFIG_LIBOPENVINO) += dnn/dnn_backend_openvino.o8DNN-OBJS-$(CONFIG_LIBOPENVINO) += dnn/dnn_backend_openvino.o
9DNN-OBJS-$(CONFIG_LIBTORCH) += dnn/dnn_backend_torch.o
9 10
10OBJS-$(CONFIG_DNN) += $(DNN-OBJS-yes)11OBJS-$(CONFIG_DNN) += $(DNN-OBJS-yes)
libavfilter/dnn/dnn_backend_torch.cpp+597−0
@@ -0,0 +1,597 @@
1/*
2 * Copyright (c) 2024
3 *
4 * This file is part of FFmpeg.
5 *
6 * FFmpeg is free software; you can redistribute it and/or
7 * modify it under the terms of the GNU Lesser General Public
8 * License as published by the Free Software Foundation; either
9 * version 2.1 of the License, or (at your option) any later version.
10 *
11 * FFmpeg is distributed in the hope that it will be useful,
12 * but WITHOUT ANY WARRANTY; without even the implied warranty of
13 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
14 * Lesser General Public License for more details.
15 *
16 * You should have received a copy of the GNU Lesser General Public
17 * License along with FFmpeg; if not, write to the Free Software
18 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
19 */
20
21/**
22 * @file
23 * DNN Torch backend implementation.
24 */
25
26#include <torch/torch.h>
27#include <torch/script.h>
28
29extern "C" {
30#include "../internal.h"
31#include "dnn_io_proc.h"
32#include "dnn_backend_common.h"
33#include "libavutil/opt.h"
34#include "queue.h"
35#include "safe_queue.h"
36}
37
38typedef struct THOptions{
39 char *device_name;
40 int optimize;
41} THOptions;
42
43typedef struct THContext {
44 const AVClass *c_class;
45 THOptions options;
46} THContext;
47
48typedef struct THModel {
49 THContext ctx;
50 DNNModel *model;
51 torch::jit::Module *jit_model;
52 SafeQueue *request_queue;
53 Queue *task_queue;
54 Queue *lltask_queue;
55} THModel;
56
57typedef struct THInferRequest {
58 torch::Tensor *output;
59 torch::Tensor *input_tensor;
60} THInferRequest;
61
62typedef struct THRequestItem {
63 THInferRequest *infer_request;
64 LastLevelTaskItem *lltask;
65 DNNAsyncExecModule exec_module;
66} THRequestItem;
67
68
69#define OFFSET(x) offsetof(THContext, x)
70#define FLAGS AV_OPT_FLAG_FILTERING_PARAM
71static const AVOption dnn_th_options[] = {
72 { "device", "device to run model", OFFSET(options.device_name), AV_OPT_TYPE_STRING, { .str = "cpu" }, 0, 0, FLAGS },
73 { "optimize", "turn on graph executor optimization", OFFSET(options.optimize), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 1, FLAGS},
74 { NULL }
75};
76
77AVFILTER_DEFINE_CLASS(dnn_th);
78
79static int extract_lltask_from_task(TaskItem *task, Queue *lltask_queue)
80{
81 THModel *th_model = (THModel *)task->model;
82 THContext *ctx = &th_model->ctx;
83 LastLevelTaskItem *lltask = (LastLevelTaskItem *)av_malloc(sizeof(*lltask));
84 if (!lltask) {
85 av_log(ctx, AV_LOG_ERROR, "Failed to allocate memory for LastLevelTaskItem\n");
86 return AVERROR(ENOMEM);
87 }
88 task->inference_todo = 1;
89 task->inference_done = 0;
90 lltask->task = task;
91 if (ff_queue_push_back(lltask_queue, lltask) < 0) {
92 av_log(ctx, AV_LOG_ERROR, "Failed to push back lltask_queue.\n");
93 av_freep(&lltask);
94 return AVERROR(ENOMEM);
95 }
96 return 0;
97}
98
99static void th_free_request(THInferRequest *request)
100{
101 if (!request)
102 return;
103 if (request->output) {
104 delete(request->output);
105 request->output = NULL;
106 }
107 if (request->input_tensor) {
108 delete(request->input_tensor);
109 request->input_tensor = NULL;
110 }
111 return;
112}
113
114static inline void destroy_request_item(THRequestItem **arg)
115{
116 THRequestItem *item;
117 if (!arg || !*arg) {
118 return;
119 }
120 item = *arg;
121 th_free_request(item->infer_request);
122 av_freep(&item->infer_request);
123 av_freep(&item->lltask);
124 ff_dnn_async_module_cleanup(&item->exec_module);
125 av_freep(arg);
126}
127
128static void dnn_free_model_th(DNNModel **model)
129{
130 THModel *th_model;
131 if (!model || !*model)
132 return;
133
134 th_model = (THModel *) (*model)->model;
135 while (ff_safe_queue_size(th_model->request_queue) != 0) {
136 THRequestItem *item = (THRequestItem *)ff_safe_queue_pop_front(th_model->request_queue);
137 destroy_request_item(&item);
138 }
139 ff_safe_queue_destroy(th_model->request_queue);
140
141 while (ff_queue_size(th_model->lltask_queue) != 0) {
142 LastLevelTaskItem *item = (LastLevelTaskItem *)ff_queue_pop_front(th_model->lltask_queue);
143 av_freep(&item);
144 }
145 ff_queue_destroy(th_model->lltask_queue);
146
147 while (ff_queue_size(th_model->task_queue) != 0) {
148 TaskItem *item = (TaskItem *)ff_queue_pop_front(th_model->task_queue);
149 av_frame_free(&item->in_frame);
150 av_frame_free(&item->out_frame);
151 av_freep(&item);
152 }
153 ff_queue_destroy(th_model->task_queue);
154 delete th_model->jit_model;
155 av_opt_free(&th_model->ctx);
156 av_freep(&th_model);
157 av_freep(model);
158}
159
160static int get_input_th(void *model, DNNData *input, const char *input_name)
161{
162 input->dt = DNN_FLOAT;
163 input->order = DCO_RGB;
164 input->layout = DL_NCHW;
165 input->dims[0] = 1;
166 input->dims[1] = 3;
167 input->dims[2] = -1;
168 input->dims[3] = -1;
169 return 0;
170}
171
172static void deleter(void *arg)
173{
174 av_freep(&arg);
175}
176
177static int fill_model_input_th(THModel *th_model, THRequestItem *request)
178{
179 LastLevelTaskItem *lltask = NULL;
180 TaskItem *task = NULL;
181 THInferRequest *infer_request = NULL;
182 DNNData input = { 0 };
183 THContext *ctx = &th_model->ctx;
184 int ret, width_idx, height_idx, channel_idx;
185
186 lltask = (LastLevelTaskItem *)ff_queue_pop_front(th_model->lltask_queue);
187 if (!lltask) {
188 ret = AVERROR(EINVAL);
189 goto err;
190 }
191 request->lltask = lltask;
192 task = lltask->task;
193 infer_request = request->infer_request;
194
195 ret = get_input_th(th_model, &input, NULL);
196 if ( ret != 0) {
197 goto err;
198 }
199 width_idx = dnn_get_width_idx_by_layout(input.layout);
200 height_idx = dnn_get_height_idx_by_layout(input.layout);
201 channel_idx = dnn_get_channel_idx_by_layout(input.layout);
202 input.dims[height_idx] = task->in_frame->height;
203 input.dims[width_idx] = task->in_frame->width;
204 input.data = av_malloc(input.dims[height_idx] * input.dims[width_idx] *
205 input.dims[channel_idx] * sizeof(float));
206 if (!input.data)
207 return AVERROR(ENOMEM);
208 infer_request->input_tensor = new torch::Tensor();
209 infer_request->output = new torch::Tensor();
210
211 switch (th_model->model->func_type) {
212 case DFT_PROCESS_FRAME:
213 input.scale = 255;
214 if (task->do_ioproc) {
215 if (th_model->model->frame_pre_proc != NULL) {
216 th_model->model->frame_pre_proc(task->in_frame, &input, th_model->model->filter_ctx);
217 } else {
218 ff_proc_from_frame_to_dnn(task->in_frame, &input, ctx);
219 }
220 }
221 break;
222 default:
223 avpriv_report_missing_feature(NULL, "model function type %d", th_model->model->func_type);
224 break;
225 }
226 *infer_request->input_tensor = torch::from_blob(input.data,
227 {1, input.dims[channel_idx], input.dims[height_idx], input.dims[width_idx]},
228 deleter, torch::kFloat32);
229 return 0;
230
231err:
232 th_free_request(infer_request);
233 return ret;
234}
235
236static int th_start_inference(void *args)
237{
238 THRequestItem *request = (THRequestItem *)args;
239 THInferRequest *infer_request = NULL;
240 LastLevelTaskItem *lltask = NULL;
241 TaskItem *task = NULL;
242 THModel *th_model = NULL;
243 THContext *ctx = NULL;
244 std::vector<torch::jit::IValue> inputs;
245 torch::NoGradGuard no_grad;
246
247 if (!request) {
248 av_log(NULL, AV_LOG_ERROR, "THRequestItem is NULL\n");
249 return AVERROR(EINVAL);
250 }
251 infer_request = request->infer_request;
252 lltask = request->lltask;
253 task = lltask->task;
254 th_model = (THModel *)task->model;
255 ctx = &th_model->ctx;
256
257 if (ctx->options.optimize)
258 torch::jit::setGraphExecutorOptimize(true);
259 else
260 torch::jit::setGraphExecutorOptimize(false);
261
262 if (!infer_request->input_tensor || !infer_request->output) {
263 av_log(ctx, AV_LOG_ERROR, "input or output tensor is NULL\n");
264 return DNN_GENERIC_ERROR;
265 }
266 inputs.push_back(*infer_request->input_tensor);
267
268 *infer_request->output = th_model->jit_model->forward(inputs).toTensor();
269
270 return 0;
271}
272
273static void infer_completion_callback(void *args) {
274 THRequestItem *request = (THRequestItem*)args;
275 LastLevelTaskItem *lltask = request->lltask;
276 TaskItem *task = lltask->task;
277 DNNData outputs = { 0 };
278 THInferRequest *infer_request = request->infer_request;
279 THModel *th_model = (THModel *)task->model;
280 torch::Tensor *output = infer_request->output;
281
282 c10::IntArrayRef sizes = output->sizes();
283 outputs.order = DCO_RGB;
284 outputs.layout = DL_NCHW;
285 outputs.dt = DNN_FLOAT;
286 if (sizes.size() == 4) {
287 // 4 dimensions: [batch_size, channel, height, width]
288 // this format of data is normally used for video frame SR
289 outputs.dims[0] = sizes.at(0); // N
290 outputs.dims[1] = sizes.at(1); // C
291 outputs.dims[2] = sizes.at(2); // H
292 outputs.dims[3] = sizes.at(3); // W
293 } else {
294 avpriv_report_missing_feature(&th_model->ctx, "Support of this kind of model");
295 goto err;
296 }
297
298 switch (th_model->model->func_type) {
299 case DFT_PROCESS_FRAME:
300 if (task->do_ioproc) {
301 outputs.scale = 255;
302 outputs.data = output->data_ptr();
303 if (th_model->model->frame_post_proc != NULL) {
304 th_model->model->frame_post_proc(task->out_frame, &outputs, th_model->model->filter_ctx);
305 } else {
306 ff_proc_from_dnn_to_frame(task->out_frame, &outputs, &th_model->ctx);
307 }
308 } else {
309 task->out_frame->width = outputs.dims[dnn_get_width_idx_by_layout(outputs.layout)];
310 task->out_frame->height = outputs.dims[dnn_get_height_idx_by_layout(outputs.layout)];
311 }
312 break;
313 default:
314 avpriv_report_missing_feature(&th_model->ctx, "model function type %d", th_model->model->func_type);
315 goto err;
316 }
317 task->inference_done++;
318 av_freep(&request->lltask);
319err:
320 th_free_request(infer_request);
321
322 if (ff_safe_queue_push_back(th_model->request_queue, request) < 0) {
323 destroy_request_item(&request);
324 av_log(&th_model->ctx, AV_LOG_ERROR, "Unable to push back request_queue when failed to start inference.\n");
325 }
326}
327
328static int execute_model_th(THRequestItem *request, Queue *lltask_queue)
329{
330 THModel *th_model = NULL;
331 LastLevelTaskItem *lltask;
332 TaskItem *task = NULL;
333 int ret = 0;
334
335 if (ff_queue_size(lltask_queue) == 0) {
336 destroy_request_item(&request);
337 return 0;
338 }
339
340 lltask = (LastLevelTaskItem *)ff_queue_peek_front(lltask_queue);
341 if (lltask == NULL) {
342 av_log(NULL, AV_LOG_ERROR, "Failed to get LastLevelTaskItem\n");
343 ret = AVERROR(EINVAL);
344 goto err;
345 }
346 task = lltask->task;
347 th_model = (THModel *)task->model;
348
349 ret = fill_model_input_th(th_model, request);
350 if ( ret != 0) {
351 goto err;
352 }
353 if (task->async) {
354 avpriv_report_missing_feature(&th_model->ctx, "LibTorch async");
355 } else {
356 ret = th_start_inference((void *)(request));
357 if (ret != 0) {
358 goto err;
359 }
360 infer_completion_callback(request);
361 return (task->inference_done == task->inference_todo) ? 0 : DNN_GENERIC_ERROR;
362 }
363
364err:
365 th_free_request(request->infer_request);
366 if (ff_safe_queue_push_back(th_model->request_queue, request) < 0) {
367 destroy_request_item(&request);
368 }
369 return ret;
370}
371
372static int get_output_th(void *model, const char *input_name, int input_width, int input_height,
373 const char *output_name, int *output_width, int *output_height)
374{
375 int ret = 0;
376 THModel *th_model = (THModel*) model;
377 THContext *ctx = &th_model->ctx;
378 TaskItem task = { 0 };
379 THRequestItem *request = NULL;
380 DNNExecBaseParams exec_params = {
381 .input_name = input_name,
382 .output_names = &output_name,
383 .nb_output = 1,
384 .in_frame = NULL,
385 .out_frame = NULL,
386 };
387 ret = ff_dnn_fill_gettingoutput_task(&task, &exec_params, th_model, input_height, input_width, ctx);
388 if ( ret != 0) {
389 goto err;
390 }
391
392 ret = extract_lltask_from_task(&task, th_model->lltask_queue);
393 if ( ret != 0) {
394 av_log(ctx, AV_LOG_ERROR, "unable to extract last level task from task.\n");
395 goto err;
396 }
397
398 request = (THRequestItem*) ff_safe_queue_pop_front(th_model->request_queue);
399 if (!request) {
400 av_log(ctx, AV_LOG_ERROR, "unable to get infer request.\n");
401 ret = AVERROR(EINVAL);
402 goto err;
403 }
404
405 ret = execute_model_th(request, th_model->lltask_queue);
406 *output_width = task.out_frame->width;
407 *output_height = task.out_frame->height;
408
409err:
410 av_frame_free(&task.out_frame);
411 av_frame_free(&task.in_frame);
412 return ret;
413}
414
415static THInferRequest *th_create_inference_request(void)
416{
417 THInferRequest *request = (THInferRequest *)av_malloc(sizeof(THInferRequest));
418 if (!request) {
419 return NULL;
420 }
421 request->input_tensor = NULL;
422 request->output = NULL;
423 return request;
424}
425
426static DNNModel *dnn_load_model_th(const char *model_filename, DNNFunctionType func_type, const char *options, AVFilterContext *filter_ctx)
427{
428 DNNModel *model = NULL;
429 THModel *th_model = NULL;
430 THRequestItem *item = NULL;
431 THContext *ctx;
432
433 model = (DNNModel *)av_mallocz(sizeof(DNNModel));
434 if (!model) {
435 return NULL;
436 }
437
438 th_model = (THModel *)av_mallocz(sizeof(THModel));
439 if (!th_model) {
440 av_freep(&model);
441 return NULL;
442 }
443 th_model->model = model;
444 model->model = th_model;
445 th_model->ctx.c_class = &dnn_th_class;
446 ctx = &th_model->ctx;
447 //parse options
448 av_opt_set_defaults(ctx);
449 if (av_opt_set_from_string(ctx, options, NULL, "=", "&") < 0) {
450 av_log(ctx, AV_LOG_ERROR, "Failed to parse options \"%s\"\n", options);
451 return NULL;
452 }
453
454 c10::Device device = c10::Device(ctx->options.device_name);
455 if (!device.is_cpu()) {
456 av_log(ctx, AV_LOG_ERROR, "Not supported device:\"%s\"\n", ctx->options.device_name);
457 goto fail;
458 }
459
460 try {
461 th_model->jit_model = new torch::jit::Module;
462 (*th_model->jit_model) = torch::jit::load(model_filename);
463 } catch (const c10::Error& e) {
464 av_log(ctx, AV_LOG_ERROR, "Failed to load torch model\n");
465 goto fail;
466 }
467
468 th_model->request_queue = ff_safe_queue_create();
469 if (!th_model->request_queue) {
470 goto fail;
471 }
472
473 item = (THRequestItem *)av_mallocz(sizeof(THRequestItem));
474 if (!item) {
475 goto fail;
476 }
477 item->lltask = NULL;
478 item->infer_request = th_create_inference_request();
479 if (!item->infer_request) {
480 av_log(NULL, AV_LOG_ERROR, "Failed to allocate memory for Torch inference request\n");
481 goto fail;
482 }
483 item->exec_module.start_inference = &th_start_inference;
484 item->exec_module.callback = &infer_completion_callback;
485 item->exec_module.args = item;
486
487 if (ff_safe_queue_push_back(th_model->request_queue, item) < 0) {
488 goto fail;
489 }
490 item = NULL;
491
492 th_model->task_queue = ff_queue_create();
493 if (!th_model->task_queue) {
494 goto fail;
495 }
496
497 th_model->lltask_queue = ff_queue_create();
498 if (!th_model->lltask_queue) {
499 goto fail;
500 }
501
502 model->get_input = &get_input_th;
503 model->get_output = &get_output_th;
504 model->options = NULL;
505 model->filter_ctx = filter_ctx;
506 model->func_type = func_type;
507 return model;
508
509fail:
510 if (item) {
511 destroy_request_item(&item);
512 av_freep(&item);
513 }
514 dnn_free_model_th(&model);
515 return NULL;
516}
517
518static int dnn_execute_model_th(const DNNModel *model, DNNExecBaseParams *exec_params)
519{
520 THModel *th_model = (THModel *)model->model;
521 THContext *ctx = &th_model->ctx;
522 TaskItem *task;
523 THRequestItem *request;
524 int ret = 0;
525
526 ret = ff_check_exec_params(ctx, DNN_TH, model->func_type, exec_params);
527 if (ret != 0) {
528 av_log(ctx, AV_LOG_ERROR, "exec parameter checking fail.\n");
529 return ret;
530 }
531
532 task = (TaskItem *)av_malloc(sizeof(TaskItem));
533 if (!task) {
534 av_log(ctx, AV_LOG_ERROR, "unable to alloc memory for task item.\n");
535 return AVERROR(ENOMEM);
536 }
537
538 ret = ff_dnn_fill_task(task, exec_params, th_model, 0, 1);
539 if (ret != 0) {
540 av_freep(&task);
541 av_log(ctx, AV_LOG_ERROR, "unable to fill task.\n");
542 return ret;
543 }
544
545 ret = ff_queue_push_back(th_model->task_queue, task);
546 if (ret < 0) {
547 av_freep(&task);
548 av_log(ctx, AV_LOG_ERROR, "unable to push back task_queue.\n");
549 return ret;
550 }
551
552 ret = extract_lltask_from_task(task, th_model->lltask_queue);
553 if (ret != 0) {
554 av_log(ctx, AV_LOG_ERROR, "unable to extract last level task from task.\n");
555 return ret;
556 }
557
558 request = (THRequestItem *)ff_safe_queue_pop_front(th_model->request_queue);
559 if (!request) {
560 av_log(ctx, AV_LOG_ERROR, "unable to get infer request.\n");
561 return AVERROR(EINVAL);
562 }
563
564 return execute_model_th(request, th_model->lltask_queue);
565}
566
567static DNNAsyncStatusType dnn_get_result_th(const DNNModel *model, AVFrame **in, AVFrame **out)
568{
569 THModel *th_model = (THModel *)model->model;
570 return ff_dnn_get_result_common(th_model->task_queue, in, out);
571}
572
573static int dnn_flush_th(const DNNModel *model)
574{
575 THModel *th_model = (THModel *)model->model;
576 THRequestItem *request;
577
578 if (ff_queue_size(th_model->lltask_queue) == 0)
579 // no pending task need to flush
580 return 0;
581
582 request = (THRequestItem *)ff_safe_queue_pop_front(th_model->request_queue);
583 if (!request) {
584 av_log(&th_model->ctx, AV_LOG_ERROR, "unable to get infer request.\n");
585 return AVERROR(EINVAL);
586 }
587
588 return execute_model_th(request, th_model->lltask_queue);
589}
590
591extern const DNNModule ff_dnn_backend_torch = {
592 .load_model = dnn_load_model_th,
593 .execute_model = dnn_execute_model_th,
594 .get_result = dnn_get_result_th,
595 .flush = dnn_flush_th,
596 .free_model = dnn_free_model_th,
597};
libavfilter/dnn/dnn_interface.c+5−0
@@ -28,6 +28,7 @@ · 27 unmodified lines
28 28
29extern const DNNModule ff_dnn_backend_openvino;29extern const DNNModule ff_dnn_backend_openvino;
30extern const DNNModule ff_dnn_backend_tf;30extern const DNNModule ff_dnn_backend_tf;
31extern const DNNModule ff_dnn_backend_torch;
31 32
32const DNNModule *ff_get_dnn_module(DNNBackendType backend_type, void *log_ctx)33const DNNModule *ff_get_dnn_module(DNNBackendType backend_type, void *log_ctx)
33{34{
@@ -40,6 +41,10 @@ const DNNModule *ff_get_dnn_module(DNNBackendType backend_type, void *log_ctx) · 6 unmodified lines
40 case DNN_OV:41 case DNN_OV:
41 return &ff_dnn_backend_openvino;42 return &ff_dnn_backend_openvino;
42 #endif43 #endif
44 #if (CONFIG_LIBTORCH == 1)
45 case DNN_TH:
46 return &ff_dnn_backend_torch;
47 #endif
43 default:48 default:
44 av_log(log_ctx, AV_LOG_ERROR,49 av_log(log_ctx, AV_LOG_ERROR,
45 "Module backend_type %d is not supported or enabled.\n",50 "Module backend_type %d is not supported or enabled.\n",
libavfilter/dnn_filter_common.c+13−2
@@ -53,12 +53,22 @@ static char **separate_output_names(const char *expr, const char *val_sep, int * · 52 unmodified lines
53 53
54int ff_dnn_init(DnnContext *ctx, DNNFunctionType func_type, AVFilterContext *filter_ctx)54int ff_dnn_init(DnnContext *ctx, DNNFunctionType func_type, AVFilterContext *filter_ctx)
55{55{
56 DNNBackendType backend = ctx->backend_type;
57
56 if (!ctx->model_filename) {58 if (!ctx->model_filename) {
57 av_log(filter_ctx, AV_LOG_ERROR, "model file for network is not specified\n");59 av_log(filter_ctx, AV_LOG_ERROR, "model file for network is not specified\n");
58 return AVERROR(EINVAL);60 return AVERROR(EINVAL);
59 }61 }
60 62
61 if (ctx->backend_type == DNN_TF) {63 if (backend == DNN_TH) {
64 if (ctx->model_inputname)
65 av_log(filter_ctx, AV_LOG_WARNING, "LibTorch backend do not require inputname, "\
66 "inputname will be ignored.\n");
67 if (ctx->model_outputnames)
68 av_log(filter_ctx, AV_LOG_WARNING, "LibTorch backend do not require outputname(s), "\
69 "all outputname(s) will be ignored.\n");
70 ctx->nb_outputs = 1;
71 } else if (backend == DNN_TF) {
62 if (!ctx->model_inputname) {72 if (!ctx->model_inputname) {
63 av_log(filter_ctx, AV_LOG_ERROR, "input name of the model network is not specified\n");73 av_log(filter_ctx, AV_LOG_ERROR, "input name of the model network is not specified\n");
64 return AVERROR(EINVAL);74 return AVERROR(EINVAL);
@@ -115,7 +125,8 @@ int ff_dnn_get_input(DnnContext *ctx, DNNData *input) · 50 unmodified lines
115 125
116int ff_dnn_get_output(DnnContext *ctx, int input_width, int input_height, int *output_width, int *output_height)126int ff_dnn_get_output(DnnContext *ctx, int input_width, int input_height, int *output_width, int *output_height)
117{127{
118 char * output_name = ctx->model_outputnames ? ctx->model_outputnames[0] : NULL;128 char * output_name = ctx->model_outputnames && ctx->backend_type != DNN_TH ?
129 ctx->model_outputnames[0] : NULL;
119 return ctx->model->get_output(ctx->model->model, ctx->model_inputname, input_width, input_height,130 return ctx->model->get_output(ctx->model->model, ctx->model_inputname, input_width, input_height,
120 (const char *)output_name, output_width, output_height);131 (const char *)output_name, output_width, output_height);
121}132}
libavfilter/dnn_interface.h+1−1
@@ -32,7 +32,7 @@ · 31 unmodified lines
32 32
33#define DNN_GENERIC_ERROR FFERRTAG('D','N','N','!')33#define DNN_GENERIC_ERROR FFERRTAG('D','N','N','!')
34 34
35typedef enum {DNN_TF = 1, DNN_OV} DNNBackendType;35typedef enum {DNN_TF = 1, DNN_OV, DNN_TH} DNNBackendType;
36 36
37typedef enum {DNN_FLOAT = 1, DNN_UINT8 = 4} DNNDataType;37typedef enum {DNN_FLOAT = 1, DNN_UINT8 = 4} DNNDataType;
38 38
libavfilter/vf_dnn_processing.c+3−0
@@ -50,6 +50,9 @@ static const AVOption dnn_processing_options[] = { · 49 unmodified lines
50#endif50#endif
51#if (CONFIG_LIBOPENVINO == 1)51#if (CONFIG_LIBOPENVINO == 1)
52 { "openvino", "openvino backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = DNN_OV }, 0, 0, FLAGS, .unit = "backend" },52 { "openvino", "openvino backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = DNN_OV }, 0, 0, FLAGS, .unit = "backend" },
53#endif
54#if (CONFIG_LIBTORCH == 1)
55 { "torch", "torch backend flag", 0, AV_OPT_TYPE_CONST, { .i64 = DNN_TH }, 0, 0, FLAGS, "backend" },
53#endif56#endif
54 DNN_COMMON_OPTIONS57 DNN_COMMON_OPTIONS
55 { NULL }58 { NULL }

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t