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