git.sesse.net Git - ffmpeg/blob - libavfilter/vf_hwupload_cuda.c

   1 /*
   2  * This file is part of FFmpeg.
   3  *
   4  * FFmpeg is free software; you can redistribute it and/or
   5  * modify it under the terms of the GNU Lesser General Public
   6  * License as published by the Free Software Foundation; either
   7  * version 2.1 of the License, or (at your option) any later version.
   8  *
   9  * FFmpeg is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  12  * Lesser General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU Lesser General Public
  15  * License along with FFmpeg; if not, write to the Free Software
  16  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  17  */
  18
  19 #include "libavutil/buffer.h"
  20 #include "libavutil/hwcontext.h"
  21 #include "libavutil/log.h"
  22 #include "libavutil/opt.h"
  23
  24 #include "avfilter.h"
  25 #include "formats.h"
  26 #include "internal.h"
  27 #include "video.h"
  28
  29 typedef struct CudaUploadContext {
  30     const AVClass *class;
  31     int device_idx;
  32
  33     AVBufferRef *hwdevice;
  34     AVBufferRef *hwframe;
  35 } CudaUploadContext;
  36
  37 static av_cold int cudaupload_init(AVFilterContext *ctx)
  38 {
  39     CudaUploadContext *s = ctx->priv;
  40     char buf[64] = { 0 };
  41
  42     snprintf(buf, sizeof(buf), "%d", s->device_idx);
  43
  44     return av_hwdevice_ctx_create(&s->hwdevice, AV_HWDEVICE_TYPE_CUDA, buf, NULL, 0);
  45 }
  46
  47 static av_cold void cudaupload_uninit(AVFilterContext *ctx)
  48 {
  49     CudaUploadContext *s = ctx->priv;
  50
  51     av_buffer_unref(&s->hwframe);
  52     av_buffer_unref(&s->hwdevice);
  53 }
  54
  55 static int cudaupload_query_formats(AVFilterContext *ctx)
  56 {
  57     int ret;
  58
  59     static const enum AVPixelFormat input_pix_fmts[] = {
  60         AV_PIX_FMT_NV12, AV_PIX_FMT_YUV420P, AV_PIX_FMT_YUV444P,
  61         AV_PIX_FMT_P010, AV_PIX_FMT_P016, AV_PIX_FMT_YUV444P16,
  62         AV_PIX_FMT_NONE,
  63     };
  64     static const enum AVPixelFormat output_pix_fmts[] = {
  65         AV_PIX_FMT_CUDA, AV_PIX_FMT_NONE,
  66     };
  67     AVFilterFormats *in_fmts  = ff_make_format_list(input_pix_fmts);
  68     AVFilterFormats *out_fmts;
  69
  70     ret = ff_formats_ref(in_fmts, &ctx->inputs[0]->out_formats);
  71     if (ret < 0)
  72         return ret;
  73
  74     out_fmts = ff_make_format_list(output_pix_fmts);
  75
  76     ret = ff_formats_ref(out_fmts, &ctx->outputs[0]->in_formats);
  77     if (ret < 0)
  78         return ret;
  79
  80     return 0;
  81 }
  82
  83 static int cudaupload_config_output(AVFilterLink *outlink)
  84 {
  85     AVFilterContext *ctx = outlink->src;
  86     AVFilterLink *inlink = ctx->inputs[0];
  87     CudaUploadContext *s = ctx->priv;
  88
  89     AVHWFramesContext *hwframe_ctx;
  90     int ret;
  91
  92     av_buffer_unref(&s->hwframe);
  93     s->hwframe = av_hwframe_ctx_alloc(s->hwdevice);
  94     if (!s->hwframe)
  95         return AVERROR(ENOMEM);
  96
  97     hwframe_ctx            = (AVHWFramesContext*)s->hwframe->data;
  98     hwframe_ctx->format    = AV_PIX_FMT_CUDA;
  99     hwframe_ctx->sw_format = inlink->format;
 100     hwframe_ctx->width     = inlink->w;
 101     hwframe_ctx->height    = inlink->h;
 102
 103     ret = av_hwframe_ctx_init(s->hwframe);
 104     if (ret < 0)
 105         return ret;
 106
 107     outlink->hw_frames_ctx = av_buffer_ref(s->hwframe);
 108     if (!outlink->hw_frames_ctx)
 109         return AVERROR(ENOMEM);
 110
 111     return 0;
 112 }
 113
 114 static int cudaupload_filter_frame(AVFilterLink *link, AVFrame *in)
 115 {
 116     AVFilterContext   *ctx = link->dst;
 117     AVFilterLink  *outlink = ctx->outputs[0];
 118
 119     AVFrame *out = NULL;
 120     int ret;
 121
 122     out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
 123     if (!out) {
 124         ret = AVERROR(ENOMEM);
 125         goto fail;
 126     }
 127
 128     out->width  = in->width;
 129     out->height = in->height;
 130
 131     ret = av_hwframe_transfer_data(out, in, 0);
 132     if (ret < 0) {
 133         av_log(ctx, AV_LOG_ERROR, "Error transferring data to the GPU\n");
 134         goto fail;
 135     }
 136
 137     ret = av_frame_copy_props(out, in);
 138     if (ret < 0)
 139         goto fail;
 140
 141     av_frame_free(&in);
 142
 143     return ff_filter_frame(ctx->outputs[0], out);
 144 fail:
 145     av_frame_free(&in);
 146     av_frame_free(&out);
 147     return ret;
 148 }
 149
 150 #define OFFSET(x) offsetof(CudaUploadContext, x)
 151 #define FLAGS (AV_OPT_FLAG_FILTERING_PARAM | AV_OPT_FLAG_VIDEO_PARAM)
 152 static const AVOption cudaupload_options[] = {
 153     { "device", "Number of the device to use", OFFSET(device_idx), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, INT_MAX, FLAGS },
 154     { NULL },
 155 };
 156
 157 AVFILTER_DEFINE_CLASS(cudaupload);
 158
 159 static const AVFilterPad cudaupload_inputs[] = {
 160     {
 161         .name         = "default",
 162         .type         = AVMEDIA_TYPE_VIDEO,
 163         .filter_frame = cudaupload_filter_frame,
 164     },
 165     { NULL }
 166 };
 167
 168 static const AVFilterPad cudaupload_outputs[] = {
 169     {
 170         .name         = "default",
 171         .type         = AVMEDIA_TYPE_VIDEO,
 172         .config_props = cudaupload_config_output,
 173     },
 174     { NULL }
 175 };
 176
 177 AVFilter ff_vf_hwupload_cuda = {
 178     .name        = "hwupload_cuda",
 179     .description = NULL_IF_CONFIG_SMALL("Upload a system memory frame to a CUDA device."),
 180
 181     .init      = cudaupload_init,
 182     .uninit    = cudaupload_uninit,
 183
 184     .query_formats = cudaupload_query_formats,
 185
 186     .priv_size  = sizeof(CudaUploadContext),
 187     .priv_class = &cudaupload_class,
 188
 189     .inputs    = cudaupload_inputs,
 190     .outputs   = cudaupload_outputs,
 191
 192     .flags_internal = FF_FILTER_FLAG_HWFRAME_AWARE,
 193 };