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