avcodec/h264_slice: compute and export film grain seed

From SMPTE RDD 5-2006, the grain seed is to be computed from the
following definition of `pic_offset`:

> When decoding H.264 | MPEG-4 AVC bitstreams, pic_offset is defined as
> follows:
>   - pic_offset = PicOrderCnt(CurrPic) + (PicOrderCnt_offset << 5)
> where:
>   - PicOrderCnt(CurrPic) is the picture order count of the current frame,
>     which shall be derived from [the video stream].
>
>   - PicOrderCnt_offset is set to idr_pic_id on IDR frames. idr_pic_id
>     shall be read from the slice header of [the video stream]. On non-IDR I
>     frames, PicOrderCnt_offset is set to 0. A frame shall be classified as I
>     frame when all its slices are I slices, which may be optionally
>     designated by setting primary_pic_type to 0 in the access delimiter NAL
>     unit. Otherwise, PicOrderCnt_offset it not changed. PicOrderCnt_offset is
>     updated in decoding order.

Co-authored-by: James Almer <jamrial@gmail.com>
Signed-off-by: Niklas Haas <git@haasn.dev>
Signed-off-by: James Almer <jamrial@gmail.com>

3 files changed, 29 insertions, 1 deletions

diff --git a/libavcodec/h264_slice.c b/libavcodec/h264_slice.c
index 0d7107d455..9244d2d5dd 100644
--- a/libavcodec/h264_slice.c
+++ b/libavcodec/h264_slice.c
@@ -406,6 +406,7 @@ int ff_h264_update_thread_context(AVCodecContext *dst,
 
     h->next_output_pic   = h1->next_output_pic;
     h->next_outputed_poc = h1->next_outputed_poc;
+    h->poc_offset        = h1->poc_offset;
 
     memcpy(h->mmco, h1->mmco, sizeof(h->mmco));
     h->nb_mmco         = h1->nb_mmco;
@@ -1335,6 +1336,7 @@ static int h264_export_frame_props(H264Context *h)
             return AVERROR(ENOMEM);
 
         fgp->type = AV_FILM_GRAIN_PARAMS_H274;
+        fgp->seed = cur->poc + (h->poc_offset << 5);
 
         fgp->codec.h274.model_id = fgc->model_id;
         if (fgc->separate_colour_description_present_flag) {
@@ -1543,6 +1545,11 @@ static int h264_field_start(H264Context *h, const H264SliceContext *sl,
     h->poc.delta_poc[0]     = sl->delta_poc[0];
     h->poc.delta_poc[1]     = sl->delta_poc[1];
 
+    if (nal->type == H264_NAL_IDR_SLICE)
+        h->poc_offset = sl->idr_pic_id;
+    else if (h->picture_intra_only)
+        h->poc_offset = 0;
+
     /* Shorten frame num gaps so we don't have to allocate reference
      * frames just to throw them away */
     if (h->poc.frame_num != h->poc.prev_frame_num) {
@@ -1891,7 +1898,7 @@ static int h264_slice_header_parse(const H264Context *h, H264SliceContext *sl,
     }
 
     if (nal->type == H264_NAL_IDR_SLICE)
-        get_ue_golomb_long(&sl->gb); /* idr_pic_id */
+        sl->idr_pic_id = get_ue_golomb_long(&sl->gb);
 
     sl->poc_lsb = 0;
     sl->delta_poc_bottom = 0;
diff --git a/libavcodec/h264dec.c b/libavcodec/h264dec.c
index 07ff9110d2..f46d29baa6 100644
--- a/libavcodec/h264dec.c
+++ b/libavcodec/h264dec.c
@@ -485,6 +485,8 @@ static void h264_decode_flush(AVCodecContext *avctx)
 static int get_last_needed_nal(H264Context *h)
 {
     int nals_needed = 0;
+    int slice_type = 0;
+    int picture_intra_only = 1;
     int first_slice = 0;
     int i, ret;
 
@@ -516,11 +518,23 @@ static int get_last_needed_nal(H264Context *h)
                 !first_slice ||
                 first_slice != nal->type)
                 nals_needed = i;
+            slice_type = get_ue_golomb_31(&gb);
+            if (slice_type > 9) {
+                if (h->avctx->err_recognition & AV_EF_EXPLODE)
+                    return AVERROR_INVALIDDATA;
+            }
+            if (slice_type > 4)
+                slice_type -= 5;
+
+            slice_type = ff_h264_golomb_to_pict_type[slice_type];
+            picture_intra_only &= (slice_type & 3) == AV_PICTURE_TYPE_I;
             if (!first_slice)
                 first_slice = nal->type;
         }
     }
 
+    h->picture_intra_only = picture_intra_only;
+
     return nals_needed;
 }
 
diff --git a/libavcodec/h264dec.h b/libavcodec/h264dec.h
index 125966aa04..7c419de051 100644
--- a/libavcodec/h264dec.h
+++ b/libavcodec/h264dec.h
@@ -331,6 +331,7 @@ typedef struct H264SliceContext {
     int explicit_ref_marking;
 
     int frame_num;
+    int idr_pic_id;
     int poc_lsb;
     int delta_poc_bottom;
     int delta_poc[2];
@@ -384,6 +385,11 @@ typedef struct H264Context {
      */
     int picture_idr;
 
+    /*
+     * Set to 1 when the current picture contains only I slices, 0 otherwise.
+     */
+    int picture_intra_only;
+
     int crop_left;
     int crop_right;
     int crop_top;
@@ -473,6 +479,7 @@ typedef struct H264Context {
     int last_pocs[MAX_DELAYED_PIC_COUNT];
     H264Picture *next_output_pic;
     int next_outputed_poc;
+    int poc_offset;         ///< PicOrderCnt_offset from SMPTE RDD-2006
 
     /**
      * memory management control operations buffer.