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