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