xref: /third_party/ffmpeg/libavcodec/cbs_sei.c (revision cabdff1a) 
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
25static 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
32static 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
39int 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
74int 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
99void 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
109static 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
218static 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
247int 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
297int 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
330static 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
350void 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}
370