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