1/*
2 * Copyright 2009-2010, Stephan A��mus <superstippi@gmx.de>
3 * All rights reserved. Distributed under the terms of the GNU L-GPL license.
4 */
5
6#include "AVFormatWriter.h"
7
8#include <stdio.h>
9#include <string.h>
10#include <stdlib.h>
11
12#include <new>
13
14#include <Application.h>
15#include <AutoDeleter.h>
16#include <Autolock.h>
17#include <ByteOrder.h>
18#include <DataIO.h>
19#include <MediaDefs.h>
20#include <MediaFormats.h>
21#include <Roster.h>
22
23extern "C" {
24	#include "avformat.h"
25}
26
27#include "DemuxerTable.h"
28#include "EncoderTable.h"
29#include "gfx_util.h"
30
31
32//#define TRACE_AVFORMAT_WRITER
33#ifdef TRACE_AVFORMAT_WRITER
34#	define TRACE printf
35#	define TRACE_IO(a...)
36#	define TRACE_PACKET printf
37#else
38#	define TRACE(a...)
39#	define TRACE_IO(a...)
40#	define TRACE_PACKET(a...)
41#endif
42
43#define ERROR(a...) fprintf(stderr, a)
44
45
46static const size_t kIOBufferSize = 64 * 1024;
47	// TODO: This could depend on the BMediaFile creation flags, IIRC,
48	// they allow to specify a buffering mode.
49
50// NOTE: The following works around some weird bug in libavformat. We
51// have to open the AVFormatContext->AVStream->AVCodecContext, even though
52// we are not interested in donig any encoding here!!
53#define OPEN_CODEC_CONTEXT 1
54#define GET_CONTEXT_DEFAULTS 0
55
56
57// #pragma mark - AVFormatWriter::StreamCookie
58
59
60class AVFormatWriter::StreamCookie {
61public:
62								StreamCookie(AVFormatContext* context,
63									BLocker* streamLock);
64	virtual						~StreamCookie();
65
66			status_t			Init(media_format* format,
67									const media_codec_info* codecInfo);
68
69			status_t			WriteChunk(const void* chunkBuffer,
70									size_t chunkSize,
71									media_encode_info* encodeInfo);
72
73			status_t			AddTrackInfo(uint32 code, const void* data,
74									size_t size, uint32 flags);
75
76private:
77			AVFormatContext*	fContext;
78			AVStream*			fStream;
79			AVPacket			fPacket;
80			// Since different threads may write to the target,
81			// we need to protect the file position and I/O by a lock.
82			BLocker*			fStreamLock;
83};
84
85
86
87AVFormatWriter::StreamCookie::StreamCookie(AVFormatContext* context,
88		BLocker* streamLock)
89	:
90	fContext(context),
91	fStream(NULL),
92	fStreamLock(streamLock)
93{
94	av_init_packet(&fPacket);
95}
96
97
98AVFormatWriter::StreamCookie::~StreamCookie()
99{
100}
101
102
103status_t
104AVFormatWriter::StreamCookie::Init(media_format* format,
105	const media_codec_info* codecInfo)
106{
107	TRACE("AVFormatWriter::StreamCookie::Init()\n");
108
109	BAutolock _(fStreamLock);
110
111	fPacket.stream_index = fContext->nb_streams;
112	fStream = av_new_stream(fContext, fPacket.stream_index);
113
114	if (fStream == NULL) {
115		TRACE("  failed to add new stream\n");
116		return B_ERROR;
117	}
118
119//	TRACE("  fStream->codec: %p\n", fStream->codec);
120	// TODO: This is a hack for now! Use avcodec_find_encoder_by_name()
121	// or something similar...
122	fStream->codec->codec_id = (CodecID)codecInfo->sub_id;
123	if (fStream->codec->codec_id == CODEC_ID_NONE)
124		fStream->codec->codec_id = raw_audio_codec_id_for(*format);
125
126	// Setup the stream according to the media format...
127	if (format->type == B_MEDIA_RAW_VIDEO) {
128		fStream->codec->codec_type = AVMEDIA_TYPE_VIDEO;
129#if GET_CONTEXT_DEFAULTS
130// NOTE: API example does not do this:
131		avcodec_get_context_defaults(fStream->codec);
132#endif
133		// frame rate
134		fStream->codec->time_base.den = (int)format->u.raw_video.field_rate;
135		fStream->codec->time_base.num = 1;
136		// video size
137		fStream->codec->width = format->u.raw_video.display.line_width;
138		fStream->codec->height = format->u.raw_video.display.line_count;
139		// pixel aspect ratio
140		fStream->sample_aspect_ratio.num
141			= format->u.raw_video.pixel_width_aspect;
142		fStream->sample_aspect_ratio.den
143			= format->u.raw_video.pixel_height_aspect;
144		if (fStream->sample_aspect_ratio.num == 0
145			|| fStream->sample_aspect_ratio.den == 0) {
146			av_reduce(&fStream->sample_aspect_ratio.num,
147				&fStream->sample_aspect_ratio.den, fStream->codec->width,
148				fStream->codec->height, 255);
149		}
150
151		fStream->codec->gop_size = 12;
152
153		fStream->codec->sample_aspect_ratio = fStream->sample_aspect_ratio;
154
155		// Use the last supported pixel format of the AVCodec, which we hope
156		// is the one with the best quality (true for all currently supported
157		// encoders).
158//		AVCodec* codec = fStream->codec->codec;
159//		for (int i = 0; codec->pix_fmts[i] != PIX_FMT_NONE; i++)
160//			fStream->codec->pix_fmt = codec->pix_fmts[i];
161		fStream->codec->pix_fmt = PIX_FMT_YUV420P;
162
163	} else if (format->type == B_MEDIA_RAW_AUDIO) {
164		fStream->codec->codec_type = AVMEDIA_TYPE_AUDIO;
165#if GET_CONTEXT_DEFAULTS
166// NOTE: API example does not do this:
167		avcodec_get_context_defaults(fStream->codec);
168#endif
169		// frame rate
170		fStream->codec->sample_rate = (int)format->u.raw_audio.frame_rate;
171
172		// channels
173		fStream->codec->channels = format->u.raw_audio.channel_count;
174		switch (format->u.raw_audio.format) {
175			case media_raw_audio_format::B_AUDIO_FLOAT:
176				fStream->codec->sample_fmt = SAMPLE_FMT_FLT;
177				break;
178			case media_raw_audio_format::B_AUDIO_DOUBLE:
179				fStream->codec->sample_fmt = SAMPLE_FMT_DBL;
180				break;
181			case media_raw_audio_format::B_AUDIO_INT:
182				fStream->codec->sample_fmt = SAMPLE_FMT_S32;
183				break;
184			case media_raw_audio_format::B_AUDIO_SHORT:
185				fStream->codec->sample_fmt = SAMPLE_FMT_S16;
186				break;
187			case media_raw_audio_format::B_AUDIO_UCHAR:
188				fStream->codec->sample_fmt = SAMPLE_FMT_U8;
189				break;
190
191			case media_raw_audio_format::B_AUDIO_CHAR:
192			default:
193				return B_MEDIA_BAD_FORMAT;
194				break;
195		}
196		if (format->u.raw_audio.channel_mask == 0) {
197			// guess the channel mask...
198			switch (format->u.raw_audio.channel_count) {
199				default:
200				case 2:
201					fStream->codec->channel_layout = CH_LAYOUT_STEREO;
202					break;
203				case 1:
204					fStream->codec->channel_layout = CH_LAYOUT_MONO;
205					break;
206				case 3:
207					fStream->codec->channel_layout = CH_LAYOUT_SURROUND;
208					break;
209				case 4:
210					fStream->codec->channel_layout = CH_LAYOUT_QUAD;
211					break;
212				case 5:
213					fStream->codec->channel_layout = CH_LAYOUT_5POINT0;
214					break;
215				case 6:
216					fStream->codec->channel_layout = CH_LAYOUT_5POINT1;
217					break;
218				case 8:
219					fStream->codec->channel_layout = CH_LAYOUT_7POINT1;
220					break;
221				case 10:
222					fStream->codec->channel_layout = CH_LAYOUT_7POINT1_WIDE;
223					break;
224			}
225		} else {
226			// The bits match 1:1 for media_multi_channels and FFmpeg defines.
227			fStream->codec->channel_layout = format->u.raw_audio.channel_mask;
228		}
229	}
230
231	// Some formats want stream headers to be separate
232	if ((fContext->oformat->flags & AVFMT_GLOBALHEADER) != 0)
233		fStream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
234
235	TRACE("  stream->time_base: (%d/%d), codec->time_base: (%d/%d))\n",
236		fStream->time_base.num, fStream->time_base.den,
237		fStream->codec->time_base.num, fStream->codec->time_base.den);
238
239#if 0
240	// Write the AVCodecContext pointer to the user data section of the
241	// media_format. For some encoders, it seems to be necessary to use
242	// the AVCodecContext of the AVStream in order to successfully encode
243	// anything and write valid media files. For example some codecs need
244	// to store meta data or global data in the container.
245	app_info appInfo;
246	if (be_app->GetAppInfo(&appInfo) == B_OK) {
247		uchar* userData = format->user_data;
248		*(uint32*)userData = 'ffmp';
249		userData += sizeof(uint32);
250		*(team_id*)userData = appInfo.team;
251		userData += sizeof(team_id);
252		*(AVCodecContext**)userData = fStream->codec;
253	}
254#endif
255
256	return B_OK;
257}
258
259
260status_t
261AVFormatWriter::StreamCookie::WriteChunk(const void* chunkBuffer,
262	size_t chunkSize, media_encode_info* encodeInfo)
263{
264	TRACE_PACKET("AVFormatWriter::StreamCookie[%d]::WriteChunk(%p, %ld, "
265		"start_time: %lld)\n", fStream->index, chunkBuffer, chunkSize,
266		encodeInfo->start_time);
267
268	BAutolock _(fStreamLock);
269
270	fPacket.data = const_cast<uint8_t*>((const uint8_t*)chunkBuffer);
271	fPacket.size = chunkSize;
272
273	fPacket.pts = int64_t((double)encodeInfo->start_time
274		* fStream->time_base.den / (1000000.0 * fStream->time_base.num)
275		+ 0.5);
276
277	fPacket.flags = 0;
278	if ((encodeInfo->flags & B_MEDIA_KEY_FRAME) != 0)
279		fPacket.flags |= AV_PKT_FLAG_KEY;
280
281	TRACE_PACKET("  PTS: %lld (stream->time_base: (%d/%d), "
282		"codec->time_base: (%d/%d))\n", fPacket.pts,
283		fStream->time_base.num, fStream->time_base.den,
284		fStream->codec->time_base.num, fStream->codec->time_base.den);
285
286#if 0
287	// TODO: Eventually, we need to write interleaved packets, but
288	// maybe we are only supposed to use this if we have actually
289	// more than one stream. For the moment, this crashes in AVPacket
290	// shuffling inside libavformat. Maybe if we want to use this, we
291	// need to allocate a separate AVPacket and copy the chunk buffer.
292	int result = av_interleaved_write_frame(fContext, &fPacket);
293	if (result < 0)
294		TRACE("  av_interleaved_write_frame(): %d\n", result);
295#else
296	int result = av_write_frame(fContext, &fPacket);
297	if (result < 0)
298		TRACE("  av_write_frame(): %d\n", result);
299#endif
300
301	return result == 0 ? B_OK : B_ERROR;
302}
303
304
305status_t
306AVFormatWriter::StreamCookie::AddTrackInfo(uint32 code,
307	const void* data, size_t size, uint32 flags)
308{
309	TRACE("AVFormatWriter::StreamCookie::AddTrackInfo(%lu, %p, %ld, %lu)\n",
310		code, data, size, flags);
311
312	BAutolock _(fStreamLock);
313
314	return B_NOT_SUPPORTED;
315}
316
317
318// #pragma mark - AVFormatWriter
319
320
321AVFormatWriter::AVFormatWriter()
322	:
323	fContext(avformat_alloc_context()),
324	fHeaderWritten(false),
325	fIOContext(NULL),
326	fStreamLock("stream lock")
327{
328	TRACE("AVFormatWriter::AVFormatWriter\n");
329}
330
331
332AVFormatWriter::~AVFormatWriter()
333{
334	TRACE("AVFormatWriter::~AVFormatWriter\n");
335
336	// Free the streams and close the AVCodecContexts
337    for(unsigned i = 0; i < fContext->nb_streams; i++) {
338#if OPEN_CODEC_CONTEXT
339		// We only need to close the AVCodecContext when we opened it.
340		// This is experimental, see CommitHeader().
341		if (fHeaderWritten)
342			avcodec_close(fContext->streams[i]->codec);
343#endif
344		av_freep(&fContext->streams[i]->codec);
345		av_freep(&fContext->streams[i]);
346    }
347
348	av_free(fContext);
349	av_free(fIOContext->buffer);
350	av_free(fIOContext);
351}
352
353
354// #pragma mark -
355
356
357status_t
358AVFormatWriter::Init(const media_file_format* fileFormat)
359{
360	TRACE("AVFormatWriter::Init()\n");
361
362	uint8* buffer = static_cast<uint8*>(av_malloc(kIOBufferSize));
363	if (buffer == NULL)
364		return B_NO_MEMORY;
365
366	// Allocate I/O context and initialize it with buffer
367	// and hook functions, pass ourself as cookie.
368	fIOContext = av_alloc_put_byte(buffer, kIOBufferSize, 1, this,
369			0, _Write, _Seek);
370	if (fIOContext == NULL) {
371		TRACE("av_alloc_put_byte() failed!\n");
372		return B_ERROR;
373	}
374
375	// Setup I/O hooks. This seems to be enough.
376	fContext->pb = fIOContext;
377
378	// Set the AVOutputFormat according to fileFormat...
379	fContext->oformat = av_guess_format(fileFormat->short_name,
380		fileFormat->file_extension, fileFormat->mime_type);
381	if (fContext->oformat == NULL) {
382		TRACE("  failed to find AVOuputFormat for %s\n",
383			fileFormat->short_name);
384		return B_NOT_SUPPORTED;
385	}
386
387	TRACE("  found AVOuputFormat for %s: %s\n", fileFormat->short_name,
388		fContext->oformat->name);
389
390	return B_OK;
391}
392
393
394status_t
395AVFormatWriter::SetCopyright(const char* copyright)
396{
397	TRACE("AVFormatWriter::SetCopyright(%s)\n", copyright);
398
399	return B_NOT_SUPPORTED;
400}
401
402
403status_t
404AVFormatWriter::CommitHeader()
405{
406	TRACE("AVFormatWriter::CommitHeader\n");
407
408	if (fContext == NULL)
409		return B_NO_INIT;
410
411	if (fHeaderWritten)
412		return B_NOT_ALLOWED;
413
414	// According to output_example.c, the output parameters must be set even
415	// if none are specified. In the example, this call is used after the
416	// streams have been created.
417	if (av_set_parameters(fContext, NULL) < 0)
418		return B_ERROR;
419
420#if OPEN_CODEC_CONTEXT
421	for (unsigned i = 0; i < fContext->nb_streams; i++) {
422		AVStream* stream = fContext->streams[i];
423		// NOTE: Experimental, this should not be needed. Especially, since
424		// we have no idea (in the future) what CodecID some encoder uses,
425		// it may be an encoder from a different plugin.
426		AVCodecContext* codecContext = stream->codec;
427		AVCodec* codec = avcodec_find_encoder(codecContext->codec_id);
428		if (codec == NULL || avcodec_open(codecContext, codec) < 0) {
429			TRACE("  stream[%u] - failed to open AVCodecContext\n", i);
430		}
431		TRACE("  stream[%u] time_base: (%d/%d), codec->time_base: (%d/%d)\n",
432			i, stream->time_base.num, stream->time_base.den,
433			stream->codec->time_base.num, stream->codec->time_base.den);
434	}
435#endif
436
437	int result = av_write_header(fContext);
438	if (result < 0)
439		TRACE("  av_write_header(): %d\n", result);
440
441	// We need to close the codecs we opened, even in case of failure.
442	fHeaderWritten = true;
443
444	#ifdef TRACE_AVFORMAT_WRITER
445	TRACE("  wrote header\n");
446	for (unsigned i = 0; i < fContext->nb_streams; i++) {
447		AVStream* stream = fContext->streams[i];
448		TRACE("  stream[%u] time_base: (%d/%d), codec->time_base: (%d/%d)\n",
449			i, stream->time_base.num, stream->time_base.den,
450			stream->codec->time_base.num, stream->codec->time_base.den);
451	}
452	#endif // TRACE_AVFORMAT_WRITER
453
454	return result == 0 ? B_OK : B_ERROR;
455}
456
457
458status_t
459AVFormatWriter::Flush()
460{
461	TRACE("AVFormatWriter::Flush\n");
462
463	return B_NOT_SUPPORTED;
464}
465
466
467status_t
468AVFormatWriter::Close()
469{
470	TRACE("AVFormatWriter::Close\n");
471
472	if (fContext == NULL)
473		return B_NO_INIT;
474
475	if (!fHeaderWritten)
476		return B_NOT_ALLOWED;
477
478	int result = av_write_trailer(fContext);
479	if (result < 0)
480		TRACE("  av_write_trailer(): %d\n", result);
481
482	return result == 0 ? B_OK : B_ERROR;
483}
484
485
486status_t
487AVFormatWriter::AllocateCookie(void** _cookie, media_format* format,
488	const media_codec_info* codecInfo)
489{
490	TRACE("AVFormatWriter::AllocateCookie()\n");
491
492	if (fHeaderWritten)
493		return B_NOT_ALLOWED;
494
495	BAutolock _(fStreamLock);
496
497	if (_cookie == NULL)
498		return B_BAD_VALUE;
499
500	StreamCookie* cookie = new(std::nothrow) StreamCookie(fContext,
501		&fStreamLock);
502
503	status_t ret = cookie->Init(format, codecInfo);
504	if (ret != B_OK) {
505		delete cookie;
506		return ret;
507	}
508
509	*_cookie = cookie;
510	return B_OK;
511}
512
513
514status_t
515AVFormatWriter::FreeCookie(void* _cookie)
516{
517	BAutolock _(fStreamLock);
518
519	StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
520	delete cookie;
521
522	return B_OK;
523}
524
525
526// #pragma mark -
527
528
529status_t
530AVFormatWriter::SetCopyright(void* cookie, const char* copyright)
531{
532	TRACE("AVFormatWriter::SetCopyright(%p, %s)\n", cookie, copyright);
533
534	return B_NOT_SUPPORTED;
535}
536
537
538status_t
539AVFormatWriter::AddTrackInfo(void* _cookie, uint32 code,
540	const void* data, size_t size, uint32 flags)
541{
542	TRACE("AVFormatWriter::AddTrackInfo(%lu, %p, %ld, %lu)\n",
543		code, data, size, flags);
544
545	StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
546	return cookie->AddTrackInfo(code, data, size, flags);
547}
548
549
550status_t
551AVFormatWriter::WriteChunk(void* _cookie, const void* chunkBuffer,
552	size_t chunkSize, media_encode_info* encodeInfo)
553{
554	TRACE_PACKET("AVFormatWriter::WriteChunk(%p, %ld, %p)\n", chunkBuffer,
555		chunkSize, encodeInfo);
556
557	StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
558	return cookie->WriteChunk(chunkBuffer, chunkSize, encodeInfo);
559}
560
561
562// #pragma mark - I/O hooks
563
564
565/*static*/ int
566AVFormatWriter::_Write(void* cookie, uint8* buffer, int bufferSize)
567{
568	TRACE_IO("AVFormatWriter::_Write(%p, %p, %d)\n",
569		cookie, buffer, bufferSize);
570
571	AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie);
572
573	ssize_t written = writer->fTarget->Write(buffer, bufferSize);
574
575	TRACE_IO("  written: %ld\n", written);
576	return (int)written;
577
578}
579
580
581/*static*/ off_t
582AVFormatWriter::_Seek(void* cookie, off_t offset, int whence)
583{
584	TRACE_IO("AVFormatWriter::_Seek(%p, %lld, %d)\n",
585		cookie, offset, whence);
586
587	AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie);
588
589	BPositionIO* positionIO = dynamic_cast<BPositionIO*>(writer->fTarget);
590	if (positionIO == NULL)
591		return -1;
592
593	// Support for special file size retrieval API without seeking anywhere:
594	if (whence == AVSEEK_SIZE) {
595		off_t size;
596		if (positionIO->GetSize(&size) == B_OK)
597			return size;
598		return -1;
599	}
600
601	off_t position = positionIO->Seek(offset, whence);
602	TRACE_IO("  position: %lld\n", position);
603	if (position < 0)
604		return -1;
605
606	return position;
607}
608
609
610