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