1 /* 2 * Copyright 2009, 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 <AutoDeleter.h> 15 #include <Autolock.h> 16 #include <ByteOrder.h> 17 #include <DataIO.h> 18 #include <MediaDefs.h> 19 #include <MediaFormats.h> 20 21 extern "C" { 22 #include "avformat.h" 23 } 24 25 #include "DemuxerTable.h" 26 #include "EncoderTable.h" 27 #include "gfx_util.h" 28 29 30 //#define TRACE_AVFORMAT_WRITER 31 #ifdef TRACE_AVFORMAT_WRITER 32 # define TRACE printf 33 # define TRACE_IO(a...) 34 # define TRACE_PACKET printf 35 #else 36 # define TRACE(a...) 37 # define TRACE_IO(a...) 38 # define TRACE_PACKET(a...) 39 #endif 40 41 #define ERROR(a...) fprintf(stderr, a) 42 43 44 static const size_t kIOBufferSize = 64 * 1024; 45 // TODO: This could depend on the BMediaFile creation flags, IIRC, 46 // they allow to specify a buffering mode. 47 48 // NOTE: The following works around some weird bug in libavformat. We 49 // have to open the AVFormatContext->AVStream->AVCodecContext, even though 50 // we are not interested in donig any encoding here!! 51 #define OPEN_CODEC_CONTEXT 1 52 #define GET_CONTEXT_DEFAULTS 0 53 54 55 // #pragma mark - AVFormatWriter::StreamCookie 56 57 58 class AVFormatWriter::StreamCookie { 59 public: 60 StreamCookie(AVFormatContext* context, 61 BLocker* streamLock); 62 virtual ~StreamCookie(); 63 64 status_t Init(const media_format* format, 65 const media_codec_info* codecInfo); 66 67 status_t WriteChunk(const void* chunkBuffer, 68 size_t chunkSize, 69 media_encode_info* encodeInfo); 70 71 status_t AddTrackInfo(uint32 code, const void* data, 72 size_t size, uint32 flags); 73 74 private: 75 AVFormatContext* fContext; 76 AVStream* fStream; 77 AVPacket fPacket; 78 // Since different threads may write to the target, 79 // we need to protect the file position and I/O by a lock. 80 BLocker* fStreamLock; 81 }; 82 83 84 85 AVFormatWriter::StreamCookie::StreamCookie(AVFormatContext* context, 86 BLocker* streamLock) 87 : 88 fContext(context), 89 fStream(NULL), 90 fStreamLock(streamLock) 91 { 92 av_init_packet(&fPacket); 93 } 94 95 96 AVFormatWriter::StreamCookie::~StreamCookie() 97 { 98 } 99 100 101 status_t 102 AVFormatWriter::StreamCookie::Init(const media_format* format, 103 const media_codec_info* codecInfo) 104 { 105 TRACE("AVFormatWriter::StreamCookie::Init()\n"); 106 107 BAutolock _(fStreamLock); 108 109 fPacket.stream_index = fContext->nb_streams; 110 fStream = av_new_stream(fContext, fPacket.stream_index); 111 112 if (fStream == NULL) { 113 TRACE(" failed to add new stream\n"); 114 return B_ERROR; 115 } 116 117 // TRACE(" fStream->codec: %p\n", fStream->codec); 118 // TODO: This is a hack for now! Use avcodec_find_encoder_by_name() 119 // or something similar... 120 fStream->codec->codec_id = (CodecID)codecInfo->sub_id; 121 if (fStream->codec->codec_id == CODEC_ID_NONE) 122 fStream->codec->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->codec->codec_type = CODEC_TYPE_VIDEO; 127 #if GET_CONTEXT_DEFAULTS 128 // NOTE: API example does not do this: 129 avcodec_get_context_defaults(fStream->codec); 130 #endif 131 // frame rate 132 fStream->codec->time_base.den = (int)format->u.raw_video.field_rate; 133 fStream->codec->time_base.num = 1; 134 // NOTE: API example does not do this: 135 // fStream->r_frame_rate.den = (int)format->u.raw_video.field_rate; 136 // fStream->r_frame_rate.num = 1; 137 // fStream->time_base.den = (int)format->u.raw_video.field_rate; 138 // fStream->time_base.num = 1; 139 // video size 140 fStream->codec->width = format->u.raw_video.display.line_width; 141 fStream->codec->height = format->u.raw_video.display.line_count; 142 // pixel aspect ratio 143 fStream->sample_aspect_ratio.num 144 = format->u.raw_video.pixel_width_aspect; 145 fStream->sample_aspect_ratio.den 146 = format->u.raw_video.pixel_height_aspect; 147 if (fStream->sample_aspect_ratio.num == 0 148 || fStream->sample_aspect_ratio.den == 0) { 149 av_reduce(&fStream->sample_aspect_ratio.num, 150 &fStream->sample_aspect_ratio.den, fStream->codec->width, 151 fStream->codec->height, 255); 152 } 153 154 fStream->codec->sample_aspect_ratio = fStream->sample_aspect_ratio; 155 // TODO: Don't hard code this... 156 fStream->codec->pix_fmt = PIX_FMT_YUV420P; 157 158 // Some formats want stream headers to be separate 159 if ((fContext->oformat->flags & AVFMT_GLOBALHEADER) != 0) 160 fStream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER; 161 } else if (format->type == B_MEDIA_RAW_AUDIO) { 162 fStream->codec->codec_type = CODEC_TYPE_AUDIO; 163 #if GET_CONTEXT_DEFAULTS 164 // NOTE: API example does not do this: 165 avcodec_get_context_defaults(fStream->codec); 166 #endif 167 // frame rate 168 fStream->codec->sample_rate = (int)format->u.raw_audio.frame_rate; 169 170 // channels 171 fStream->codec->channels = format->u.raw_audio.channel_count; 172 switch (format->u.raw_audio.format) { 173 case media_raw_audio_format::B_AUDIO_FLOAT: 174 fStream->codec->sample_fmt = SAMPLE_FMT_FLT; 175 break; 176 case media_raw_audio_format::B_AUDIO_DOUBLE: 177 fStream->codec->sample_fmt = SAMPLE_FMT_DBL; 178 break; 179 case media_raw_audio_format::B_AUDIO_INT: 180 fStream->codec->sample_fmt = SAMPLE_FMT_S32; 181 break; 182 case media_raw_audio_format::B_AUDIO_SHORT: 183 fStream->codec->sample_fmt = SAMPLE_FMT_S16; 184 break; 185 case media_raw_audio_format::B_AUDIO_UCHAR: 186 fStream->codec->sample_fmt = SAMPLE_FMT_U8; 187 break; 188 189 case media_raw_audio_format::B_AUDIO_CHAR: 190 default: 191 return B_MEDIA_BAD_FORMAT; 192 break; 193 } 194 if (format->u.raw_audio.channel_mask == 0) { 195 // guess the channel mask... 196 switch (format->u.raw_audio.channel_count) { 197 default: 198 case 2: 199 fStream->codec->channel_layout = CH_LAYOUT_STEREO; 200 break; 201 case 1: 202 fStream->codec->channel_layout = CH_LAYOUT_MONO; 203 break; 204 case 3: 205 fStream->codec->channel_layout = CH_LAYOUT_SURROUND; 206 break; 207 case 4: 208 fStream->codec->channel_layout = CH_LAYOUT_QUAD; 209 break; 210 case 5: 211 fStream->codec->channel_layout = CH_LAYOUT_5POINT0; 212 break; 213 case 6: 214 fStream->codec->channel_layout = CH_LAYOUT_5POINT1; 215 break; 216 case 8: 217 fStream->codec->channel_layout = CH_LAYOUT_7POINT1; 218 break; 219 case 10: 220 fStream->codec->channel_layout = CH_LAYOUT_7POINT1_WIDE; 221 break; 222 } 223 } else { 224 // The bits match 1:1 for media_multi_channels and FFmpeg defines. 225 fStream->codec->channel_layout = format->u.raw_audio.channel_mask; 226 } 227 } 228 229 TRACE(" stream->time_base: (%d/%d), codec->time_base: (%d/%d))\n", 230 fStream->time_base.num, fStream->time_base.den, 231 fStream->codec->time_base.num, fStream->codec->time_base.den); 232 233 return B_OK; 234 } 235 236 237 status_t 238 AVFormatWriter::StreamCookie::WriteChunk(const void* chunkBuffer, 239 size_t chunkSize, media_encode_info* encodeInfo) 240 { 241 TRACE_PACKET("AVFormatWriter::StreamCookie[%d]::WriteChunk(%p, %ld, " 242 "start_time: %lld)\n", fStream->index, chunkBuffer, chunkSize, 243 encodeInfo->start_time); 244 245 BAutolock _(fStreamLock); 246 247 // TODO: Probably the AVCodecEncoder needs to pass packet data 248 // in encodeInfo... 249 250 fPacket.data = const_cast<uint8_t*>((const uint8_t*)chunkBuffer); 251 fPacket.size = chunkSize; 252 253 fPacket.pts = int64_t((double)encodeInfo->start_time 254 * fStream->time_base.den / (1000000.0 * fStream->time_base.num) 255 + 0.5); 256 TRACE_PACKET(" PTS: %lld (stream->time_base: (%d/%d), " 257 "codec->time_base: (%d/%d))\n", fPacket.pts, 258 fStream->time_base.num, fStream->time_base.den, 259 fStream->codec->time_base.num, fStream->codec->time_base.den); 260 261 // From ffmpeg.c::do_audio_out(): 262 // TODO: 263 // if (enc->coded_frame && enc->coded_frame->pts != AV_NOPTS_VALUE) 264 // fPacket.pts = av_rescale_q(enc->coded_frame->pts, 265 // enc->time_base, ost->st->time_base); 266 267 268 #if 0 269 // TODO: Eventually, we need to write interleaved packets, but 270 // maybe we are only supposed to use this if we have actually 271 // more than one stream. For the moment, this crashes in AVPacket 272 // shuffling inside libavformat. Maybe if we want to use this, we 273 // need to allocate a separate AVPacket and copy the chunk buffer. 274 int result = av_interleaved_write_frame(fContext, &fPacket); 275 if (result < 0) 276 TRACE(" av_interleaved_write_frame(): %d\n", result); 277 #else 278 int result = av_write_frame(fContext, &fPacket); 279 if (result < 0) 280 TRACE(" av_write_frame(): %d\n", result); 281 #endif 282 283 return result == 0 ? B_OK : B_ERROR; 284 } 285 286 287 status_t 288 AVFormatWriter::StreamCookie::AddTrackInfo(uint32 code, 289 const void* data, size_t size, uint32 flags) 290 { 291 TRACE("AVFormatWriter::StreamCookie::AddTrackInfo(%lu, %p, %ld, %lu)\n", 292 code, data, size, flags); 293 294 BAutolock _(fStreamLock); 295 296 return B_NOT_SUPPORTED; 297 } 298 299 300 // #pragma mark - AVFormatWriter 301 302 303 AVFormatWriter::AVFormatWriter() 304 : 305 fContext(avformat_alloc_context()), 306 fHeaderWritten(false), 307 fStreamLock("stream lock") 308 { 309 TRACE("AVFormatWriter::AVFormatWriter\n"); 310 } 311 312 313 AVFormatWriter::~AVFormatWriter() 314 { 315 TRACE("AVFormatWriter::~AVFormatWriter\n"); 316 317 // Free the streams and close the AVCodecContexts 318 for(unsigned i = 0; i < fContext->nb_streams; i++) { 319 #if OPEN_CODEC_CONTEXT 320 // We only need to close the AVCodecContext when we opened it. 321 // This is experimental, see CommitHeader(). 322 if (fHeaderWritten) 323 avcodec_close(fContext->streams[i]->codec); 324 #endif 325 av_freep(&fContext->streams[i]->codec); 326 av_freep(&fContext->streams[i]); 327 } 328 329 av_free(fContext); 330 av_free(fIOContext.buffer); 331 } 332 333 334 // #pragma mark - 335 336 337 status_t 338 AVFormatWriter::Init(const media_file_format* fileFormat) 339 { 340 TRACE("AVFormatWriter::Init()\n"); 341 342 uint8* buffer = static_cast<uint8*>(av_malloc(kIOBufferSize)); 343 if (buffer == NULL) 344 return B_NO_MEMORY; 345 346 // Init I/O context with buffer and hook functions, pass ourself as 347 // cookie. 348 if (init_put_byte(&fIOContext, buffer, kIOBufferSize, 1, this, 349 0, _Write, _Seek) != 0) { 350 TRACE(" init_put_byte() failed!\n"); 351 return B_ERROR; 352 } 353 354 // Setup I/O hooks. This seems to be enough. 355 fContext->pb = &fIOContext; 356 357 // Set the AVOutputFormat according to fileFormat... 358 fContext->oformat = av_guess_format(fileFormat->short_name, 359 fileFormat->file_extension, fileFormat->mime_type); 360 if (fContext->oformat == NULL) { 361 TRACE(" failed to find AVOuputFormat for %s\n", 362 fileFormat->short_name); 363 return B_NOT_SUPPORTED; 364 } 365 366 TRACE(" found AVOuputFormat for %s: %s\n", fileFormat->short_name, 367 fContext->oformat->name); 368 369 return B_OK; 370 } 371 372 373 status_t 374 AVFormatWriter::SetCopyright(const char* copyright) 375 { 376 TRACE("AVFormatWriter::SetCopyright(%s)\n", copyright); 377 378 return B_NOT_SUPPORTED; 379 } 380 381 382 status_t 383 AVFormatWriter::CommitHeader() 384 { 385 TRACE("AVFormatWriter::CommitHeader\n"); 386 387 if (fContext == NULL) 388 return B_NO_INIT; 389 390 if (fHeaderWritten) 391 return B_NOT_ALLOWED; 392 393 // According to output_example.c, the output parameters must be set even 394 // if none are specified. In the example, this call is used after the 395 // streams have been created. 396 if (av_set_parameters(fContext, NULL) < 0) 397 return B_ERROR; 398 399 for (unsigned i = 0; i < fContext->nb_streams; i++) { 400 AVStream* stream = fContext->streams[i]; 401 #if OPEN_CODEC_CONTEXT 402 // NOTE: Experimental, this should not be needed. Especially, since 403 // we have no idea (in the future) what CodecID some encoder uses, 404 // it may be an encoder from a different plugin. 405 AVCodecContext* codecContext = stream->codec; 406 AVCodec* codec = avcodec_find_encoder(codecContext->codec_id); 407 if (codec == NULL || avcodec_open(codecContext, codec) < 0) { 408 TRACE(" stream[%u] - failed to open AVCodecContext\n", i); 409 } 410 #endif 411 TRACE(" stream[%u] time_base: (%d/%d), codec->time_base: (%d/%d)\n", 412 i, stream->time_base.num, stream->time_base.den, 413 stream->codec->time_base.num, stream->codec->time_base.den); 414 } 415 416 int result = av_write_header(fContext); 417 if (result < 0) 418 TRACE(" av_write_header(): %d\n", result); 419 420 // We need to close the codecs we opened, even in case of failure. 421 fHeaderWritten = true; 422 423 #ifdef TRACE_AVFORMAT_WRITER 424 TRACE(" wrote header\n"); 425 for (unsigned i = 0; i < fContext->nb_streams; i++) { 426 AVStream* stream = fContext->streams[i]; 427 TRACE(" stream[%u] time_base: (%d/%d), codec->time_base: (%d/%d)\n", 428 i, stream->time_base.num, stream->time_base.den, 429 stream->codec->time_base.num, stream->codec->time_base.den); 430 } 431 #endif // TRACE_AVFORMAT_WRITER 432 433 return result == 0 ? B_OK : B_ERROR; 434 } 435 436 437 status_t 438 AVFormatWriter::Flush() 439 { 440 TRACE("AVFormatWriter::Flush\n"); 441 442 return B_NOT_SUPPORTED; 443 } 444 445 446 status_t 447 AVFormatWriter::Close() 448 { 449 TRACE("AVFormatWriter::Close\n"); 450 451 if (fContext == NULL) 452 return B_NO_INIT; 453 454 if (!fHeaderWritten) 455 return B_NOT_ALLOWED; 456 457 int result = av_write_trailer(fContext); 458 if (result < 0) 459 TRACE(" av_write_trailer(): %d\n", result); 460 461 return result == 0 ? B_OK : B_ERROR; 462 } 463 464 465 status_t 466 AVFormatWriter::AllocateCookie(void** _cookie, const media_format* format, 467 const media_codec_info* codecInfo) 468 { 469 TRACE("AVFormatWriter::AllocateCookie()\n"); 470 471 if (fHeaderWritten) 472 return B_NOT_ALLOWED; 473 474 BAutolock _(fStreamLock); 475 476 if (_cookie == NULL) 477 return B_BAD_VALUE; 478 479 StreamCookie* cookie = new(std::nothrow) StreamCookie(fContext, 480 &fStreamLock); 481 482 status_t ret = cookie->Init(format, codecInfo); 483 if (ret != B_OK) { 484 delete cookie; 485 return ret; 486 } 487 488 *_cookie = cookie; 489 return B_OK; 490 } 491 492 493 status_t 494 AVFormatWriter::FreeCookie(void* _cookie) 495 { 496 BAutolock _(fStreamLock); 497 498 StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie); 499 delete cookie; 500 501 return B_OK; 502 } 503 504 505 // #pragma mark - 506 507 508 status_t 509 AVFormatWriter::SetCopyright(void* cookie, const char* copyright) 510 { 511 TRACE("AVFormatWriter::SetCopyright(%p, %s)\n", cookie, copyright); 512 513 return B_NOT_SUPPORTED; 514 } 515 516 517 status_t 518 AVFormatWriter::AddTrackInfo(void* _cookie, uint32 code, 519 const void* data, size_t size, uint32 flags) 520 { 521 TRACE("AVFormatWriter::AddTrackInfo(%lu, %p, %ld, %lu)\n", 522 code, data, size, flags); 523 524 StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie); 525 return cookie->AddTrackInfo(code, data, size, flags); 526 } 527 528 529 status_t 530 AVFormatWriter::WriteChunk(void* _cookie, const void* chunkBuffer, 531 size_t chunkSize, media_encode_info* encodeInfo) 532 { 533 TRACE_PACKET("AVFormatWriter::WriteChunk(%p, %ld, %p)\n", chunkBuffer, 534 chunkSize, encodeInfo); 535 536 StreamCookie* cookie = reinterpret_cast<StreamCookie*>(_cookie); 537 return cookie->WriteChunk(chunkBuffer, chunkSize, encodeInfo); 538 } 539 540 541 // #pragma mark - I/O hooks 542 543 544 /*static*/ int 545 AVFormatWriter::_Write(void* cookie, uint8* buffer, int bufferSize) 546 { 547 TRACE_IO("AVFormatWriter::_Write(%p, %p, %d)\n", 548 cookie, buffer, bufferSize); 549 550 AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie); 551 552 ssize_t written = writer->fTarget->Write(buffer, bufferSize); 553 554 TRACE_IO(" written: %ld\n", written); 555 return (int)written; 556 557 } 558 559 560 /*static*/ off_t 561 AVFormatWriter::_Seek(void* cookie, off_t offset, int whence) 562 { 563 TRACE_IO("AVFormatWriter::_Seek(%p, %lld, %d)\n", 564 cookie, offset, whence); 565 566 AVFormatWriter* writer = reinterpret_cast<AVFormatWriter*>(cookie); 567 568 BPositionIO* positionIO = dynamic_cast<BPositionIO*>(writer->fTarget); 569 if (positionIO == NULL) 570 return -1; 571 572 // Support for special file size retrieval API without seeking anywhere: 573 if (whence == AVSEEK_SIZE) { 574 off_t size; 575 if (positionIO->GetSize(&size) == B_OK) 576 return size; 577 return -1; 578 } 579 580 off_t position = positionIO->Seek(offset, whence); 581 TRACE_IO(" position: %lld\n", position); 582 if (position < 0) 583 return -1; 584 585 return position; 586 } 587 588 589