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

Showing the changes to libavfilter (6 of 7 files). Show all

Unified · Split

6 files
  1. libavfilter/dnn/Makefile +1 −0
  2. libavfilter/dnn/dnn_backend_torch.cpp +597 −0
  3. libavfilter/dnn/dnn_interface.c +5 −0
  4. libavfilter/dnn_filter_common.c +13 −2
  5. libavfilter/dnn_interface.h +1 −1
  6. libavfilter/vf_dnn_processing.c +3 −0
libavfilter/dnn/Makefile+1−0
@@ -6,5 +6,6 @@ OBJS-$(CONFIG_DNN) += dnn/dnn_backend_common.o · 5 unmodified lines
66
77 DNN-OBJS-$(CONFIG_LIBTENSORFLOW) += dnn/dnn_backend_tf.o
88 DNN-OBJS-$(CONFIG_LIBOPENVINO) += dnn/dnn_backend_openvino.o
9+DNN-OBJS-$(CONFIG_LIBTORCH) += dnn/dnn_backend_torch.o
910
1011 OBJS-$(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+
29+extern "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+
38+typedef struct THOptions{
39+ char *device_name;
40+ int optimize;
41+} THOptions;
42+
43+typedef struct THContext {
44+ const AVClass *c_class;
45+ THOptions options;
46+} THContext;
47+
48+typedef 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+
57+typedef struct THInferRequest {
58+ torch::Tensor *output;
59+ torch::Tensor *input_tensor;
60+} THInferRequest;
61+
62+typedef 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
71+static 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+
77+AVFILTER_DEFINE_CLASS(dnn_th);
78+
79+static 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+
99+static 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+
114+static 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+
128+static 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+
160+static 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+
172+static void deleter(void *arg)
173+{
174+ av_freep(&arg);
175+}
176+
177+static 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+
231+err:
232+ th_free_request(infer_request);
233+ return ret;
234+}
235+
236+static 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+
273+static 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);
319+err:
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+
328+static 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+
364+err:
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+
372+static 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+
409+err:
410+ av_frame_free(&task.out_frame);
411+ av_frame_free(&task.in_frame);
412+ return ret;
413+}
414+
415+static 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+
426+static 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+
509+fail:
510+ if (item) {
511+ destroy_request_item(&item);
512+ av_freep(&item);
513+ }
514+ dnn_free_model_th(&model);
515+ return NULL;
516+}
517+
518+static 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+
567+static 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+
573+static 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+
591+extern 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
2828
2929 extern const DNNModule ff_dnn_backend_openvino;
3030 extern const DNNModule ff_dnn_backend_tf;
31+extern const DNNModule ff_dnn_backend_torch;
3132
3233 const DNNModule *ff_get_dnn_module(DNNBackendType backend_type, void *log_ctx)
3334 {
@@ -40,6 +41,10 @@ const DNNModule *ff_get_dnn_module(DNNBackendType backend_type, void *log_ctx) · 6 unmodified lines
4041 case DNN_OV:
4142 return &ff_dnn_backend_openvino;
4243 #endif
44+ #if (CONFIG_LIBTORCH == 1)
45+ case DNN_TH:
46+ return &ff_dnn_backend_torch;
47+ #endif
4348 default:
4449 av_log(log_ctx, AV_LOG_ERROR,
4550 "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
5353
5454 int ff_dnn_init(DnnContext *ctx, DNNFunctionType func_type, AVFilterContext *filter_ctx)
5555 {
56+ DNNBackendType backend = ctx->backend_type;
57+
5658 if (!ctx->model_filename) {
5759 av_log(filter_ctx, AV_LOG_ERROR, "model file for network is not specified\n");
5860 return AVERROR(EINVAL);
5961 }
6062
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) {
6272 if (!ctx->model_inputname) {
6373 av_log(filter_ctx, AV_LOG_ERROR, "input name of the model network is not specified\n");
6474 return AVERROR(EINVAL);
@@ -115,7 +125,8 @@ int ff_dnn_get_input(DnnContext *ctx, DNNData *input) · 50 unmodified lines
115125
116126 int ff_dnn_get_output(DnnContext *ctx, int input_width, int input_height, int *output_width, int *output_height)
117127 {
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;
119130 return ctx->model->get_output(ctx->model->model, ctx->model_inputname, input_width, input_height,
120131 (const char *)output_name, output_width, output_height);
121132 }
libavfilter/dnn_interface.h+1−1
@@ -32,7 +32,7 @@ · 31 unmodified lines
3232
3333 #define DNN_GENERIC_ERROR FFERRTAG('D','N','N','!')
3434
35−typedef enum {DNN_TF = 1, DNN_OV} DNNBackendType;
35+typedef enum {DNN_TF = 1, DNN_OV, DNN_TH} DNNBackendType;
3636
3737 typedef enum {DNN_FLOAT = 1, DNN_UINT8 = 4} DNNDataType;
3838
libavfilter/vf_dnn_processing.c+3−0
@@ -50,6 +50,9 @@ static const AVOption dnn_processing_options[] = { · 49 unmodified lines
5050 #endif
5151 #if (CONFIG_LIBOPENVINO == 1)
5252 { "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" },
5356 #endif
5457 DNN_COMMON_OPTIONS
5558 { 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