54e83b66ad028af1b4ea73e8e71d7de364b17dae
[ffmpeg.git] / libavcodec / wmaenc.c
1 /*
2  * WMA compatible encoder
3  * Copyright (c) 2007 Michael Niedermayer
4  *
5  * This file is part of FFmpeg.
6  *
7  * FFmpeg 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  * FFmpeg 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 FFmpeg; 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/attributes.h"
23 #include "libavutil/internal.h"
24
25 #include "avcodec.h"
26 #include "internal.h"
27 #include "wma.h"
28 #include "libavutil/avassert.h"
29
30
31 static av_cold int encode_init(AVCodecContext *avctx)
32 {
33     WMACodecContext *s = avctx->priv_data;
34     int i, flags1, flags2, block_align;
35     uint8_t *extradata;
36
37     s->avctx = avctx;
38
39     if (avctx->channels > MAX_CHANNELS) {
40         av_log(avctx, AV_LOG_ERROR,
41                "too many channels: got %i, need %i or fewer\n",
42                avctx->channels, MAX_CHANNELS);
43         return AVERROR(EINVAL);
44     }
45
46     if (avctx->sample_rate > 48000) {
47         av_log(avctx, AV_LOG_ERROR, "sample rate is too high: %d > 48kHz\n",
48                avctx->sample_rate);
49         return AVERROR(EINVAL);
50     }
51
52     if (avctx->bit_rate < 24 * 1000) {
53         av_log(avctx, AV_LOG_ERROR,
54                "bitrate too low: got %"PRId64", need 24000 or higher\n",
55                (int64_t)avctx->bit_rate);
56         return AVERROR(EINVAL);
57     }
58
59     /* extract flag infos */
60     flags1 = 0;
61     flags2 = 1;
62     if (avctx->codec->id == AV_CODEC_ID_WMAV1) {
63         extradata             = av_malloc(4);
64         if (!extradata)
65             return AVERROR(ENOMEM);
66         avctx->extradata_size = 4;
67         AV_WL16(extradata, flags1);
68         AV_WL16(extradata + 2, flags2);
69     } else if (avctx->codec->id == AV_CODEC_ID_WMAV2) {
70         extradata             = av_mallocz(10);
71         if (!extradata)
72             return AVERROR(ENOMEM);
73         avctx->extradata_size = 10;
74         AV_WL32(extradata, flags1);
75         AV_WL16(extradata + 4, flags2);
76     } else {
77         av_assert0(0);
78     }
79     avctx->extradata          = extradata;
80     s->use_exp_vlc            = flags2 & 0x0001;
81     s->use_bit_reservoir      = flags2 & 0x0002;
82     s->use_variable_block_len = flags2 & 0x0004;
83     if (avctx->channels == 2)
84         s->ms_stereo = 1;
85
86     ff_wma_init(avctx, flags2);
87
88     /* init MDCT */
89     for (i = 0; i < s->nb_block_sizes; i++)
90         ff_mdct_init(&s->mdct_ctx[i], s->frame_len_bits - i + 1, 0, 1.0);
91
92     block_align        = avctx->bit_rate * (int64_t) s->frame_len /
93                          (avctx->sample_rate * 8);
94     block_align        = FFMIN(block_align, MAX_CODED_SUPERFRAME_SIZE);
95     avctx->block_align = block_align;
96     avctx->frame_size = avctx->initial_padding = s->frame_len;
97
98     return 0;
99 }
100
101 static void apply_window_and_mdct(AVCodecContext *avctx, const AVFrame *frame)
102 {
103     WMACodecContext *s = avctx->priv_data;
104     float **audio      = (float **) frame->extended_data;
105     int len            = frame->nb_samples;
106     int window_index   = s->frame_len_bits - s->block_len_bits;
107     FFTContext *mdct   = &s->mdct_ctx[window_index];
108     int ch;
109     const float *win   = s->windows[window_index];
110     int window_len     = 1 << s->block_len_bits;
111     float n            = 2.0 * 32768.0 / window_len;
112
113     for (ch = 0; ch < avctx->channels; ch++) {
114         memcpy(s->output, s->frame_out[ch], window_len * sizeof(*s->output));
115         s->fdsp->vector_fmul_scalar(s->frame_out[ch], audio[ch], n, len);
116         s->fdsp->vector_fmul_reverse(&s->output[window_len], s->frame_out[ch],
117                                     win, len);
118         s->fdsp->vector_fmul(s->frame_out[ch], s->frame_out[ch], win, len);
119         mdct->mdct_calc(mdct, s->coefs[ch], s->output);
120     }
121 }
122
123 // FIXME use for decoding too
124 static void init_exp(WMACodecContext *s, int ch, const int *exp_param)
125 {
126     int n;
127     const uint16_t *ptr;
128     float v, *q, max_scale, *q_end;
129
130     ptr       = s->exponent_bands[s->frame_len_bits - s->block_len_bits];
131     q         = s->exponents[ch];
132     q_end     = q + s->block_len;
133     max_scale = 0;
134     while (q < q_end) {
135         /* XXX: use a table */
136         v         = ff_exp10(*exp_param++ *(1.0 / 16.0));
137         max_scale = FFMAX(max_scale, v);
138         n         = *ptr++;
139         do {
140             *q++ = v;
141         } while (--n);
142     }
143     s->max_exponent[ch] = max_scale;
144 }
145
146 static void encode_exp_vlc(WMACodecContext *s, int ch, const int *exp_param)
147 {
148     int last_exp;
149     const uint16_t *ptr;
150     float *q, *q_end;
151
152     ptr   = s->exponent_bands[s->frame_len_bits - s->block_len_bits];
153     q     = s->exponents[ch];
154     q_end = q + s->block_len;
155     if (s->version == 1) {
156         last_exp = *exp_param++;
157         av_assert0(last_exp - 10 >= 0 && last_exp - 10 < 32);
158         put_bits(&s->pb, 5, last_exp - 10);
159         q += *ptr++;
160     } else
161         last_exp = 36;
162     while (q < q_end) {
163         int exp  = *exp_param++;
164         int code = exp - last_exp + 60;
165         av_assert1(code >= 0 && code < 120);
166         put_bits(&s->pb, ff_aac_scalefactor_bits[code],
167                  ff_aac_scalefactor_code[code]);
168         /* XXX: use a table */
169         q       += *ptr++;
170         last_exp = exp;
171     }
172 }
173
174 static int encode_block(WMACodecContext *s, float (*src_coefs)[BLOCK_MAX_SIZE],
175                         int total_gain)
176 {
177     int v, bsize, ch, coef_nb_bits, parse_exponents;
178     float mdct_norm;
179     int nb_coefs[MAX_CHANNELS];
180     static const int fixed_exp[25] = {
181         20, 20, 20, 20, 20,
182         20, 20, 20, 20, 20,
183         20, 20, 20, 20, 20,
184         20, 20, 20, 20, 20,
185         20, 20, 20, 20, 20
186     };
187
188     // FIXME remove duplication relative to decoder
189     if (s->use_variable_block_len) {
190         av_assert0(0); // FIXME not implemented
191     } else {
192         /* fixed block len */
193         s->next_block_len_bits = s->frame_len_bits;
194         s->prev_block_len_bits = s->frame_len_bits;
195         s->block_len_bits      = s->frame_len_bits;
196     }
197
198     s->block_len = 1 << s->block_len_bits;
199 //     av_assert0((s->block_pos + s->block_len) <= s->frame_len);
200     bsize = s->frame_len_bits - s->block_len_bits;
201
202     // FIXME factor
203     v = s->coefs_end[bsize] - s->coefs_start;
204     for (ch = 0; ch < s->avctx->channels; ch++)
205         nb_coefs[ch] = v;
206     {
207         int n4 = s->block_len / 2;
208         mdct_norm = 1.0 / (float) n4;
209         if (s->version == 1)
210             mdct_norm *= sqrt(n4);
211     }
212
213     if (s->avctx->channels == 2)
214         put_bits(&s->pb, 1, !!s->ms_stereo);
215
216     for (ch = 0; ch < s->avctx->channels; ch++) {
217         // FIXME only set channel_coded when needed, instead of always
218         s->channel_coded[ch] = 1;
219         if (s->channel_coded[ch])
220             init_exp(s, ch, fixed_exp);
221     }
222
223     for (ch = 0; ch < s->avctx->channels; ch++) {
224         if (s->channel_coded[ch]) {
225             WMACoef *coefs1;
226             float *coefs, *exponents, mult;
227             int i, n;
228
229             coefs1    = s->coefs1[ch];
230             exponents = s->exponents[ch];
231             mult      = ff_exp10(total_gain * 0.05) / s->max_exponent[ch];
232             mult     *= mdct_norm;
233             coefs     = src_coefs[ch];
234             if (s->use_noise_coding && 0) {
235                 av_assert0(0); // FIXME not implemented
236             } else {
237                 coefs += s->coefs_start;
238                 n      = nb_coefs[ch];
239                 for (i = 0; i < n; i++) {
240                     double t = *coefs++ / (exponents[i] * mult);
241                     if (t < -32768 || t > 32767)
242                         return -1;
243
244                     coefs1[i] = lrint(t);
245                 }
246             }
247         }
248     }
249
250     v = 0;
251     for (ch = 0; ch < s->avctx->channels; ch++) {
252         int a = s->channel_coded[ch];
253         put_bits(&s->pb, 1, a);
254         v |= a;
255     }
256
257     if (!v)
258         return 1;
259
260     for (v = total_gain - 1; v >= 127; v -= 127)
261         put_bits(&s->pb, 7, 127);
262     put_bits(&s->pb, 7, v);
263
264     coef_nb_bits = ff_wma_total_gain_to_bits(total_gain);
265
266     if (s->use_noise_coding) {
267         for (ch = 0; ch < s->avctx->channels; ch++) {
268             if (s->channel_coded[ch]) {
269                 int i, n;
270                 n = s->exponent_high_sizes[bsize];
271                 for (i = 0; i < n; i++) {
272                     put_bits(&s->pb, 1, s->high_band_coded[ch][i] = 0);
273                     if (0)
274                         nb_coefs[ch] -= s->exponent_high_bands[bsize][i];
275                 }
276             }
277         }
278     }
279
280     parse_exponents = 1;
281     if (s->block_len_bits != s->frame_len_bits)
282         put_bits(&s->pb, 1, parse_exponents);
283
284     if (parse_exponents) {
285         for (ch = 0; ch < s->avctx->channels; ch++) {
286             if (s->channel_coded[ch]) {
287                 if (s->use_exp_vlc) {
288                     encode_exp_vlc(s, ch, fixed_exp);
289                 } else {
290                     av_assert0(0); // FIXME not implemented
291 //                    encode_exp_lsp(s, ch);
292                 }
293             }
294         }
295     } else
296         av_assert0(0); // FIXME not implemented
297
298     for (ch = 0; ch < s->avctx->channels; ch++) {
299         if (s->channel_coded[ch]) {
300             int run, tindex;
301             WMACoef *ptr, *eptr;
302             tindex = (ch == 1 && s->ms_stereo);
303             ptr    = &s->coefs1[ch][0];
304             eptr   = ptr + nb_coefs[ch];
305
306             run = 0;
307             for (; ptr < eptr; ptr++) {
308                 if (*ptr) {
309                     int level     = *ptr;
310                     int abs_level = FFABS(level);
311                     int code      = 0;
312                     if (abs_level <= s->coef_vlcs[tindex]->max_level)
313                         if (run < s->coef_vlcs[tindex]->levels[abs_level - 1])
314                             code = run + s->int_table[tindex][abs_level - 1];
315
316                     av_assert2(code < s->coef_vlcs[tindex]->n);
317                     put_bits(&s->pb, s->coef_vlcs[tindex]->huffbits[code],
318                              s->coef_vlcs[tindex]->huffcodes[code]);
319
320                     if (code == 0) {
321                         if (1 << coef_nb_bits <= abs_level)
322                             return -1;
323
324                         put_bits(&s->pb, coef_nb_bits, abs_level);
325                         put_bits(&s->pb, s->frame_len_bits, run);
326                     }
327                     // FIXME the sign is flipped somewhere
328                     put_bits(&s->pb, 1, level < 0);
329                     run = 0;
330                 } else
331                     run++;
332             }
333             if (run)
334                 put_bits(&s->pb, s->coef_vlcs[tindex]->huffbits[1],
335                          s->coef_vlcs[tindex]->huffcodes[1]);
336         }
337         if (s->version == 1 && s->avctx->channels >= 2)
338             avpriv_align_put_bits(&s->pb);
339     }
340     return 0;
341 }
342
343 static int encode_frame(WMACodecContext *s, float (*src_coefs)[BLOCK_MAX_SIZE],
344                         uint8_t *buf, int buf_size, int total_gain)
345 {
346     init_put_bits(&s->pb, buf, buf_size);
347
348     if (s->use_bit_reservoir)
349         av_assert0(0); // FIXME not implemented
350     else if (encode_block(s, src_coefs, total_gain) < 0)
351         return INT_MAX;
352
353     avpriv_align_put_bits(&s->pb);
354
355     return put_bits_count(&s->pb) / 8 - s->avctx->block_align;
356 }
357
358 static int encode_superframe(AVCodecContext *avctx, AVPacket *avpkt,
359                              const AVFrame *frame, int *got_packet_ptr)
360 {
361     WMACodecContext *s = avctx->priv_data;
362     int i, total_gain, ret, error;
363
364     s->block_len_bits = s->frame_len_bits; // required by non variable block len
365     s->block_len      = 1 << s->block_len_bits;
366
367     apply_window_and_mdct(avctx, frame);
368
369     if (s->ms_stereo) {
370         float a, b;
371         int i;
372
373         for (i = 0; i < s->block_len; i++) {
374             a              = s->coefs[0][i] * 0.5;
375             b              = s->coefs[1][i] * 0.5;
376             s->coefs[0][i] = a + b;
377             s->coefs[1][i] = a - b;
378         }
379     }
380
381     if ((ret = ff_alloc_packet2(avctx, avpkt, 2 * MAX_CODED_SUPERFRAME_SIZE, 0)) < 0)
382         return ret;
383
384     total_gain = 128;
385     for (i = 64; i; i >>= 1) {
386         error = encode_frame(s, s->coefs, avpkt->data, avpkt->size,
387                                  total_gain - i);
388         if (error <= 0)
389             total_gain -= i;
390     }
391
392     while(total_gain <= 128 && error > 0)
393         error = encode_frame(s, s->coefs, avpkt->data, avpkt->size, total_gain++);
394     if (error > 0) {
395         av_log(avctx, AV_LOG_ERROR, "Invalid input data or requested bitrate too low, cannot encode\n");
396         avpkt->size = 0;
397         return AVERROR(EINVAL);
398     }
399     av_assert0((put_bits_count(&s->pb) & 7) == 0);
400     i= avctx->block_align - (put_bits_count(&s->pb)+7)/8;
401     av_assert0(i>=0);
402     while(i--)
403         put_bits(&s->pb, 8, 'N');
404
405     flush_put_bits(&s->pb);
406     av_assert0(put_bits_ptr(&s->pb) - s->pb.buf == avctx->block_align);
407
408     if (frame->pts != AV_NOPTS_VALUE)
409         avpkt->pts = frame->pts - ff_samples_to_time_base(avctx, avctx->initial_padding);
410
411     avpkt->size     = avctx->block_align;
412     *got_packet_ptr = 1;
413     return 0;
414 }
415
416 #if CONFIG_WMAV1_ENCODER
417 AVCodec ff_wmav1_encoder = {
418     .name           = "wmav1",
419     .long_name      = NULL_IF_CONFIG_SMALL("Windows Media Audio 1"),
420     .type           = AVMEDIA_TYPE_AUDIO,
421     .id             = AV_CODEC_ID_WMAV1,
422     .priv_data_size = sizeof(WMACodecContext),
423     .init           = encode_init,
424     .encode2        = encode_superframe,
425     .close          = ff_wma_end,
426     .sample_fmts    = (const enum AVSampleFormat[]) { AV_SAMPLE_FMT_FLTP,
427                                                       AV_SAMPLE_FMT_NONE },
428 };
429 #endif
430 #if CONFIG_WMAV2_ENCODER
431 AVCodec ff_wmav2_encoder = {
432     .name           = "wmav2",
433     .long_name      = NULL_IF_CONFIG_SMALL("Windows Media Audio 2"),
434     .type           = AVMEDIA_TYPE_AUDIO,
435     .id             = AV_CODEC_ID_WMAV2,
436     .priv_data_size = sizeof(WMACodecContext),
437     .init           = encode_init,
438     .encode2        = encode_superframe,
439     .close          = ff_wma_end,
440     .sample_fmts    = (const enum AVSampleFormat[]) { AV_SAMPLE_FMT_FLTP,
441                                                       AV_SAMPLE_FMT_NONE },
442 };
443 #endif