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 <DataIO.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 fHeaderWritten(false), 384 fIOContext(NULL), 385 fStreamLock("stream lock") 386 { 387 TRACE("AVFormatWriter::AVFormatWriter\n"); 388 } 389 390 391 AVFormatWriter::~AVFormatWriter() 392 { 393 TRACE("AVFormatWriter::~AVFormatWriter\n"); 394 395 // Free the streams and close the AVCodecContexts 396 for(unsigned i = 0; i < fContext->nb_streams; i++) { 397 #if OPEN_CODEC_CONTEXT 398 // We only need to close the AVCodecContext when we opened it. 399 // This is experimental, see CommitHeader(). 400 if (fHeaderWritten) 401 avcodec_close(fContext->streams[i]->codec); 402 #endif 403 av_freep(&fContext->streams[i]->codec); 404 av_freep(&fContext->streams[i]); 405 } 406 407 av_free(fContext); 408 av_free(fIOContext->buffer); 409 av_free(fIOContext); 410 } 411 412 413 // #pragma mark - 414 415 416 status_t 417 AVFormatWriter::Init(const media_file_format* fileFormat) 418 { 419 TRACE("AVFormatWriter::Init()\n"); 420 421 uint8* buffer = static_cast<uint8*>(av_malloc(kIOBufferSize)); 422 if (buffer == NULL) 423 return B_NO_MEMORY; 424 425 // Allocate I/O context and initialize it with buffer 426 // and hook functions, pass ourself as cookie. 427 fIOContext = avio_alloc_context(buffer, kIOBufferSize, 1, this, 428 0, _Write, _Seek); 429 if (fIOContext == NULL) { 430 TRACE("av_alloc_put_byte() failed!\n"); 431 return B_ERROR; 432 } 433 434 // Setup I/O hooks. This seems to be enough. 435 fContext->pb = fIOContext; 436 437 // Set the AVOutputFormat according to fileFormat... 438 fContext->oformat = av_guess_format(fileFormat->short_name, 439 fileFormat->file_extension, fileFormat->mime_type); 440 if (fContext->oformat == NULL) { 441 TRACE(" failed to find AVOuputFormat for %s\n", 442 fileFormat->short_name); 443 return B_NOT_SUPPORTED; 444 } 445 446 TRACE(" found AVOuputFormat for %s: %s\n", fileFormat->short_name, 447 fContext->oformat->name); 448 449 return B_OK; 450 } 451 452 453 status_t 454 AVFormatWriter::SetCopyright(const char* copyright) 455 { 456 TRACE("AVFormatWriter::SetCopyright(%s)\n", copyright); 457 458 return B_NOT_SUPPORTED; 459 } 460 461 462 status_t 463 AVFormatWriter::CommitHeader() 464 { 465 TRACE("AVFormatWriter::CommitHeader\n"); 466 467 if (fContext == NULL) 468 return B_NO_INIT; 469 470 if (fHeaderWritten) 471 return B_NOT_ALLOWED; 472 473 #if OPEN_CODEC_CONTEXT 474 for (unsigned i = 0; i < fContext->nb_streams; i++) { 475 AVStream* stream = fContext->streams[i]; 476 // NOTE: Experimental, this should not be needed. Especially, since 477 // we have no idea (in the future) what CodecID some encoder uses, 478 // it may be an encoder from a different plugin. 479 AVCodecContext* codecContext = stream->codec; 480 codecContext->strict_std_compliance = -2; 481 AVCodec* codec = avcodec_find_encoder(codecContext->codec_id); 482 if (codec == NULL || avcodec_open2(codecContext, codec, NULL) < 0) { 483 TRACE(" stream[%u] - failed to open AVCodecContext\n", i); 484 } 485 TRACE(" stream[%u] time_base: (%d/%d), codec->time_base: (%d/%d)\n", 486 i, stream->time_base.num, stream->time_base.den, 487 stream->codec->time_base.num, stream->codec->time_base.den); 488 } 489 #endif 490 491 int result = avformat_write_header(fContext, NULL); 492 if (result < 0) 493 TRACE(" avformat_write_header(): %d\n", result); 494 495 // We need to close the codecs we opened, even in case of failure. 496 fHeaderWritten = true; 497 498 #ifdef TRACE_AVFORMAT_WRITER 499 TRACE(" wrote header\n"); 500 for (unsigned i = 0; i < fContext->nb_streams; i++) { 501 AVStream* stream = fContext->streams[i]; 502 TRACE(" stream[%u] time_base: (%d/%d), codec->time_base: (%d/%d)\n", 503 i, stream->time_base.num, stream->time_base.den, 504 stream->codec->time_base.num, stream->codec->time_base.den); 505 } 506 #endif // TRACE_AVFORMAT_WRITER 507 508 return result == 0 ? B_OK : B_ERROR; 509 } 510 511 512 status_t 513 AVFormatWriter::Flush() 514 { 515 TRACE("AVFormatWriter::Flush\n"); 516 517 return B_NOT_SUPPORTED; 518 } 519 520 521 status_t 522 AVFormatWriter::Close() 523 { 524 TRACE("AVFormatWriter::Close\n"); 525 526 if (fContext == NULL) 527 return B_NO_INIT; 528 529 if (!fHeaderWritten) 530 return B_NOT_ALLOWED; 531 532 int result = av_write_trailer(fContext); 533 if (result < 0) 534 TRACE(" av_write_trailer(): %d\n", result); 535 536 return result == 0 ? B_OK : B_ERROR; 537 } 538 539 540 status_t 541 AVFormatWriter::AllocateCookie(void** _cookie, media_format* format, 542 const media_codec_info* codecInfo) 543 { 544 TRACE("AVFormatWriter::AllocateCookie()\n"); 545 546 if (fHeaderWritten) 547 return B_NOT_ALLOWED; 548 549 BAutolock _(fStreamLock); 550 551 if (_cookie == NULL) 552 return B_BAD_VALUE; 553 554 StreamCookie* cookie = new(std::nothrow) StreamCookie(fContext, 555 &fStreamLock); 556 557 status_t ret = cookie->Init(format, codecInfo); 558 if (ret != B_OK) { 559 delete cookie; 560 return ret; 561 } 562 563 *_cookie = cookie; 564 return B_OK; 565 } 566 567 568 status_t 569 AVFormatWriter::FreeCookie(void* _cookie) 570 { 571 BAutolock _(fStreamLock); 572 573 StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie); 574 delete cookie; 575 576 return B_OK; 577 } 578 579 580 // #pragma mark - 581 582 583 status_t 584 AVFormatWriter::SetCopyright(void* cookie, const char* copyright) 585 { 586 TRACE("AVFormatWriter::SetCopyright(%p, %s)\n", cookie, copyright); 587 588 return B_NOT_SUPPORTED; 589 } 590 591 592 status_t 593 AVFormatWriter::AddTrackInfo(void* _cookie, uint32 code, 594 const void* data, size_t size, uint32 flags) 595 { 596 TRACE("AVFormatWriter::AddTrackInfo(%lu, %p, %ld, %lu)\n", 597 code, data, size, flags); 598 599 StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie); 600 return cookie->AddTrackInfo(code, data, size, flags); 601 } 602 603 604 status_t 605 AVFormatWriter::WriteChunk(void* _cookie, const void* chunkBuffer, 606 size_t chunkSize, media_encode_info* encodeInfo) 607 { 608 TRACE_PACKET("AVFormatWriter::WriteChunk(%p, %ld, %p)\n", chunkBuffer, 609 chunkSize, encodeInfo); 610 611 StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie); 612 return cookie->WriteChunk(chunkBuffer, chunkSize, encodeInfo); 613 } 614 615 616 // #pragma mark - I/O hooks 617 618 619 /*static*/ int 620 AVFormatWriter::_Write(void* cookie, uint8* buffer, int bufferSize) 621 { 622 TRACE_IO("AVFormatWriter::_Write(%p, %p, %d)\n", 623 cookie, buffer, bufferSize); 624 625 AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie); 626 627 ssize_t written = writer->fTarget->Write(buffer, bufferSize); 628 629 TRACE_IO(" written: %ld\n", written); 630 return (int)written; 631 632 } 633 634 635 /*static*/ off_t 636 AVFormatWriter::_Seek(void* cookie, off_t offset, int whence) 637 { 638 TRACE_IO("AVFormatWriter::_Seek(%p, %lld, %d)\n", 639 cookie, offset, whence); 640 641 AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie); 642 643 BPositionIO* positionIO = dynamic_cast<BPositionIO*>(writer->fTarget); 644 if (positionIO == NULL) 645 return -1; 646 647 // Support for special file size retrieval API without seeking anywhere: 648 if (whence == AVSEEK_SIZE) { 649 off_t size; 650 if (positionIO->GetSize(&size) == B_OK) 651 return size; 652 return -1; 653 } 654 655 off_t position = positionIO->Seek(offset, whence); 656 TRACE_IO(" position: %lld\n", position); 657 if (position < 0) 658 return -1; 659 660 return position; 661 } 662 663 664