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