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