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