git.sesse.net Git - ffmpeg/blob - libavcodec/cbs_sei.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 "cbs.h"
  20 #include "cbs_internal.h"
  21 #include "cbs_h264.h"
  22 #include "cbs_h265.h"
  23 #include "cbs_sei.h"
  24
  25 static void cbs_free_user_data_registered(void *opaque, uint8_t *data)
  26 {
  27     SEIRawUserDataRegistered *udr = (SEIRawUserDataRegistered*)data;
  28     av_buffer_unref(&udr->data_ref);
  29     av_free(udr);
  30 }
  31
  32 static void cbs_free_user_data_unregistered(void *opaque, uint8_t *data)
  33 {
  34     SEIRawUserDataUnregistered *udu = (SEIRawUserDataUnregistered*)data;
  35     av_buffer_unref(&udu->data_ref);
  36     av_free(udu);
  37 }
  38
  39 int ff_cbs_sei_alloc_message_payload(SEIRawMessage *message,
  40                                      const SEIMessageTypeDescriptor *desc)
  41 {
  42     void (*free_func)(void*, uint8_t*);
  43
  44     av_assert0(message->payload     == NULL &&
  45                message->payload_ref == NULL);
  46     message->payload_type = desc->type;
  47
  48     if (desc->type == SEI_TYPE_USER_DATA_REGISTERED_ITU_T_T35)
  49         free_func = &cbs_free_user_data_registered;
  50     else if (desc->type == SEI_TYPE_USER_DATA_UNREGISTERED)
  51         free_func = &cbs_free_user_data_unregistered;
  52     else
  53         free_func = NULL;
  54
  55     if (free_func) {
  56         message->payload = av_mallocz(desc->size);
  57         if (!message->payload)
  58             return AVERROR(ENOMEM);
  59         message->payload_ref =
  60             av_buffer_create(message->payload, desc->size,
  61                              free_func, NULL, 0);
  62     } else {
  63         message->payload_ref = av_buffer_alloc(desc->size);
  64     }
  65     if (!message->payload_ref) {
  66         av_freep(&message->payload);
  67         return AVERROR(ENOMEM);
  68     }
  69     message->payload = message->payload_ref->data;
  70
  71     return 0;
  72 }
  73
  74 int ff_cbs_sei_list_add(SEIRawMessageList *list)
  75 {
  76     void *ptr;
  77     int old_count = list->nb_messages_allocated;
  78
  79     av_assert0(list->nb_messages <= old_count);
  80     if (list->nb_messages + 1 > old_count) {
  81         int new_count = 2 * old_count + 1;
  82
  83         ptr = av_realloc_array(list->messages,
  84                                new_count, sizeof(*list->messages));
  85         if (!ptr)
  86             return AVERROR(ENOMEM);
  87
  88         list->messages = ptr;
  89         list->nb_messages_allocated = new_count;
  90
  91         // Zero the newly-added entries.
  92         memset(list->messages + old_count, 0,
  93                (new_count - old_count) * sizeof(*list->messages));
  94     }
  95     ++list->nb_messages;
  96     return 0;
  97 }
  98
  99 void ff_cbs_sei_free_message_list(SEIRawMessageList *list)
 100 {
 101     for (int i = 0; i < list->nb_messages; i++) {
 102         SEIRawMessage *message = &list->messages[i];
 103         av_buffer_unref(&message->payload_ref);
 104         av_buffer_unref(&message->extension_data_ref);
 105     }
 106     av_free(list->messages);
 107 }
 108
 109 static int cbs_sei_get_unit(CodedBitstreamContext *ctx,
 110                             CodedBitstreamFragment *au,
 111                             int prefix,
 112                             CodedBitstreamUnit **sei_unit)
 113 {
 114     CodedBitstreamUnit *unit;
 115     int sei_type, highest_vcl_type, err, i, position;
 116
 117     switch (ctx->codec->codec_id) {
 118     case AV_CODEC_ID_H264:
 119         // (We can ignore auxiliary slices because we only have prefix
 120         // SEI in H.264 and an auxiliary picture must always follow a
 121         // primary picture.)
 122         highest_vcl_type = H264_NAL_IDR_SLICE;
 123         if (prefix)
 124             sei_type = H264_NAL_SEI;
 125         else
 126             return AVERROR(EINVAL);
 127         break;
 128     case AV_CODEC_ID_H265:
 129         highest_vcl_type = HEVC_NAL_RSV_VCL31;
 130         if (prefix)
 131             sei_type = HEVC_NAL_SEI_PREFIX;
 132         else
 133             sei_type = HEVC_NAL_SEI_SUFFIX;
 134         break;
 135     default:
 136         return AVERROR(EINVAL);
 137     }
 138
 139     // Find an existing SEI NAL unit of the right type.
 140     unit = NULL;
 141     for (i = 0; i < au->nb_units; i++) {
 142         if (au->units[i].type == sei_type) {
 143             unit = &au->units[i];
 144             break;
 145         }
 146     }
 147
 148     if (unit) {
 149         *sei_unit = unit;
 150         return 0;
 151     }
 152
 153     // Need to add a new SEI NAL unit ...
 154     if (prefix) {
 155         // ... before the first VCL NAL unit.
 156         for (i = 0; i < au->nb_units; i++) {
 157             if (au->units[i].type < highest_vcl_type)
 158                 break;
 159         }
 160         position = i;
 161     } else {
 162         // ... after the last VCL NAL unit.
 163         for (i = au->nb_units - 1; i >= 0; i--) {
 164             if (au->units[i].type < highest_vcl_type)
 165                 break;
 166         }
 167         if (i < 0) {
 168             // No VCL units; just put it at the end.
 169             position = au->nb_units;
 170         } else {
 171             position = i + 1;
 172         }
 173     }
 174
 175     err = ff_cbs_insert_unit_content(au, position, sei_type,
 176                                      NULL, NULL);
 177     if (err < 0)
 178         return err;
 179     unit = &au->units[position];
 180     unit->type = sei_type;
 181
 182     err = ff_cbs_alloc_unit_content2(ctx, unit);
 183     if (err < 0)
 184         return err;
 185
 186     switch (ctx->codec->codec_id) {
 187     case AV_CODEC_ID_H264:
 188         {
 189             H264RawSEI sei = {
 190                 .nal_unit_header = {
 191                     .nal_ref_idc   = 0,
 192                     .nal_unit_type = sei_type,
 193                 },
 194             };
 195             memcpy(unit->content, &sei, sizeof(sei));
 196         }
 197         break;
 198     case AV_CODEC_ID_H265:
 199         {
 200             H265RawSEI sei = {
 201                 .nal_unit_header = {
 202                     .nal_unit_type         = sei_type,
 203                     .nuh_layer_id          = 0,
 204                     .nuh_temporal_id_plus1 = 1,
 205                 },
 206             };
 207             memcpy(unit->content, &sei, sizeof(sei));
 208         }
 209         break;
 210     default:
 211         av_assert0(0);
 212     }
 213
 214     *sei_unit = unit;
 215     return 0;
 216 }
 217
 218 static int cbs_sei_get_message_list(CodedBitstreamContext *ctx,
 219                                     CodedBitstreamUnit *unit,
 220                                     SEIRawMessageList **list)
 221 {
 222     switch (ctx->codec->codec_id) {
 223     case AV_CODEC_ID_H264:
 224         {
 225             H264RawSEI *sei = unit->content;
 226             if (unit->type != H264_NAL_SEI)
 227                 return AVERROR(EINVAL);
 228             *list = &sei->message_list;
 229         }
 230         break;
 231     case AV_CODEC_ID_H265:
 232         {
 233             H265RawSEI *sei = unit->content;
 234             if (unit->type != HEVC_NAL_SEI_PREFIX &&
 235                 unit->type != HEVC_NAL_SEI_SUFFIX)
 236                 return AVERROR(EINVAL);
 237             *list = &sei->message_list;
 238         }
 239         break;
 240     default:
 241         return AVERROR(EINVAL);
 242     }
 243
 244     return 0;
 245 }
 246
 247 int ff_cbs_sei_add_message(CodedBitstreamContext *ctx,
 248                            CodedBitstreamFragment *au,
 249                            int prefix,
 250                            uint32_t     payload_type,
 251                            void        *payload_data,
 252                            AVBufferRef *payload_buf)
 253 {
 254     const SEIMessageTypeDescriptor *desc;
 255     CodedBitstreamUnit *unit;
 256     SEIRawMessageList *list;
 257     SEIRawMessage *message;
 258     AVBufferRef *payload_ref;
 259     int err;
 260
 261     desc = ff_cbs_sei_find_type(ctx, payload_type);
 262     if (!desc)
 263         return AVERROR(EINVAL);
 264
 265     // Find an existing SEI unit or make a new one to add to.
 266     err = cbs_sei_get_unit(ctx, au, prefix, &unit);
 267     if (err < 0)
 268         return err;
 269
 270     // Find the message list inside the codec-dependent unit.
 271     err = cbs_sei_get_message_list(ctx, unit, &list);
 272     if (err < 0)
 273         return err;
 274
 275     // Add a new message to the message list.
 276     err = ff_cbs_sei_list_add(list);
 277     if (err < 0)
 278         return err;
 279
 280     if (payload_buf) {
 281         payload_ref = av_buffer_ref(payload_buf);
 282         if (!payload_ref)
 283             return AVERROR(ENOMEM);
 284     } else {
 285         payload_ref = NULL;
 286     }
 287
 288     message = &list->messages[list->nb_messages - 1];
 289
 290     message->payload_type = payload_type;
 291     message->payload      = payload_data;
 292     message->payload_ref  = payload_ref;
 293
 294     return 0;
 295 }
 296
 297 int ff_cbs_sei_find_message(CodedBitstreamContext *ctx,
 298                             CodedBitstreamFragment *au,
 299                             uint32_t payload_type,
 300                             SEIRawMessage **iter)
 301 {
 302     int err, i, j, found;
 303
 304     found = 0;
 305     for (i = 0; i < au->nb_units; i++) {
 306         CodedBitstreamUnit *unit = &au->units[i];
 307         SEIRawMessageList *list;
 308
 309         err = cbs_sei_get_message_list(ctx, unit, &list);
 310         if (err < 0)
 311             continue;
 312
 313         for (j = 0; j < list->nb_messages; j++) {
 314             SEIRawMessage *message = &list->messages[j];
 315
 316             if (message->payload_type == payload_type) {
 317                 if (!*iter || found) {
 318                     *iter = message;
 319                     return 0;
 320                 }
 321                 if (message == *iter)
 322                     found = 1;
 323             }
 324         }
 325     }
 326
 327     return AVERROR(ENOENT);
 328 }
 329
 330 static void cbs_sei_delete_message(SEIRawMessageList *list,
 331                                    int position)
 332 {
 333     SEIRawMessage *message;
 334
 335     av_assert0(0 <= position && position < list->nb_messages);
 336
 337     message = &list->messages[position];
 338     av_buffer_unref(&message->payload_ref);
 339     av_buffer_unref(&message->extension_data_ref);
 340
 341     --list->nb_messages;
 342
 343     if (list->nb_messages > 0) {
 344         memmove(list->messages + position,
 345                 list->messages + position + 1,
 346                 (list->nb_messages - position) * sizeof(*list->messages));
 347     }
 348 }
 349
 350 void ff_cbs_sei_delete_message_type(CodedBitstreamContext *ctx,
 351                                     CodedBitstreamFragment *au,
 352                                     uint32_t payload_type)
 353 {
 354     int err, i, j;
 355
 356     for (i = 0; i < au->nb_units; i++) {
 357         CodedBitstreamUnit *unit = &au->units[i];
 358         SEIRawMessageList *list;
 359
 360         err = cbs_sei_get_message_list(ctx, unit, &list);
 361         if (err < 0)
 362             continue;
 363
 364         for (j = list->nb_messages - 1; j >= 0; j--) {
 365             if (list->messages[j].payload_type == payload_type)
 366                 cbs_sei_delete_message(list, j);
 367         }
 368     }
 369 }