vp9: Add bsf to merge superframes
[ffmpeg.git] / libavcodec / vp9_superframe_bsf.c
1 /*
2  * VP9 invisible (alt-ref) frame to superframe merge bitstream filter
3  * Copyright (c) 2016 Ronald S. Bultje <rsbultje@gmail.com>
4  *
5  * This file is part of Libav.
6  *
7  * Libav is free software; you can redistribute it and/or
8  * modify it under the terms of the GNU Lesser General Public
9  * License as published by the Free Software Foundation; either
10  * version 2.1 of the License, or (at your option) any later version.
11  *
12  * Libav is distributed in the hope that it will be useful,
13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
15  * Lesser General Public License for more details.
16  *
17  * You should have received a copy of the GNU Lesser General Public
18  * License along with Libav; if not, write to the Free Software
19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
20  */
21
22 #include "libavutil/avassert.h"
23 #include "avcodec.h"
24 #include "bsf.h"
25 #include "get_bits.h"
26
27 #define MAX_CACHE 8
28 typedef struct VP9BSFContext {
29     int n_cache;
30     struct CachedBuf {
31         uint8_t *data;
32         int size;
33     } cache[MAX_CACHE];
34 } VP9BSFContext;
35
36 static void stats(const struct CachedBuf *in, int n_in,
37                   unsigned *_max, unsigned *_sum)
38 {
39     int n;
40     unsigned max = 0, sum = 0;
41
42     for (n = 0; n < n_in; n++) {
43         unsigned sz = in[n].size;
44
45         if (sz > max)
46             max = sz;
47         sum += sz;
48     }
49
50     *_max = max;
51     *_sum = sum;
52 }
53
54 static int merge_superframe(const struct CachedBuf *in, int n_in, AVPacket *out)
55 {
56     unsigned max, sum, mag, marker, n, sz;
57     uint8_t *ptr;
58     int res;
59
60     stats(in, n_in, &max, &sum);
61     mag = av_log2(max) >> 3;
62     marker = 0xC0 + (mag << 3) + (n_in - 1);
63     sz = sum + 2 + (mag + 1) * n_in;
64     res = av_new_packet(out, sz);
65     if (res < 0)
66         return res;
67     ptr = out->data;
68     for (n = 0; n < n_in; n++) {
69         memcpy(ptr, in[n].data, in[n].size);
70         ptr += in[n].size;
71     }
72
73 #define wloop(mag, wr) do { \
74         for (n = 0; n < n_in; n++) { \
75             wr; \
76             ptr += mag + 1; \
77         } \
78     } while (0)
79
80     // write superframe with marker 110[mag:2][nframes:3]
81     *ptr++ = marker;
82     switch (mag) {
83     case 0:
84         wloop(mag, *ptr = in[n].size);
85         break;
86     case 1:
87         wloop(mag, AV_WL16(ptr, in[n].size));
88         break;
89     case 2:
90         wloop(mag, AV_WL24(ptr, in[n].size));
91         break;
92     case 3:
93         wloop(mag, AV_WL32(ptr, in[n].size));
94         break;
95     }
96     *ptr++ = marker;
97     av_assert0(ptr == &out->data[out->size]);
98
99     return 0;
100 }
101
102 static int vp9_superframe_filter(AVBSFContext *ctx, AVPacket *out)
103 {
104     GetBitContext gb;
105     VP9BSFContext *s = ctx->priv_data;
106     AVPacket *in;
107     int res, invisible, profile, marker, uses_superframe_syntax = 0, n;
108
109     res = ff_bsf_get_packet(ctx, &in);
110     if (res < 0)
111         return res;
112
113     marker = in->data[in->size - 1];
114     if ((marker & 0xe0) == 0xc0) {
115         int nbytes = 1 + ((marker >> 3) & 0x3);
116         int n_frames = 1 + (marker & 0x7), idx_sz = 2 + n_frames * nbytes;
117
118         uses_superframe_syntax = in->size >= idx_sz && in->data[in->size - idx_sz] == marker;
119     }
120
121     if ((res = init_get_bits8(&gb, in->data, in->size)) < 0)
122         goto done;
123
124     get_bits(&gb, 2); // frame marker
125     profile  = get_bits1(&gb);
126     profile |= get_bits1(&gb) << 1;
127     if (profile == 3) profile += get_bits1(&gb);
128
129     if (get_bits1(&gb)) {
130         invisible = 0;
131     } else {
132         get_bits1(&gb); // keyframe
133         invisible = !get_bits1(&gb);
134     }
135
136     if (uses_superframe_syntax && s->n_cache > 0) {
137         av_log(ctx, AV_LOG_ERROR,
138                "Mixing of superframe syntax and naked VP9 frames not supported");
139         res = AVERROR(ENOSYS);
140         goto done;
141     } else if ((!invisible || uses_superframe_syntax) && !s->n_cache) {
142         // passthrough
143         av_packet_move_ref(out, in);
144         goto done;
145     } else if (s->n_cache + 1 >= MAX_CACHE) {
146         av_log(ctx, AV_LOG_ERROR,
147                "Too many invisible frames");
148         res = AVERROR_INVALIDDATA;
149         goto done;
150     }
151
152     if (invisible) {
153         s->cache[s->n_cache].data = av_malloc(in->size);
154         if (!s->cache[s->n_cache].data) {
155             res = AVERROR(ENOMEM);
156             goto done;
157         }
158         memcpy(s->cache[s->n_cache].data, in->data, in->size);
159         s->cache[s->n_cache++].size = in->size;
160         res = AVERROR(EAGAIN);
161         goto done;
162     }
163     av_assert0(s->n_cache > 0);
164
165     s->cache[s->n_cache].data = in->data;
166     s->cache[s->n_cache].size = in->size;
167
168     // build superframe
169     if ((res = merge_superframe(s->cache, s->n_cache + 1, out)) < 0)
170         goto done;
171
172     for (n = 0; n < s->n_cache; n++)
173         av_freep(&s->cache[n].data);
174     s->n_cache = 0;
175
176     res = av_packet_copy_props(out, in);
177     if (res < 0)
178         goto done;
179
180 done:
181     if (res < 0)
182         av_packet_unref(out);
183     av_packet_free(&in);
184     return res;
185 }
186
187 static void vp9_superframe_close(AVBSFContext *ctx)
188 {
189     VP9BSFContext *s = ctx->priv_data;
190     int n;
191
192     // free cached data
193     for (n = 0; n < s->n_cache; n++)
194         av_freep(&s->cache[n].data);
195 }
196
197 static const enum AVCodecID codec_ids[] = {
198     AV_CODEC_ID_VP9, AV_CODEC_ID_NONE,
199 };
200
201 const AVBitStreamFilter ff_vp9_superframe_bsf = {
202     .name           = "vp9_superframe",
203     .priv_data_size = sizeof(VP9BSFContext),
204     .filter         = vp9_superframe_filter,
205     .close          = vp9_superframe_close,
206     .codec_ids      = codec_ids,
207 };