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