xref: /haiku/src/add-ons/media/plugins/ffmpeg/AVFormatWriter.cpp (revision b6b0567fbd186f8ce8a0c90bdc7a7b5b4c649678)
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 WriteHeader().
331 		avcodec_close(fContext->streams[i]->codec);
332 #endif
333 		av_freep(&fContext->streams[i]->codec);
334 		av_freep(&fContext->streams[i]);
335     }
336 
337 	av_free(fContext);
338 
339 	delete[] fIOBuffer;
340 }
341 
342 
343 // #pragma mark -
344 
345 
346 status_t
347 AVFormatWriter::Init(const media_file_format* fileFormat)
348 {
349 	TRACE("AVFormatWriter::Init()\n");
350 
351 	delete[] fIOBuffer;
352 	fIOBuffer = new(std::nothrow) uint8[kIOBufferSize];
353 	if (fIOBuffer == NULL)
354 		return B_NO_MEMORY;
355 
356 	// Init I/O context with buffer and hook functions, pass ourself as
357 	// cookie.
358 	if (init_put_byte(&fIOContext, fIOBuffer, kIOBufferSize, 0, this,
359 			0, _Write, _Seek) != 0) {
360 		TRACE("  init_put_byte() failed!\n");
361 		return B_ERROR;
362 	}
363 
364 	// Setup I/O hooks. This seems to be enough.
365 	fContext->pb = &fIOContext;
366 
367 	// Set the AVOutputFormat according to fileFormat...
368 	fContext->oformat = guess_format(fileFormat->short_name,
369 		fileFormat->file_extension, fileFormat->mime_type);
370 	if (fContext->oformat == NULL) {
371 		TRACE("  failed to find AVOuputFormat for %s\n",
372 			fileFormat->short_name);
373 		return B_NOT_SUPPORTED;
374 	}
375 
376 	TRACE("  found AVOuputFormat for %s: %s\n", fileFormat->short_name,
377 		fContext->oformat->name);
378 
379 	return B_OK;
380 }
381 
382 
383 status_t
384 AVFormatWriter::SetCopyright(const char* copyright)
385 {
386 	TRACE("AVFormatWriter::SetCopyright(%s)\n", copyright);
387 
388 	return B_NOT_SUPPORTED;
389 }
390 
391 
392 status_t
393 AVFormatWriter::CommitHeader()
394 {
395 	TRACE("AVFormatWriter::CommitHeader\n");
396 
397 	if (fContext == NULL)
398 		return B_NO_INIT;
399 
400 	if (fHeaderWritten)
401 		return B_NOT_ALLOWED;
402 
403 	// According to output_example.c, the output parameters must be set even
404 	// if none are specified. In the example, this call is used after the
405 	// streams have been created.
406 	if (av_set_parameters(fContext, NULL) < 0)
407 		return B_ERROR;
408 
409 	for (unsigned i = 0; i < fContext->nb_streams; i++) {
410 		AVStream* stream = fContext->streams[i];
411 #if OPEN_CODEC_CONTEXT
412 		// NOTE: Experimental, this should not be needed. Especially, since
413 		// we have no idea (in the future) what CodecID some encoder uses,
414 		// it may be an encoder from a different plugin.
415 		AVCodecContext* codecContext = stream->codec;
416 		AVCodec* codec = avcodec_find_encoder(codecContext->codec_id);
417 		if (codec == NULL || avcodec_open(codecContext, codec) < 0) {
418 			TRACE("  stream[%u] - failed to open AVCodecContext\n", i);
419 		}
420 #endif
421 		TRACE("  stream[%u] time_base: (%d/%d), codec->time_base: (%d/%d)\n",
422 			i, stream->time_base.num, stream->time_base.den,
423 			stream->codec->time_base.num, stream->codec->time_base.den);
424 	}
425 
426 	int result = av_write_header(fContext);
427 	if (result < 0)
428 		TRACE("  av_write_header(): %d\n", result);
429 	else
430 		fHeaderWritten = true;
431 
432 	TRACE("  wrote header\n");
433 	for (unsigned i = 0; i < fContext->nb_streams; i++) {
434 		AVStream* stream = fContext->streams[i];
435 		TRACE("  stream[%u] time_base: (%d/%d), codec->time_base: (%d/%d)\n",
436 			i, stream->time_base.num, stream->time_base.den,
437 			stream->codec->time_base.num, stream->codec->time_base.den);
438 	}
439 
440 	return result == 0 ? B_OK : B_ERROR;
441 }
442 
443 
444 status_t
445 AVFormatWriter::Flush()
446 {
447 	TRACE("AVFormatWriter::Flush\n");
448 
449 	return B_NOT_SUPPORTED;
450 }
451 
452 
453 status_t
454 AVFormatWriter::Close()
455 {
456 	TRACE("AVFormatWriter::Close\n");
457 
458 	if (fContext == NULL)
459 		return B_NO_INIT;
460 
461 	if (!fHeaderWritten)
462 		return B_NOT_ALLOWED;
463 
464 	int result = av_write_trailer(fContext);
465 	if (result < 0)
466 		TRACE("  av_write_trailer(): %d\n", result);
467 
468 	return result == 0 ? B_OK : B_ERROR;
469 }
470 
471 
472 status_t
473 AVFormatWriter::AllocateCookie(void** _cookie, const media_format* format,
474 	const media_codec_info* codecInfo)
475 {
476 	TRACE("AVFormatWriter::AllocateCookie()\n");
477 
478 	if (fHeaderWritten)
479 		return B_NOT_ALLOWED;
480 
481 	BAutolock _(fStreamLock);
482 
483 	if (_cookie == NULL)
484 		return B_BAD_VALUE;
485 
486 	StreamCookie* cookie = new(std::nothrow) StreamCookie(fContext,
487 		&fStreamLock);
488 
489 	status_t ret = cookie->Init(format, codecInfo);
490 	if (ret != B_OK) {
491 		delete cookie;
492 		return ret;
493 	}
494 
495 	*_cookie = cookie;
496 	return B_OK;
497 }
498 
499 
500 status_t
501 AVFormatWriter::FreeCookie(void* _cookie)
502 {
503 	BAutolock _(fStreamLock);
504 
505 	StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
506 	delete cookie;
507 
508 	return B_OK;
509 }
510 
511 
512 // #pragma mark -
513 
514 
515 status_t
516 AVFormatWriter::SetCopyright(void* cookie, const char* copyright)
517 {
518 	TRACE("AVFormatWriter::SetCopyright(%p, %s)\n", cookie, copyright);
519 
520 	return B_NOT_SUPPORTED;
521 }
522 
523 
524 status_t
525 AVFormatWriter::AddTrackInfo(void* _cookie, uint32 code,
526 	const void* data, size_t size, uint32 flags)
527 {
528 	TRACE("AVFormatWriter::AddTrackInfo(%lu, %p, %ld, %lu)\n",
529 		code, data, size, flags);
530 
531 	StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
532 	return cookie->AddTrackInfo(code, data, size, flags);
533 }
534 
535 
536 status_t
537 AVFormatWriter::WriteChunk(void* _cookie, const void* chunkBuffer,
538 	size_t chunkSize, media_encode_info* encodeInfo)
539 {
540 	TRACE_PACKET("AVFormatWriter::WriteChunk(%p, %ld, %p)\n", chunkBuffer,
541 		chunkSize, encodeInfo);
542 
543 	StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie);
544 	return cookie->WriteChunk(chunkBuffer, chunkSize, encodeInfo);
545 }
546 
547 
548 // #pragma mark - I/O hooks
549 
550 
551 /*static*/ int
552 AVFormatWriter::_Write(void* cookie, uint8* buffer, int bufferSize)
553 {
554 	TRACE_IO("AVFormatWriter::_Write(%p, %p, %d)\n",
555 		cookie, buffer, bufferSize);
556 
557 	AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie);
558 
559 	ssize_t written = writer->fTarget->Write(buffer, bufferSize);
560 
561 	TRACE_IO("  written: %ld\n", written);
562 	return (int)written;
563 
564 }
565 
566 
567 /*static*/ off_t
568 AVFormatWriter::_Seek(void* cookie, off_t offset, int whence)
569 {
570 	TRACE_IO("AVFormatWriter::_Seek(%p, %lld, %d)\n",
571 		cookie, offset, whence);
572 
573 	AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie);
574 
575 	BPositionIO* positionIO = dynamic_cast<BPositionIO*>(writer->fTarget);
576 	if (positionIO == NULL)
577 		return -1;
578 
579 	// Support for special file size retrieval API without seeking anywhere:
580 	if (whence == AVSEEK_SIZE) {
581 		off_t size;
582 		if (positionIO->GetSize(&size) == B_OK)
583 			return size;
584 		return -1;
585 	}
586 
587 	off_t position = positionIO->Seek(offset, whence);
588 	TRACE_IO("  position: %lld\n", position);
589 	if (position < 0)
590 		return -1;
591 
592 	return position;
593 }
594 
595 
596