avcodec/av1dec: support exporting Film Grain params as frame side data

Signed-off-by: James Almer <jamrial@gmail.com>

4 files changed, 83 insertions, 9 deletions

diff --git a/libavcodec/av1dec.c b/libavcodec/av1dec.c
index bc897af9cf..f1cc5150c3 100644
--- a/libavcodec/av1dec.c
+++ b/libavcodec/av1dec.c
@@ -18,6 +18,7 @@
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  */
 
+#include "libavutil/film_grain_params.h"
 #include "libavutil/pixdesc.h"
 #include "avcodec.h"
 #include "av1dec.h"
@@ -737,6 +738,66 @@ fail:
     return ret;
 }
 
+static int export_film_grain(AVCodecContext *avctx, AVFrame *frame)
+{
+    AV1DecContext *s = avctx->priv_data;
+    const AV1RawFilmGrainParams *film_grain = &s->cur_frame.film_grain;
+    AVFilmGrainParams *fgp;
+    AVFilmGrainAOMParams *aom;
+
+    if (!film_grain->apply_grain)
+        return 0;
+
+    fgp = av_film_grain_params_create_side_data(frame);
+    if (!fgp)
+        return AVERROR(ENOMEM);
+
+    fgp->type = AV_FILM_GRAIN_PARAMS_AV1;
+    fgp->seed = film_grain->grain_seed;
+
+    aom = &fgp->codec.aom;
+    aom->chroma_scaling_from_luma = film_grain->chroma_scaling_from_luma;
+    aom->scaling_shift = film_grain->grain_scaling_minus_8 + 8;
+    aom->ar_coeff_lag = film_grain->ar_coeff_lag;
+    aom->ar_coeff_shift = film_grain->ar_coeff_shift_minus_6 + 6;
+    aom->grain_scale_shift = film_grain->grain_scale_shift;
+    aom->overlap_flag = film_grain->overlap_flag;
+    aom->limit_output_range = film_grain->clip_to_restricted_range;
+
+    aom->num_y_points = film_grain->num_y_points;
+    for (int i = 0; i < film_grain->num_y_points; i++) {
+        aom->y_points[i][0] = film_grain->point_y_value[i];
+        aom->y_points[i][1] = film_grain->point_y_scaling[i];
+    }
+    aom->num_uv_points[0] = film_grain->num_cb_points;
+    for (int i = 0; i < film_grain->num_cb_points; i++) {
+        aom->uv_points[0][i][0] = film_grain->point_cb_value[i];
+        aom->uv_points[0][i][1] = film_grain->point_cb_scaling[i];
+    }
+    aom->num_uv_points[1] = film_grain->num_cr_points;
+    for (int i = 0; i < film_grain->num_cr_points; i++) {
+        aom->uv_points[1][i][0] = film_grain->point_cr_value[i];
+        aom->uv_points[1][i][1] = film_grain->point_cr_scaling[i];
+    }
+
+    for (int i = 0; i < 24; i++) {
+        aom->ar_coeffs_y[i] = film_grain->ar_coeffs_y_plus_128[i] - 128;
+    }
+    for (int i = 0; i < 25; i++) {
+        aom->ar_coeffs_uv[0][i] = film_grain->ar_coeffs_cb_plus_128[i] - 128;
+        aom->ar_coeffs_uv[1][i] = film_grain->ar_coeffs_cr_plus_128[i] - 128;
+    }
+
+    aom->uv_mult[0] = film_grain->cb_mult;
+    aom->uv_mult[1] = film_grain->cr_mult;
+    aom->uv_mult_luma[0] = film_grain->cb_luma_mult;
+    aom->uv_mult_luma[1] = film_grain->cr_luma_mult;
+    aom->uv_offset[0] = film_grain->cb_offset;
+    aom->uv_offset[1] = film_grain->cr_offset;
+
+    return 0;
+}
+
 static int set_output_frame(AVCodecContext *avctx, AVFrame *frame,
                             const AVPacket *pkt, int *got_frame)
 {
@@ -748,6 +809,14 @@ static int set_output_frame(AVCodecContext *avctx, AVFrame *frame,
     if (ret < 0)
         return ret;
 
+    if (avctx->export_side_data & AV_CODEC_EXPORT_DATA_FILM_GRAIN) {
+        ret = export_film_grain(avctx, frame);
+        if (ret < 0) {
+            av_frame_unref(frame);
+            return ret;
+        }
+    }
+
     frame->pts = pkt->pts;
     frame->pkt_dts = pkt->dts;
     frame->pkt_size = pkt->size;
diff --git a/libavcodec/dxva2_av1.c b/libavcodec/dxva2_av1.c
index 72c38f5d12..aa14e473df 100644
--- a/libavcodec/dxva2_av1.c
+++ b/libavcodec/dxva2_av1.c
@@ -62,6 +62,7 @@ static int fill_picture_parameters(const AVCodecContext *avctx, AVDXVAContext *c
     const AV1RawFilmGrainParams *film_grain = &h->cur_frame.film_grain;
 
     unsigned char remap_lr_type[4] = { AV1_RESTORE_NONE, AV1_RESTORE_SWITCHABLE, AV1_RESTORE_WIENER, AV1_RESTORE_SGRPROJ };
+    int apply_grain = !(avctx->export_side_data & AV_CODEC_EXPORT_DATA_FILM_GRAIN) && film_grain->apply_grain;
 
     memset(pp, 0, sizeof(*pp));
 
@@ -99,7 +100,7 @@ static int fill_picture_parameters(const AVCodecContext *avctx, AVDXVAContext *c
     pp->coding.integer_mv                   = frame_header->force_integer_mv || !(frame_header->frame_type & 1);
     pp->coding.cdef                         = seq->enable_cdef;
     pp->coding.restoration                  = seq->enable_restoration;
-    pp->coding.film_grain                   = seq->film_grain_params_present;
+    pp->coding.film_grain                   = seq->film_grain_params_present && !(avctx->export_side_data & AV_CODEC_EXPORT_DATA_FILM_GRAIN);
     pp->coding.intrabc                      = frame_header->allow_intrabc;
     pp->coding.high_precision_mv            = frame_header->allow_high_precision_mv;
     pp->coding.switchable_motion_mode       = frame_header->is_motion_mode_switchable;
@@ -215,7 +216,7 @@ static int fill_picture_parameters(const AVCodecContext *avctx, AVDXVAContext *c
     }
 
     /* Film grain */
-    if (film_grain->apply_grain) {
+    if (apply_grain) {
         pp->film_grain.apply_grain              = 1;
         pp->film_grain.scaling_shift_minus8     = film_grain->grain_scaling_minus_8;
         pp->film_grain.chroma_scaling_from_luma = film_grain->chroma_scaling_from_luma;
diff --git a/libavcodec/nvdec_av1.c b/libavcodec/nvdec_av1.c
index bd920d814a..1ce846a60d 100644
--- a/libavcodec/nvdec_av1.c
+++ b/libavcodec/nvdec_av1.c
@@ -53,9 +53,10 @@ static int nvdec_av1_start_frame(AVCodecContext *avctx, const uint8_t *buffer, u
 
     unsigned char remap_lr_type[4] = { AV1_RESTORE_NONE, AV1_RESTORE_SWITCHABLE, AV1_RESTORE_WIENER, AV1_RESTORE_SGRPROJ };
 
+    int apply_grain = !(avctx->export_side_data & AV_CODEC_EXPORT_DATA_FILM_GRAIN) && film_grain->apply_grain;
     int ret, i, j;
 
-    ret = ff_nvdec_start_frame_sep_ref(avctx, cur_frame, film_grain->apply_grain);
+    ret = ff_nvdec_start_frame_sep_ref(avctx, cur_frame, apply_grain);
     if (ret < 0)
         return ret;
 
@@ -95,7 +96,8 @@ static int nvdec_av1_start_frame(AVCodecContext *avctx, const uint8_t *buffer, u
             .enable_superres            = seq->enable_superres,
             .enable_cdef                = seq->enable_cdef,
             .enable_restoration         = seq->enable_restoration,
-            .enable_fgs                 = seq->film_grain_params_present,
+            .enable_fgs                 = seq->film_grain_params_present &&
+                                          !(avctx->export_side_data & AV_CODEC_EXPORT_DATA_FILM_GRAIN),
 
             /* Frame Header */
             .frame_type                   = frame_header->frame_type,
@@ -182,7 +184,7 @@ static int nvdec_av1_start_frame(AVCodecContext *avctx, const uint8_t *buffer, u
             .spatial_layer_id  = s->cur_frame.spatial_id,
 
             /* Film Grain Params */
-            .apply_grain              = film_grain->apply_grain,
+            .apply_grain              = apply_grain,
             .overlap_flag             = film_grain->overlap_flag,
             .scaling_shift_minus8     = film_grain->grain_scaling_minus_8,
             .chroma_scaling_from_luma = film_grain->chroma_scaling_from_luma,
@@ -259,7 +261,7 @@ static int nvdec_av1_start_frame(AVCodecContext *avctx, const uint8_t *buffer, u
     }
 
     /* Film Grain Params */
-    if (film_grain->apply_grain) {
+    if (apply_grain) {
         for (i = 0; i < 14; ++i) {
             ppc->scaling_points_y[i][0] = film_grain->point_y_value[i];
             ppc->scaling_points_y[i][1] = film_grain->point_y_scaling[i];
diff --git a/libavcodec/vaapi_av1.c b/libavcodec/vaapi_av1.c
index b76e217daa..1809b485aa 100644
--- a/libavcodec/vaapi_av1.c
+++ b/libavcodec/vaapi_av1.c
@@ -61,6 +61,7 @@ static int vaapi_av1_start_frame(AVCodecContext *avctx,
     VADecPictureParameterBufferAV1 pic_param;
     int8_t bit_depth_idx;
     int err = 0;
+    int apply_grain = !(avctx->export_side_data & AV_CODEC_EXPORT_DATA_FILM_GRAIN) && film_grain->apply_grain;
     uint8_t remap_lr_type[4] = {AV1_RESTORE_NONE, AV1_RESTORE_SWITCHABLE, AV1_RESTORE_WIENER, AV1_RESTORE_SGRPROJ};
 
     pic->output_surface = vaapi_av1_surface_id(&s->cur_frame);
@@ -107,7 +108,8 @@ static int vaapi_av1_start_frame(AVCodecContext *avctx,
             .subsampling_x              = seq->color_config.subsampling_x,
             .subsampling_y              = seq->color_config.subsampling_y,
             .chroma_sample_position     = seq->color_config.chroma_sample_position,
-            .film_grain_params_present  = seq->film_grain_params_present,
+            .film_grain_params_present  = seq->film_grain_params_present &&
+                                          !(avctx->export_side_data & AV_CODEC_EXPORT_DATA_FILM_GRAIN),
         },
         .seg_info.segment_info_fields.bits = {
             .enabled         = frame_header->segmentation_enabled,
@@ -117,7 +119,7 @@ static int vaapi_av1_start_frame(AVCodecContext *avctx,
         },
         .film_grain_info = {
             .film_grain_info_fields.bits = {
-                .apply_grain              = film_grain->apply_grain,
+                .apply_grain              = apply_grain,
                 .chroma_scaling_from_luma = film_grain->chroma_scaling_from_luma,
                 .grain_scaling_minus_8    = film_grain->grain_scaling_minus_8,
                 .ar_coeff_lag             = film_grain->ar_coeff_lag,
@@ -215,7 +217,7 @@ static int vaapi_av1_start_frame(AVCodecContext *avctx,
         for (int j = 0; j < 6; j++)
             pic_param.wm[i - 1].wmmat[j] = s->cur_frame.gm_params[i][j];
     }
-    if (film_grain->apply_grain) {
+    if (apply_grain) {
         for (int i = 0; i < film_grain->num_y_points; i++) {
             pic_param.film_grain_info.point_y_value[i] =
                 film_grain->point_y_value[i];