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