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