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 switch (format->u.raw_audio.format) { 175 case media_raw_audio_format::B_AUDIO_FLOAT: 176 fStream->codec->sample_fmt = SAMPLE_FMT_FLT; 177 break; 178 case media_raw_audio_format::B_AUDIO_DOUBLE: 179 fStream->codec->sample_fmt = SAMPLE_FMT_DBL; 180 break; 181 case media_raw_audio_format::B_AUDIO_INT: 182 fStream->codec->sample_fmt = SAMPLE_FMT_S32; 183 break; 184 case media_raw_audio_format::B_AUDIO_SHORT: 185 fStream->codec->sample_fmt = SAMPLE_FMT_S16; 186 break; 187 case media_raw_audio_format::B_AUDIO_UCHAR: 188 fStream->codec->sample_fmt = SAMPLE_FMT_U8; 189 break; 190 191 case media_raw_audio_format::B_AUDIO_CHAR: 192 default: 193 return B_MEDIA_BAD_FORMAT; 194 break; 195 } 196 if (format->u.raw_audio.channel_mask == 0) { 197 // guess the channel mask... 198 switch (format->u.raw_audio.channel_count) { 199 default: 200 case 2: 201 fStream->codec->channel_layout = CH_LAYOUT_STEREO; 202 break; 203 case 1: 204 fStream->codec->channel_layout = CH_LAYOUT_MONO; 205 break; 206 case 3: 207 fStream->codec->channel_layout = CH_LAYOUT_SURROUND; 208 break; 209 case 4: 210 fStream->codec->channel_layout = CH_LAYOUT_QUAD; 211 break; 212 case 5: 213 fStream->codec->channel_layout = CH_LAYOUT_5POINT0; 214 break; 215 case 6: 216 fStream->codec->channel_layout = CH_LAYOUT_5POINT1; 217 break; 218 case 8: 219 fStream->codec->channel_layout = CH_LAYOUT_7POINT1; 220 break; 221 case 10: 222 fStream->codec->channel_layout = CH_LAYOUT_7POINT1_WIDE; 223 break; 224 } 225 } else { 226 // The bits match 1:1 for media_multi_channels and FFmpeg defines. 227 fStream->codec->channel_layout = format->u.raw_audio.channel_mask; 228 } 229 } 230 231 // Some formats want stream headers to be separate 232 if ((fContext->oformat->flags & AVFMT_GLOBALHEADER) != 0) 233 fStream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER; 234 235 TRACE(" stream->time_base: (%d/%d), codec->time_base: (%d/%d))\n", 236 fStream->time_base.num, fStream->time_base.den, 237 fStream->codec->time_base.num, fStream->codec->time_base.den); 238 239 #if 0 240 // Write the AVCodecContext pointer to the user data section of the 241 // media_format. For some encoders, it seems to be necessary to use 242 // the AVCodecContext of the AVStream in order to successfully encode 243 // anything and write valid media files. For example some codecs need 244 // to store meta data or global data in the container. 245 app_info appInfo; 246 if (be_app->GetAppInfo(&appInfo) == B_OK) { 247 uchar* userData = format->user_data; 248 *(uint32*)userData = 'ffmp'; 249 userData += sizeof(uint32); 250 *(team_id*)userData = appInfo.team; 251 userData += sizeof(team_id); 252 *(AVCodecContext**)userData = fStream->codec; 253 } 254 #endif 255 256 return B_OK; 257 } 258 259 260 status_t 261 AVFormatWriter::StreamCookie::WriteChunk(const void* chunkBuffer, 262 size_t chunkSize, media_encode_info* encodeInfo) 263 { 264 TRACE_PACKET("AVFormatWriter::StreamCookie[%d]::WriteChunk(%p, %ld, " 265 "start_time: %lld)\n", fStream->index, chunkBuffer, chunkSize, 266 encodeInfo->start_time); 267 268 BAutolock _(fStreamLock); 269 270 fPacket.data = const_cast<uint8_t*>((const uint8_t*)chunkBuffer); 271 fPacket.size = chunkSize; 272 273 fPacket.pts = int64_t((double)encodeInfo->start_time 274 * fStream->time_base.den / (1000000.0 * fStream->time_base.num) 275 + 0.5); 276 277 fPacket.flags = 0; 278 if ((encodeInfo->flags & B_MEDIA_KEY_FRAME) != 0) 279 fPacket.flags |= AV_PKT_FLAG_KEY; 280 281 TRACE_PACKET(" PTS: %lld (stream->time_base: (%d/%d), " 282 "codec->time_base: (%d/%d))\n", fPacket.pts, 283 fStream->time_base.num, fStream->time_base.den, 284 fStream->codec->time_base.num, fStream->codec->time_base.den); 285 286 #if 0 287 // TODO: Eventually, we need to write interleaved packets, but 288 // maybe we are only supposed to use this if we have actually 289 // more than one stream. For the moment, this crashes in AVPacket 290 // shuffling inside libavformat. Maybe if we want to use this, we 291 // need to allocate a separate AVPacket and copy the chunk buffer. 292 int result = av_interleaved_write_frame(fContext, &fPacket); 293 if (result < 0) 294 TRACE(" av_interleaved_write_frame(): %d\n", result); 295 #else 296 int result = av_write_frame(fContext, &fPacket); 297 if (result < 0) 298 TRACE(" av_write_frame(): %d\n", result); 299 #endif 300 301 return result == 0 ? B_OK : B_ERROR; 302 } 303 304 305 status_t 306 AVFormatWriter::StreamCookie::AddTrackInfo(uint32 code, 307 const void* data, size_t size, uint32 flags) 308 { 309 TRACE("AVFormatWriter::StreamCookie::AddTrackInfo(%lu, %p, %ld, %lu)\n", 310 code, data, size, flags); 311 312 BAutolock _(fStreamLock); 313 314 return B_NOT_SUPPORTED; 315 } 316 317 318 // #pragma mark - AVFormatWriter 319 320 321 AVFormatWriter::AVFormatWriter() 322 : 323 fContext(avformat_alloc_context()), 324 fHeaderWritten(false), 325 fIOContext(NULL), 326 fStreamLock("stream lock") 327 { 328 TRACE("AVFormatWriter::AVFormatWriter\n"); 329 } 330 331 332 AVFormatWriter::~AVFormatWriter() 333 { 334 TRACE("AVFormatWriter::~AVFormatWriter\n"); 335 336 // Free the streams and close the AVCodecContexts 337 for(unsigned i = 0; i < fContext->nb_streams; i++) { 338 #if OPEN_CODEC_CONTEXT 339 // We only need to close the AVCodecContext when we opened it. 340 // This is experimental, see CommitHeader(). 341 if (fHeaderWritten) 342 avcodec_close(fContext->streams[i]->codec); 343 #endif 344 av_freep(&fContext->streams[i]->codec); 345 av_freep(&fContext->streams[i]); 346 } 347 348 av_free(fContext); 349 av_free(fIOContext->buffer); 350 av_free(fIOContext); 351 } 352 353 354 // #pragma mark - 355 356 357 status_t 358 AVFormatWriter::Init(const media_file_format* fileFormat) 359 { 360 TRACE("AVFormatWriter::Init()\n"); 361 362 uint8* buffer = static_cast<uint8*>(av_malloc(kIOBufferSize)); 363 if (buffer == NULL) 364 return B_NO_MEMORY; 365 366 // Init I/O context with buffer and hook functions, pass ourself as 367 // cookie. 368 if (init_put_byte(fIOContext, buffer, kIOBufferSize, 1, this, 369 0, _Write, _Seek) != 0) { 370 TRACE(" init_put_byte() failed!\n"); 371 return B_ERROR; 372 } 373 374 // Setup I/O hooks. This seems to be enough. 375 fContext->pb = fIOContext; 376 377 // Set the AVOutputFormat according to fileFormat... 378 fContext->oformat = av_guess_format(fileFormat->short_name, 379 fileFormat->file_extension, fileFormat->mime_type); 380 if (fContext->oformat == NULL) { 381 TRACE(" failed to find AVOuputFormat for %s\n", 382 fileFormat->short_name); 383 return B_NOT_SUPPORTED; 384 } 385 386 TRACE(" found AVOuputFormat for %s: %s\n", fileFormat->short_name, 387 fContext->oformat->name); 388 389 return B_OK; 390 } 391 392 393 status_t 394 AVFormatWriter::SetCopyright(const char* copyright) 395 { 396 TRACE("AVFormatWriter::SetCopyright(%s)\n", copyright); 397 398 return B_NOT_SUPPORTED; 399 } 400 401 402 status_t 403 AVFormatWriter::CommitHeader() 404 { 405 TRACE("AVFormatWriter::CommitHeader\n"); 406 407 if (fContext == NULL) 408 return B_NO_INIT; 409 410 if (fHeaderWritten) 411 return B_NOT_ALLOWED; 412 413 // According to output_example.c, the output parameters must be set even 414 // if none are specified. In the example, this call is used after the 415 // streams have been created. 416 if (av_set_parameters(fContext, NULL) < 0) 417 return B_ERROR; 418 419 #if OPEN_CODEC_CONTEXT 420 for (unsigned i = 0; i < fContext->nb_streams; i++) { 421 AVStream* stream = fContext->streams[i]; 422 // NOTE: Experimental, this should not be needed. Especially, since 423 // we have no idea (in the future) what CodecID some encoder uses, 424 // it may be an encoder from a different plugin. 425 AVCodecContext* codecContext = stream->codec; 426 AVCodec* codec = avcodec_find_encoder(codecContext->codec_id); 427 if (codec == NULL || avcodec_open(codecContext, codec) < 0) { 428 TRACE(" stream[%u] - failed to open AVCodecContext\n", i); 429 } 430 TRACE(" stream[%u] time_base: (%d/%d), codec->time_base: (%d/%d)\n", 431 i, stream->time_base.num, stream->time_base.den, 432 stream->codec->time_base.num, stream->codec->time_base.den); 433 } 434 #endif 435 436 int result = av_write_header(fContext); 437 if (result < 0) 438 TRACE(" av_write_header(): %d\n", result); 439 440 // We need to close the codecs we opened, even in case of failure. 441 fHeaderWritten = true; 442 443 #ifdef TRACE_AVFORMAT_WRITER 444 TRACE(" wrote header\n"); 445 for (unsigned i = 0; i < fContext->nb_streams; i++) { 446 AVStream* stream = fContext->streams[i]; 447 TRACE(" stream[%u] time_base: (%d/%d), codec->time_base: (%d/%d)\n", 448 i, stream->time_base.num, stream->time_base.den, 449 stream->codec->time_base.num, stream->codec->time_base.den); 450 } 451 #endif // TRACE_AVFORMAT_WRITER 452 453 return result == 0 ? B_OK : B_ERROR; 454 } 455 456 457 status_t 458 AVFormatWriter::Flush() 459 { 460 TRACE("AVFormatWriter::Flush\n"); 461 462 return B_NOT_SUPPORTED; 463 } 464 465 466 status_t 467 AVFormatWriter::Close() 468 { 469 TRACE("AVFormatWriter::Close\n"); 470 471 if (fContext == NULL) 472 return B_NO_INIT; 473 474 if (!fHeaderWritten) 475 return B_NOT_ALLOWED; 476 477 int result = av_write_trailer(fContext); 478 if (result < 0) 479 TRACE(" av_write_trailer(): %d\n", result); 480 481 return result == 0 ? B_OK : B_ERROR; 482 } 483 484 485 status_t 486 AVFormatWriter::AllocateCookie(void** _cookie, media_format* format, 487 const media_codec_info* codecInfo) 488 { 489 TRACE("AVFormatWriter::AllocateCookie()\n"); 490 491 if (fHeaderWritten) 492 return B_NOT_ALLOWED; 493 494 BAutolock _(fStreamLock); 495 496 if (_cookie == NULL) 497 return B_BAD_VALUE; 498 499 StreamCookie* cookie = new(std::nothrow) StreamCookie(fContext, 500 &fStreamLock); 501 502 status_t ret = cookie->Init(format, codecInfo); 503 if (ret != B_OK) { 504 delete cookie; 505 return ret; 506 } 507 508 *_cookie = cookie; 509 return B_OK; 510 } 511 512 513 status_t 514 AVFormatWriter::FreeCookie(void* _cookie) 515 { 516 BAutolock _(fStreamLock); 517 518 StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie); 519 delete cookie; 520 521 return B_OK; 522 } 523 524 525 // #pragma mark - 526 527 528 status_t 529 AVFormatWriter::SetCopyright(void* cookie, const char* copyright) 530 { 531 TRACE("AVFormatWriter::SetCopyright(%p, %s)\n", cookie, copyright); 532 533 return B_NOT_SUPPORTED; 534 } 535 536 537 status_t 538 AVFormatWriter::AddTrackInfo(void* _cookie, uint32 code, 539 const void* data, size_t size, uint32 flags) 540 { 541 TRACE("AVFormatWriter::AddTrackInfo(%lu, %p, %ld, %lu)\n", 542 code, data, size, flags); 543 544 StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie); 545 return cookie->AddTrackInfo(code, data, size, flags); 546 } 547 548 549 status_t 550 AVFormatWriter::WriteChunk(void* _cookie, const void* chunkBuffer, 551 size_t chunkSize, media_encode_info* encodeInfo) 552 { 553 TRACE_PACKET("AVFormatWriter::WriteChunk(%p, %ld, %p)\n", chunkBuffer, 554 chunkSize, encodeInfo); 555 556 StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie); 557 return cookie->WriteChunk(chunkBuffer, chunkSize, encodeInfo); 558 } 559 560 561 // #pragma mark - I/O hooks 562 563 564 /*static*/ int 565 AVFormatWriter::_Write(void* cookie, uint8* buffer, int bufferSize) 566 { 567 TRACE_IO("AVFormatWriter::_Write(%p, %p, %d)\n", 568 cookie, buffer, bufferSize); 569 570 AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie); 571 572 ssize_t written = writer->fTarget->Write(buffer, bufferSize); 573 574 TRACE_IO(" written: %ld\n", written); 575 return (int)written; 576 577 } 578 579 580 /*static*/ off_t 581 AVFormatWriter::_Seek(void* cookie, off_t offset, int whence) 582 { 583 TRACE_IO("AVFormatWriter::_Seek(%p, %lld, %d)\n", 584 cookie, offset, whence); 585 586 AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie); 587 588 BPositionIO* positionIO = dynamic_cast<BPositionIO*>(writer->fTarget); 589 if (positionIO == NULL) 590 return -1; 591 592 // Support for special file size retrieval API without seeking anywhere: 593 if (whence == AVSEEK_SIZE) { 594 off_t size; 595 if (positionIO->GetSize(&size) == B_OK) 596 return size; 597 return -1; 598 } 599 600 off_t position = positionIO->Seek(offset, whence); 601 TRACE_IO(" position: %lld\n", position); 602 if (position < 0) 603 return -1; 604 605 return position; 606 } 607 608 609