[FFmpeg-devel,v5,21/22] h264_metadata_bsf: Refactor the filter function into smaller parts

Message ID 20200503220548.23543-21-sw@jkqxz.net
State New
Headers
Series [FFmpeg-devel,v5,01/22] cbs: Mention all codecs in unit type comment |

Checks

Context Check Description
andriy/default pending
andriy/configure warning Failed to apply patch

Commit Message

Mark Thompson May 3, 2020, 10:05 p.m. UTC
---
 libavcodec/h264_metadata_bsf.c | 447 ++++++++++++++++++---------------
 1 file changed, 244 insertions(+), 203 deletions(-)
  

Patch

diff --git a/libavcodec/h264_metadata_bsf.c b/libavcodec/h264_metadata_bsf.c
index 451eaa0b67..bfb5b7d7e4 100644
--- a/libavcodec/h264_metadata_bsf.c
+++ b/libavcodec/h264_metadata_bsf.c
@@ -54,6 +54,7 @@  typedef struct H264MetadataContext {
     int done_first_au;
 
     int aud;
+    H264RawAUD aud_nal;
 
     AVRational sample_aspect_ratio;
 
@@ -76,6 +77,7 @@  typedef struct H264MetadataContext {
     int crop_bottom;
 
     const char *sei_user_data;
+    H264RawSEIPayload sei_user_data_payload;
 
     int delete_filler;
 
@@ -87,6 +89,59 @@  typedef struct H264MetadataContext {
 } H264MetadataContext;
 
 
+static int h264_metadata_insert_aud(AVBSFContext *bsf,
+                                    CodedBitstreamFragment *au)
+{
+    H264MetadataContext *ctx = bsf->priv_data;
+    int primary_pic_type_mask = 0xff;
+    int err, i, j;
+
+    static const int primary_pic_type_table[] = {
+        0x084, // 2, 7
+        0x0a5, // 0, 2, 5, 7
+        0x0e7, // 0, 1, 2, 5, 6, 7
+        0x210, // 4, 9
+        0x318, // 3, 4, 8, 9
+        0x294, // 2, 4, 7, 9
+        0x3bd, // 0, 2, 3, 4, 5, 7, 8, 9
+        0x3ff, // 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
+    };
+
+    for (i = 0; i < au->nb_units; i++) {
+        if (au->units[i].type == H264_NAL_SLICE ||
+            au->units[i].type == H264_NAL_IDR_SLICE) {
+            H264RawSlice *slice = au->units[i].content;
+            for (j = 0; j < FF_ARRAY_ELEMS(primary_pic_type_table); j++) {
+                if (!(primary_pic_type_table[j] &
+                      (1 << slice->header.slice_type)))
+                    primary_pic_type_mask &= ~(1 << j);
+            }
+        }
+    }
+    for (j = 0; j < FF_ARRAY_ELEMS(primary_pic_type_table); j++)
+        if (primary_pic_type_mask & (1 << j))
+            break;
+    if (j >= FF_ARRAY_ELEMS(primary_pic_type_table)) {
+        av_log(bsf, AV_LOG_ERROR, "No usable primary_pic_type: "
+               "invalid slice types?\n");
+        return AVERROR_INVALIDDATA;
+    }
+
+    ctx->aud_nal = (H264RawAUD) {
+        .nal_unit_header.nal_unit_type = H264_NAL_AUD,
+        .primary_pic_type = j,
+    };
+
+    err = ff_cbs_insert_unit_content(ctx->cbc, au,
+                                     0, H264_NAL_AUD, &ctx->aud_nal, NULL);
+    if (err < 0) {
+        av_log(bsf, AV_LOG_ERROR, "Failed to insert AUD.\n");
+        return err;
+    }
+
+    return 0;
+}
+
 static int h264_metadata_update_sps(AVBSFContext *bsf,
                                     H264RawSPS *sps)
 {
@@ -318,221 +373,60 @@  static int h264_metadata_update_side_data(AVBSFContext *bsf, AVPacket *pkt)
     return 0;
 }
 
-static int h264_metadata_filter(AVBSFContext *bsf, AVPacket *pkt)
+static int h264_metadata_handle_display_orientation(AVBSFContext *bsf,
+                                                    AVPacket *pkt,
+                                                    CodedBitstreamFragment *au,
+                                                    int seek_point)
 {
     H264MetadataContext *ctx = bsf->priv_data;
-    CodedBitstreamFragment *au = &ctx->access_unit;
-    int err, i, j, has_sps;
-    H264RawAUD aud;
-
-    err = ff_bsf_get_packet_ref(bsf, pkt);
-    if (err < 0)
-        return err;
-
-    err = h264_metadata_update_side_data(bsf, pkt);
-    if (err < 0)
-        goto fail;
-
-    err = ff_cbs_read_packet(ctx->cbc, au, pkt);
-    if (err < 0) {
-        av_log(bsf, AV_LOG_ERROR, "Failed to read packet.\n");
-        goto fail;
-    }
-
-    if (au->nb_units == 0) {
-        av_log(bsf, AV_LOG_ERROR, "No NAL units in packet.\n");
-        err = AVERROR_INVALIDDATA;
-        goto fail;
-    }
-
-    // If an AUD is present, it must be the first NAL unit.
-    if (au->units[0].type == H264_NAL_AUD) {
-        if (ctx->aud == REMOVE)
-            ff_cbs_delete_unit(ctx->cbc, au, 0);
-    } else {
-        if (ctx->aud == INSERT) {
-            static const int primary_pic_type_table[] = {
-                0x084, // 2, 7
-                0x0a5, // 0, 2, 5, 7
-                0x0e7, // 0, 1, 2, 5, 6, 7
-                0x210, // 4, 9
-                0x318, // 3, 4, 8, 9
-                0x294, // 2, 4, 7, 9
-                0x3bd, // 0, 2, 3, 4, 5, 7, 8, 9
-                0x3ff, // 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
-            };
-            int primary_pic_type_mask = 0xff;
-
-            for (i = 0; i < au->nb_units; i++) {
-                if (au->units[i].type == H264_NAL_SLICE ||
-                    au->units[i].type == H264_NAL_IDR_SLICE) {
-                    H264RawSlice *slice = au->units[i].content;
-                    for (j = 0; j < FF_ARRAY_ELEMS(primary_pic_type_table); j++) {
-                         if (!(primary_pic_type_table[j] &
-                               (1 << slice->header.slice_type)))
-                             primary_pic_type_mask &= ~(1 << j);
-                    }
-                }
-            }
-            for (j = 0; j < FF_ARRAY_ELEMS(primary_pic_type_table); j++)
-                if (primary_pic_type_mask & (1 << j))
-                    break;
-            if (j >= FF_ARRAY_ELEMS(primary_pic_type_table)) {
-                av_log(bsf, AV_LOG_ERROR, "No usable primary_pic_type: "
-                       "invalid slice types?\n");
-                err = AVERROR_INVALIDDATA;
-                goto fail;
-            }
-
-            aud = (H264RawAUD) {
-                .nal_unit_header.nal_unit_type = H264_NAL_AUD,
-                .primary_pic_type = j,
-            };
+    int err, i, j;
 
-            err = ff_cbs_insert_unit_content(ctx->cbc, au,
-                                             0, H264_NAL_AUD, &aud, NULL);
-            if (err < 0) {
-                av_log(bsf, AV_LOG_ERROR, "Failed to insert AUD.\n");
-                goto fail;
-            }
-        }
-    }
-
-    has_sps = 0;
-    for (i = 0; i < au->nb_units; i++) {
-        if (au->units[i].type == H264_NAL_SPS) {
-            err = h264_metadata_update_sps(bsf, au->units[i].content);
-            if (err < 0)
-                goto fail;
-            has_sps = 1;
-        }
-    }
+    for (i = au->nb_units - 1; i >= 0; i--) {
+        H264RawSEI *sei;
+        if (au->units[i].type != H264_NAL_SEI)
+            continue;
+        sei = au->units[i].content;
 
-    // Only insert the SEI in access units containing SPSs, and also
-    // unconditionally in the first access unit we ever see.
-    if (ctx->sei_user_data && (has_sps || !ctx->done_first_au)) {
-        H264RawSEIPayload payload = {
-            .payload_type = H264_SEI_TYPE_USER_DATA_UNREGISTERED,
-        };
-        H264RawSEIUserDataUnregistered *udu =
-            &payload.payload.user_data_unregistered;
+        for (j = sei->payload_count - 1; j >= 0; j--) {
+            H264RawSEIDisplayOrientation *disp;
+            int32_t *matrix;
 
-        for (i = j = 0; j < 32 && ctx->sei_user_data[i]; i++) {
-            int c, v;
-            c = ctx->sei_user_data[i];
-            if (c == '-') {
+            if (sei->payload[j].payload_type !=
+                H264_SEI_TYPE_DISPLAY_ORIENTATION)
                 continue;
-            } else if (av_isxdigit(c)) {
-                c = av_tolower(c);
-                v = (c <= '9' ? c - '0' : c - 'a' + 10);
-            } else {
-                goto invalid_user_data;
-            }
-            if (j & 1)
-                udu->uuid_iso_iec_11578[j / 2] |= v;
-            else
-                udu->uuid_iso_iec_11578[j / 2] = v << 4;
-            ++j;
-        }
-        if (j == 32 && ctx->sei_user_data[i] == '+') {
-            size_t len = strlen(ctx->sei_user_data + i + 1);
-
-            udu->data_ref = av_buffer_alloc(len + 1);
-            if (!udu->data_ref) {
-                err = AVERROR(ENOMEM);
-                goto fail;
-            }
-
-            udu->data        = udu->data_ref->data;
-            udu->data_length = len + 1;
-            memcpy(udu->data, ctx->sei_user_data + i + 1, len + 1);
+            disp = &sei->payload[j].payload.display_orientation;
 
-            err = ff_cbs_h264_add_sei_message(ctx->cbc, au, &payload);
-            if (err < 0) {
-                av_log(bsf, AV_LOG_ERROR, "Failed to add user data SEI "
-                       "message to access unit.\n");
-                goto fail;
-            }
-
-        } else {
-        invalid_user_data:
-            av_log(bsf, AV_LOG_ERROR, "Invalid user data: "
-                   "must be \"UUID+string\".\n");
-            err = AVERROR(EINVAL);
-            goto fail;
-        }
-    }
-
-    if (ctx->delete_filler) {
-        for (i = au->nb_units - 1; i >= 0; i--) {
-            if (au->units[i].type == H264_NAL_FILLER_DATA) {
-                ff_cbs_delete_unit(ctx->cbc, au, i);
+            if (ctx->display_orientation == REMOVE ||
+                ctx->display_orientation == INSERT) {
+                ff_cbs_h264_delete_sei_message(ctx->cbc, au,
+                                               &au->units[i], j);
                 continue;
             }
 
-            if (au->units[i].type == H264_NAL_SEI) {
-                // Filler SEI messages.
-                H264RawSEI *sei = au->units[i].content;
-
-                for (j = sei->payload_count - 1; j >= 0; j--) {
-                    if (sei->payload[j].payload_type ==
-                        H264_SEI_TYPE_FILLER_PAYLOAD)
-                        ff_cbs_h264_delete_sei_message(ctx->cbc, au,
-                                                       &au->units[i], j);
-                }
+            matrix = av_malloc(9 * sizeof(int32_t));
+            if (!matrix)
+                return AVERROR(ENOMEM);
+
+            av_display_rotation_set(matrix,
+                                    disp->anticlockwise_rotation *
+                                    180.0 / 65536.0);
+            av_display_matrix_flip(matrix, disp->hor_flip, disp->ver_flip);
+
+            // If there are multiple display orientation messages in an
+            // access unit, then the last one added to the packet (i.e.
+            // the first one in the access unit) will prevail.
+            err = av_packet_add_side_data(pkt, AV_PKT_DATA_DISPLAYMATRIX,
+                                          (uint8_t*)matrix,
+                                          9 * sizeof(int32_t));
+            if (err < 0) {
+                av_log(bsf, AV_LOG_ERROR, "Failed to attach extracted "
+                       "displaymatrix side data to packet.\n");
+                av_freep(matrix);
+                return AVERROR(ENOMEM);
             }
         }
     }
 
-    if (ctx->display_orientation != PASS) {
-        for (i = au->nb_units - 1; i >= 0; i--) {
-            H264RawSEI *sei;
-            if (au->units[i].type != H264_NAL_SEI)
-                continue;
-            sei = au->units[i].content;
-
-            for (j = sei->payload_count - 1; j >= 0; j--) {
-                H264RawSEIDisplayOrientation *disp;
-                int32_t *matrix;
-
-                if (sei->payload[j].payload_type !=
-                    H264_SEI_TYPE_DISPLAY_ORIENTATION)
-                    continue;
-                disp = &sei->payload[j].payload.display_orientation;
-
-                if (ctx->display_orientation == REMOVE ||
-                    ctx->display_orientation == INSERT) {
-                    ff_cbs_h264_delete_sei_message(ctx->cbc, au,
-                                                   &au->units[i], j);
-                    continue;
-                }
-
-                matrix = av_malloc(9 * sizeof(int32_t));
-                if (!matrix) {
-                    err = AVERROR(ENOMEM);
-                    goto fail;
-                }
-
-                av_display_rotation_set(matrix,
-                                        disp->anticlockwise_rotation *
-                                        180.0 / 65536.0);
-                av_display_matrix_flip(matrix, disp->hor_flip, disp->ver_flip);
-
-                // If there are multiple display orientation messages in an
-                // access unit, then the last one added to the packet (i.e.
-                // the first one in the access unit) will prevail.
-                err = av_packet_add_side_data(pkt, AV_PKT_DATA_DISPLAYMATRIX,
-                                              (uint8_t*)matrix,
-                                              9 * sizeof(int32_t));
-                if (err < 0) {
-                    av_log(bsf, AV_LOG_ERROR, "Failed to attach extracted "
-                           "displaymatrix side data to packet.\n");
-                    av_freep(matrix);
-                    goto fail;
-                }
-            }
-        }
-    }
     if (ctx->display_orientation == INSERT) {
         H264RawSEIPayload payload = {
             .payload_type = H264_SEI_TYPE_DISPLAY_ORIENTATION,
@@ -576,7 +470,7 @@  static int h264_metadata_filter(AVBSFContext *bsf, AVPacket *pkt)
             }
         }
 
-        if (has_sps || !ctx->done_first_au) {
+        if (seek_point) {
             if (!isnan(ctx->rotate)) {
                 disp->anticlockwise_rotation =
                     (uint16_t)rint((ctx->rotate >= 0.0 ? ctx->rotate
@@ -598,11 +492,108 @@  static int h264_metadata_filter(AVBSFContext *bsf, AVPacket *pkt)
             if (err < 0) {
                 av_log(bsf, AV_LOG_ERROR, "Failed to add display orientation "
                        "SEI message to access unit.\n");
+                return err;
+            }
+        }
+    }
+
+    return 0;
+}
+
+static int h264_metadata_filter(AVBSFContext *bsf, AVPacket *pkt)
+{
+    H264MetadataContext *ctx = bsf->priv_data;
+    CodedBitstreamFragment *au = &ctx->access_unit;
+    int err, i, j, has_sps, seek_point;
+
+    err = ff_bsf_get_packet_ref(bsf, pkt);
+    if (err < 0)
+        return err;
+
+    err = h264_metadata_update_side_data(bsf, pkt);
+    if (err < 0)
+        goto fail;
+
+    err = ff_cbs_read_packet(ctx->cbc, au, pkt);
+    if (err < 0) {
+        av_log(bsf, AV_LOG_ERROR, "Failed to read packet.\n");
+        goto fail;
+    }
+
+    if (au->nb_units == 0) {
+        av_log(bsf, AV_LOG_ERROR, "No NAL units in packet.\n");
+        err = AVERROR_INVALIDDATA;
+        goto fail;
+    }
+
+    // If an AUD is present, it must be the first NAL unit.
+    if (au->units[0].type == H264_NAL_AUD) {
+        if (ctx->aud == REMOVE)
+            ff_cbs_delete_unit(ctx->cbc, au, 0);
+    } else {
+        if (ctx->aud == INSERT) {
+            err = h264_metadata_insert_aud(bsf, au);
+            if (err < 0)
                 goto fail;
+        }
+    }
+
+    has_sps = 0;
+    for (i = 0; i < au->nb_units; i++) {
+        if (au->units[i].type == H264_NAL_SPS) {
+            err = h264_metadata_update_sps(bsf, au->units[i].content);
+            if (err < 0)
+                goto fail;
+            has_sps = 1;
+        }
+    }
+
+    // The current packet should be treated as a seek point for metadata
+    // insertion if any of:
+    // - It is the first packet in the stream.
+    // - It contains an SPS, indicating that a sequence might start here.
+    // - It is marked as containing a key frame.
+    seek_point = !ctx->done_first_au || has_sps ||
+        (pkt->flags & AV_PKT_FLAG_KEY);
+
+    if (ctx->sei_user_data && seek_point) {
+        err = ff_cbs_h264_add_sei_message(ctx->cbc, au,
+                                          &ctx->sei_user_data_payload);
+        if (err < 0) {
+            av_log(bsf, AV_LOG_ERROR, "Failed to add user data SEI "
+                   "message to access unit.\n");
+            goto fail;
+        }
+    }
+
+    if (ctx->delete_filler) {
+        for (i = au->nb_units - 1; i >= 0; i--) {
+            if (au->units[i].type == H264_NAL_FILLER_DATA) {
+                ff_cbs_delete_unit(ctx->cbc, au, i);
+                continue;
+            }
+
+            if (au->units[i].type == H264_NAL_SEI) {
+                // Filler SEI messages.
+                H264RawSEI *sei = au->units[i].content;
+
+                for (j = sei->payload_count - 1; j >= 0; j--) {
+                    if (sei->payload[j].payload_type ==
+                        H264_SEI_TYPE_FILLER_PAYLOAD)
+                        ff_cbs_h264_delete_sei_message(ctx->cbc, au,
+                                                       &au->units[i], j);
+                }
             }
         }
     }
 
+    if (ctx->display_orientation != PASS) {
+        err = h264_metadata_handle_display_orientation(bsf, pkt, au,
+                                                       seek_point);
+        if (err < 0)
+            goto fail;
+    }
+
     err = ff_cbs_write_packet(ctx->cbc, pkt, au);
     if (err < 0) {
         av_log(bsf, AV_LOG_ERROR, "Failed to write packet.\n");
@@ -625,7 +616,57 @@  static int h264_metadata_init(AVBSFContext *bsf)
 {
     H264MetadataContext *ctx = bsf->priv_data;
     CodedBitstreamFragment *au = &ctx->access_unit;
-    int err, i;
+    int err, i, j;
+
+    if (ctx->sei_user_data) {
+        H264RawSEIUserDataUnregistered *udu;
+        int c, v;
+        size_t len;
+
+        ctx->sei_user_data_payload = (H264RawSEIPayload) {
+            .payload_type = H264_SEI_TYPE_USER_DATA_UNREGISTERED,
+        };
+        udu = &ctx->sei_user_data_payload.payload.user_data_unregistered;
+
+        // Parse UUID.  It must be a hex string of length 32, possibly
+        // containing '-'s which we ignore.
+        for (i = j = 0; j < 32 && ctx->sei_user_data[i]; i++) {
+            c = ctx->sei_user_data[i];
+            if (c == '-') {
+                continue;
+            } else if (av_isxdigit(c)) {
+                c = av_tolower(c);
+                v = (c <= '9' ? c - '0' : c - 'a' + 10);
+            } else {
+                break;
+            }
+            if (j & 1)
+                udu->uuid_iso_iec_11578[j / 2] |= v;
+            else
+                udu->uuid_iso_iec_11578[j / 2] = v << 4;
+            ++j;
+        }
+        if (j < 32 || ctx->sei_user_data[i] != '+') {
+            av_log(bsf, AV_LOG_ERROR, "Invalid user data: "
+                   "must be of the form \"UUID+string\".\n");
+            return AVERROR(EINVAL);
+        } else {
+            // Skip over the '+'.
+            ++i;
+
+            // Length of the actual data to insert (could be zero).
+            len = strlen(ctx->sei_user_data + i);
+
+            udu->data_ref = av_buffer_alloc(len + 1);
+            if (!udu->data_ref)
+                return AVERROR(ENOMEM);
+
+            udu->data_length = len + 1;
+            udu->data        = udu->data_ref->data;
+            memcpy(udu->data, ctx->sei_user_data + 1, len);
+            udu->data[len]   = 0;
+        }
+    }
 
     err = ff_cbs_init(&ctx->cbc, AV_CODEC_ID_H264, bsf);
     if (err < 0)