Support planar sample formats in audio mixdown
[blender.git] / intern / audaspace / ffmpeg / AUD_FFMPEGWriter.cpp
1 /*
2  * ***** BEGIN GPL LICENSE BLOCK *****
3  *
4  * Copyright 2009-2011 Jörg Hermann Müller
5  *
6  * This file is part of AudaSpace.
7  *
8  * Audaspace is free software; you can redistribute it and/or modify
9  * it under the terms of the GNU General Public License as published by
10  * the Free Software Foundation; either version 2 of the License, or
11  * (at your option) any later version.
12  *
13  * AudaSpace is distributed in the hope that it will be useful,
14  * but WITHOUT ANY WARRANTY; without even the implied warranty of
15  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
16  * GNU General Public License for more details.
17  *
18  * You should have received a copy of the GNU General Public License
19  * along with Audaspace; if not, write to the Free Software Foundation,
20  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
21  *
22  * ***** END GPL LICENSE BLOCK *****
23  */
24
25 /** \file audaspace/ffmpeg/AUD_FFMPEGWriter.cpp
26  *  \ingroup audffmpeg
27  */
28
29
30 // needed for INT64_C
31 #ifndef __STDC_CONSTANT_MACROS
32 #define __STDC_CONSTANT_MACROS
33 #endif
34
35 #include "AUD_FFMPEGWriter.h"
36
37 extern "C" {
38 #include <libavcodec/avcodec.h>
39 #include <libavformat/avformat.h>
40 #include <libavformat/avio.h>
41 #include "ffmpeg_compat.h"
42 }
43
44 static const char* context_error = "AUD_FFMPEGWriter: Couldn't allocate context.";
45 static const char* codec_error = "AUD_FFMPEGWriter: Invalid codec or codec not found.";
46 static const char* stream_error = "AUD_FFMPEGWriter: Couldn't allocate stream.";
47 static const char* format_error = "AUD_FFMPEGWriter: Unsupported sample format.";
48 static const char* file_error = "AUD_FFMPEGWriter: File couldn't be written.";
49 static const char* write_error = "AUD_FFMPEGWriter: Error writing packet.";
50
51 AUD_FFMPEGWriter::AUD_FFMPEGWriter(std::string filename, AUD_DeviceSpecs specs, AUD_Container format, AUD_Codec codec, unsigned int bitrate) :
52         m_position(0),
53         m_specs(specs),
54         m_input_samples(0)
55 {
56         static const char* formats[] = { NULL, "ac3", "flac", "matroska", "mp2", "mp3", "ogg", "wav" };
57
58         m_formatCtx = avformat_alloc_context();
59         if (!m_formatCtx) AUD_THROW(AUD_ERROR_FFMPEG, context_error);
60
61         strcpy(m_formatCtx->filename, filename.c_str());
62         m_outputFmt = m_formatCtx->oformat = av_guess_format(formats[format], filename.c_str(), NULL);
63         if (!m_outputFmt) {
64                 avformat_free_context(m_formatCtx);
65                 AUD_THROW(AUD_ERROR_FFMPEG, context_error);
66         }
67
68         switch(codec)
69         {
70         case AUD_CODEC_AAC:
71                 m_outputFmt->audio_codec = AV_CODEC_ID_AAC;
72                 break;
73         case AUD_CODEC_AC3:
74                 m_outputFmt->audio_codec = AV_CODEC_ID_AC3;
75                 break;
76         case AUD_CODEC_FLAC:
77                 m_outputFmt->audio_codec = AV_CODEC_ID_FLAC;
78                 break;
79         case AUD_CODEC_MP2:
80                 m_outputFmt->audio_codec = AV_CODEC_ID_MP2;
81                 break;
82         case AUD_CODEC_MP3:
83                 m_outputFmt->audio_codec = AV_CODEC_ID_MP3;
84                 break;
85         case AUD_CODEC_PCM:
86                 switch(specs.format)
87                 {
88                 case AUD_FORMAT_U8:
89                         m_outputFmt->audio_codec = AV_CODEC_ID_PCM_U8;
90                         break;
91                 case AUD_FORMAT_S16:
92                         m_outputFmt->audio_codec = AV_CODEC_ID_PCM_S16LE;
93                         break;
94                 case AUD_FORMAT_S24:
95                         m_outputFmt->audio_codec = AV_CODEC_ID_PCM_S24LE;
96                         break;
97                 case AUD_FORMAT_S32:
98                         m_outputFmt->audio_codec = AV_CODEC_ID_PCM_S32LE;
99                         break;
100                 case AUD_FORMAT_FLOAT32:
101                         m_outputFmt->audio_codec = AV_CODEC_ID_PCM_F32LE;
102                         break;
103                 case AUD_FORMAT_FLOAT64:
104                         m_outputFmt->audio_codec = AV_CODEC_ID_PCM_F64LE;
105                         break;
106                 default:
107                         m_outputFmt->audio_codec = AV_CODEC_ID_NONE;
108                         break;
109                 }
110                 break;
111         case AUD_CODEC_VORBIS:
112                 m_outputFmt->audio_codec = AV_CODEC_ID_VORBIS;
113                 break;
114         default:
115                 m_outputFmt->audio_codec = AV_CODEC_ID_NONE;
116                 break;
117         }
118
119         try
120         {
121                 if(m_outputFmt->audio_codec == AV_CODEC_ID_NONE)
122                         AUD_THROW(AUD_ERROR_SPECS, codec_error);
123
124                 m_stream = avformat_new_stream(m_formatCtx, NULL);
125                 if(!m_stream)
126                         AUD_THROW(AUD_ERROR_FFMPEG, stream_error);
127
128                 m_codecCtx = m_stream->codec;
129                 m_codecCtx->codec_id = m_outputFmt->audio_codec;
130                 m_codecCtx->codec_type = AVMEDIA_TYPE_AUDIO;
131                 m_codecCtx->bit_rate = bitrate;
132                 m_codecCtx->sample_rate = int(m_specs.rate);
133                 m_codecCtx->channels = m_specs.channels;
134                 m_codecCtx->time_base.num = 1;
135                 m_codecCtx->time_base.den = m_codecCtx->sample_rate;
136
137                 switch(m_specs.format)
138                 {
139                 case AUD_FORMAT_U8:
140                         m_convert = AUD_convert_float_u8;
141                         m_codecCtx->sample_fmt = AV_SAMPLE_FMT_U8;
142                         break;
143                 case AUD_FORMAT_S16:
144                         m_convert = AUD_convert_float_s16;
145                         m_codecCtx->sample_fmt = AV_SAMPLE_FMT_S16;
146                         break;
147                 case AUD_FORMAT_S32:
148                         m_convert = AUD_convert_float_s32;
149                         m_codecCtx->sample_fmt = AV_SAMPLE_FMT_S32;
150                         break;
151                 case AUD_FORMAT_FLOAT32:
152                         m_convert = AUD_convert_copy<float>;
153                         m_codecCtx->sample_fmt = AV_SAMPLE_FMT_FLT;
154                         break;
155                 case AUD_FORMAT_FLOAT64:
156                         m_convert = AUD_convert_float_double;
157                         m_codecCtx->sample_fmt = AV_SAMPLE_FMT_DBL;
158                         break;
159                 default:
160                         AUD_THROW(AUD_ERROR_FFMPEG, format_error);
161                 }
162
163                 try
164                 {
165                         if(m_formatCtx->oformat->flags & AVFMT_GLOBALHEADER)
166                                 m_codecCtx->flags |= CODEC_FLAG_GLOBAL_HEADER;
167
168                         AVCodec* codec = avcodec_find_encoder(m_codecCtx->codec_id);
169                         if(!codec)
170                                 AUD_THROW(AUD_ERROR_FFMPEG, codec_error);
171
172                         if(codec->sample_fmts) {
173                                 // Check if the prefered sample format for this codec is supported.
174                                 const enum AVSampleFormat *p = codec->sample_fmts;
175                                 for(; *p != -1; p++) {
176                                         if(*p == m_stream->codec->sample_fmt)
177                                                 break;
178                                 }
179                                 if(*p == -1) {
180                                         // Sample format incompatible with codec. Defaulting to a format known to work.
181                                         m_stream->codec->sample_fmt = codec->sample_fmts[0];
182                                 }
183                         }
184
185                         if(avcodec_open2(m_codecCtx, codec, NULL))
186                                 AUD_THROW(AUD_ERROR_FFMPEG, codec_error);
187
188                         m_output_buffer.resize(FF_MIN_BUFFER_SIZE);
189                         int samplesize = AUD_MAX(AUD_SAMPLE_SIZE(m_specs), AUD_DEVICE_SAMPLE_SIZE(m_specs));
190
191                         if(m_codecCtx->frame_size <= 1) {
192                                 m_input_size = FF_MIN_BUFFER_SIZE * 8 / m_codecCtx->bits_per_coded_sample / m_codecCtx->channels;
193                                 m_input_buffer.resize(m_input_size * samplesize);
194                         }
195                         else
196                         {
197                                 m_input_buffer.resize(m_codecCtx->frame_size * samplesize);
198                                 m_input_size = m_codecCtx->frame_size;
199                         }
200
201 #ifdef FFMPEG_HAVE_ENCODE_AUDIO2
202                         m_frame = av_frame_alloc();
203                         if (!m_frame)
204                                 AUD_THROW(AUD_ERROR_FFMPEG, codec_error);
205                         avcodec_get_frame_defaults(m_frame);
206                         m_frame->linesize[0]    = m_input_size * samplesize;
207                         m_frame->format         = m_codecCtx->sample_fmt;
208                         m_frame->nb_samples     = m_input_size;
209 #  ifdef FFMPEG_HAVE_AVFRAME_SAMPLE_RATE
210                         m_frame->sample_rate    = m_codecCtx->sample_rate;
211 #  endif
212 #  ifdef FFMPEG_HAVE_FRAME_CHANNEL_LAYOUT
213                         m_frame->channel_layout = m_codecCtx->channel_layout;
214 #  endif
215                         m_sample_size = av_get_bytes_per_sample(m_codecCtx->sample_fmt);
216                         m_frame_pts = 0;
217                         m_deinterleave = av_sample_fmt_is_planar(m_codecCtx->sample_fmt);
218                         if(m_deinterleave)
219                                 m_deinterleave_buffer.resize(m_input_size * m_codecCtx->channels * m_sample_size);
220 #endif
221
222                         try
223                         {
224                                 if(avio_open(&m_formatCtx->pb, filename.c_str(), AVIO_FLAG_WRITE))
225                                         AUD_THROW(AUD_ERROR_FILE, file_error);
226
227                                 avformat_write_header(m_formatCtx, NULL);
228                         }
229                         catch(AUD_Exception&)
230                         {
231                                 avcodec_close(m_codecCtx);
232                                 av_freep(&m_formatCtx->streams[0]->codec);
233                                 throw;
234                         }
235                 }
236                 catch(AUD_Exception&)
237                 {
238                         av_freep(&m_formatCtx->streams[0]);
239                         throw;
240                 }
241         }
242         catch(AUD_Exception&)
243         {
244                 av_free(m_formatCtx);
245                 throw;
246         }
247 }
248
249 AUD_FFMPEGWriter::~AUD_FFMPEGWriter()
250 {
251         // writte missing data
252         if(m_input_samples)
253         {
254                 sample_t* buf = m_input_buffer.getBuffer();
255                 memset(buf + m_specs.channels * m_input_samples, 0,
256                            (m_input_size - m_input_samples) * AUD_DEVICE_SAMPLE_SIZE(m_specs));
257
258                 encode(buf);
259         }
260
261         av_write_trailer(m_formatCtx);
262
263         avcodec_close(m_codecCtx);
264
265         av_freep(&m_formatCtx->streams[0]->codec);
266         av_freep(&m_formatCtx->streams[0]);
267
268 #ifdef FFMPEG_HAVE_ENCODE_AUDIO2
269         av_frame_free(&m_frame);
270 #endif
271
272         avio_close(m_formatCtx->pb);
273         av_free(m_formatCtx);
274 }
275
276 int AUD_FFMPEGWriter::getPosition() const
277 {
278         return m_position;
279 }
280
281 AUD_DeviceSpecs AUD_FFMPEGWriter::getSpecs() const
282 {
283         return m_specs;
284 }
285
286 void AUD_FFMPEGWriter::encode(sample_t* data)
287 {
288         // convert first
289         if(m_input_size)
290                 m_convert(reinterpret_cast<data_t*>(data), reinterpret_cast<data_t*>(data), m_input_size * m_specs.channels);
291
292         AVPacket packet = { 0 };
293         av_init_packet(&packet);
294
295 #ifdef FFMPEG_HAVE_ENCODE_AUDIO2
296         int got_output, ret;
297         m_frame->pts = m_frame_pts / av_q2d(m_codecCtx->time_base);
298         m_frame_pts++;
299 #ifdef FFMPEG_HAVE_FRAME_CHANNEL_LAYOUT
300         m_frame->channel_layout = m_codecCtx->channel_layout;
301 #endif
302
303         if(m_deinterleave) {
304                 for(int channel = 0; channel < m_codecCtx->channels; channel++) {
305                         for(int i = 0; i < m_frame->nb_samples; i++) {
306                                 memcpy(reinterpret_cast<uint8_t*>(m_deinterleave_buffer.getBuffer()) + (i + channel * m_frame->nb_samples) * m_sample_size,
307                                            reinterpret_cast<uint8_t*>(data) + (m_codecCtx->channels * i + channel) * m_sample_size, m_sample_size);
308                         }
309                 }
310
311                 data = m_deinterleave_buffer.getBuffer();
312         }
313
314         avcodec_fill_audio_frame(m_frame, m_codecCtx->channels, m_codecCtx->sample_fmt, reinterpret_cast<uint8_t*>(data),
315                                  m_frame->nb_samples * av_get_bytes_per_sample(m_codecCtx->sample_fmt) * m_codecCtx->channels, 1);
316
317         ret = avcodec_encode_audio2(m_codecCtx, &packet, m_frame, &got_output);
318         if(ret < 0)
319                 AUD_THROW(AUD_ERROR_FFMPEG, codec_error);
320
321         if(!got_output)
322                 return;
323 #else
324         sample_t* outbuf = m_output_buffer.getBuffer();
325
326         packet.size = avcodec_encode_audio(m_codecCtx, reinterpret_cast<uint8_t*>(outbuf), m_output_buffer.getSize(), reinterpret_cast<short*>(data));
327         if(m_codecCtx->coded_frame && m_codecCtx->coded_frame->pts != AV_NOPTS_VALUE)
328                 packet.pts = av_rescale_q(m_codecCtx->coded_frame->pts, m_codecCtx->time_base, m_stream->time_base);
329         packet.flags |= AV_PKT_FLAG_KEY;
330         packet.data = reinterpret_cast<uint8_t*>(outbuf);
331 #endif
332
333         if(packet.pts != AV_NOPTS_VALUE)
334                 packet.pts = av_rescale_q(packet.pts, m_codecCtx->time_base, m_stream->time_base);
335         if(packet.dts != AV_NOPTS_VALUE)
336                 packet.dts = av_rescale_q(packet.dts, m_codecCtx->time_base, m_stream->time_base);
337         if(packet.duration > 0)
338                 packet.duration = av_rescale_q(packet.duration, m_codecCtx->time_base, m_stream->time_base);
339
340         packet.stream_index = m_stream->index;
341
342         packet.flags |= AV_PKT_FLAG_KEY;
343
344         if(av_interleaved_write_frame(m_formatCtx, &packet)) {
345                 av_free_packet(&packet);
346                 AUD_THROW(AUD_ERROR_FFMPEG, write_error);
347         }
348
349         av_free_packet(&packet);
350 }
351
352 void AUD_FFMPEGWriter::write(unsigned int length, sample_t* buffer)
353 {
354         unsigned int samplesize = AUD_SAMPLE_SIZE(m_specs);
355
356         if(m_input_size)
357         {
358                 sample_t* inbuf = m_input_buffer.getBuffer();
359
360                 while(length)
361                 {
362                         unsigned int len = AUD_MIN(m_input_size - m_input_samples, length);
363
364                         memcpy(inbuf + m_input_samples * m_specs.channels, buffer, len * samplesize);
365
366                         buffer += len * m_specs.channels;
367                         m_input_samples += len;
368                         m_position += len;
369                         length -= len;
370
371                         if(m_input_samples == m_input_size)
372                         {
373                                 encode(inbuf);
374
375                                 m_input_samples = 0;
376                         }
377                 }
378         }
379         else // PCM data, can write directly!
380         {
381                 int samplesize = AUD_SAMPLE_SIZE(m_specs);
382                 if(m_output_buffer.getSize() != length * m_specs.channels * m_codecCtx->bits_per_coded_sample / 8)
383                         m_output_buffer.resize(length * m_specs.channels * m_codecCtx->bits_per_coded_sample / 8);
384                 m_input_buffer.assureSize(length * AUD_MAX(AUD_DEVICE_SAMPLE_SIZE(m_specs), samplesize));
385
386                 sample_t* buf = m_input_buffer.getBuffer();
387                 m_convert(reinterpret_cast<data_t*>(buf), reinterpret_cast<data_t*>(buffer), length * m_specs.channels);
388
389                 encode(buf);
390
391                 m_position += length;
392         }
393 }