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_0RGB32, AV_PIX_FMT_0BGR32,
  63         AV_PIX_FMT_NONE,
  64     };
  65     static const enum AVPixelFormat output_pix_fmts[] = {
  66         AV_PIX_FMT_CUDA, AV_PIX_FMT_NONE,
  67     };
  68     AVFilterFormats *in_fmts  = ff_make_format_list(input_pix_fmts);
  69     AVFilterFormats *out_fmts;
  70
  71     ret = ff_formats_ref(in_fmts, &ctx->inputs[0]->out_formats);
  72     if (ret < 0)
  73         return ret;
  74
  75     out_fmts = ff_make_format_list(output_pix_fmts);
  76
  77     ret = ff_formats_ref(out_fmts, &ctx->outputs[0]->in_formats);
  78     if (ret < 0)
  79         return ret;
  80
  81     return 0;
  82 }
  83
  84 static int cudaupload_config_output(AVFilterLink *outlink)
  85 {
  86     AVFilterContext *ctx = outlink->src;
  87     AVFilterLink *inlink = ctx->inputs[0];
  88     CudaUploadContext *s = ctx->priv;
  89
  90     AVHWFramesContext *hwframe_ctx;
  91     int ret;
  92
  93     av_buffer_unref(&s->hwframe);
  94     s->hwframe = av_hwframe_ctx_alloc(s->hwdevice);
  95     if (!s->hwframe)
  96         return AVERROR(ENOMEM);
  97
  98     hwframe_ctx            = (AVHWFramesContext*)s->hwframe->data;
  99     hwframe_ctx->format    = AV_PIX_FMT_CUDA;
 100     hwframe_ctx->sw_format = inlink->format;
 101     hwframe_ctx->width     = inlink->w;
 102     hwframe_ctx->height    = inlink->h;
 103
 104     ret = av_hwframe_ctx_init(s->hwframe);
 105     if (ret < 0)
 106         return ret;
 107
 108     outlink->hw_frames_ctx = av_buffer_ref(s->hwframe);
 109     if (!outlink->hw_frames_ctx)
 110         return AVERROR(ENOMEM);
 111
 112     return 0;
 113 }
 114
 115 static int cudaupload_filter_frame(AVFilterLink *link, AVFrame *in)
 116 {
 117     AVFilterContext   *ctx = link->dst;
 118     AVFilterLink  *outlink = ctx->outputs[0];
 119
 120     AVFrame *out = NULL;
 121     int ret;
 122
 123     out = ff_get_video_buffer(outlink, outlink->w, outlink->h);
 124     if (!out) {
 125         ret = AVERROR(ENOMEM);
 126         goto fail;
 127     }
 128
 129     out->width  = in->width;
 130     out->height = in->height;
 131
 132     ret = av_hwframe_transfer_data(out, in, 0);
 133     if (ret < 0) {
 134         av_log(ctx, AV_LOG_ERROR, "Error transferring data to the GPU\n");
 135         goto fail;
 136     }
 137
 138     ret = av_frame_copy_props(out, in);
 139     if (ret < 0)
 140         goto fail;
 141
 142     av_frame_free(&in);
 143
 144     return ff_filter_frame(ctx->outputs[0], out);
 145 fail:
 146     av_frame_free(&in);
 147     av_frame_free(&out);
 148     return ret;
 149 }
 150
 151 #define OFFSET(x) offsetof(CudaUploadContext, x)
 152 #define FLAGS (AV_OPT_FLAG_FILTERING_PARAM | AV_OPT_FLAG_VIDEO_PARAM)
 153 static const AVOption cudaupload_options[] = {
 154     { "device", "Number of the device to use", OFFSET(device_idx), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, INT_MAX, FLAGS },
 155     { NULL },
 156 };
 157
 158 AVFILTER_DEFINE_CLASS(cudaupload);
 159
 160 static const AVFilterPad cudaupload_inputs[] = {
 161     {
 162         .name         = "default",
 163         .type         = AVMEDIA_TYPE_VIDEO,
 164         .filter_frame = cudaupload_filter_frame,
 165     },
 166     { NULL }
 167 };
 168
 169 static const AVFilterPad cudaupload_outputs[] = {
 170     {
 171         .name         = "default",
 172         .type         = AVMEDIA_TYPE_VIDEO,
 173         .config_props = cudaupload_config_output,
 174     },
 175     { NULL }
 176 };
 177
 178 AVFilter ff_vf_hwupload_cuda = {
 179     .name        = "hwupload_cuda",
 180     .description = NULL_IF_CONFIG_SMALL("Upload a system memory frame to a CUDA device."),
 181
 182     .init      = cudaupload_init,
 183     .uninit    = cudaupload_uninit,
 184
 185     .query_formats = cudaupload_query_formats,
 186
 187     .priv_size  = sizeof(CudaUploadContext),
 188     .priv_class = &cudaupload_class,
 189
 190     .inputs    = cudaupload_inputs,
 191     .outputs   = cudaupload_outputs,
 192
 193     .flags_internal = FF_FILTER_FLAG_HWFRAME_AWARE,
 194 };