git.sesse.net Git - ffmpeg/blob - libavcodec/videotoolboxenc.c

   1 /*
   2  * copyright (c) 2015 Rick Kern <kernrj@gmail.com>
   3  *
   4  * This file is part of FFmpeg.
   5  *
   6  * FFmpeg is free software; you can redistribute it and/or
   7  * modify it under the terms of the GNU Lesser General Public
   8  * License as published by the Free Software Foundation; either
   9  * version 2.1 of the License, or (at your option) any later version.
  10  *
  11  * FFmpeg is distributed in the hope that it will be useful,
  12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14  * Lesser General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU Lesser General Public
  17  * License along with FFmpeg; if not, write to the Free Software
  18  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  19  */
  20
  21 #include <VideoToolbox/VideoToolbox.h>
  22 #include <CoreVideo/CoreVideo.h>
  23 #include <CoreMedia/CoreMedia.h>
  24 #include <TargetConditionals.h>
  25 #include <Availability.h>
  26 #include "avcodec.h"
  27 #include "libavutil/opt.h"
  28 #include "libavutil/avassert.h"
  29 #include "libavutil/atomic.h"
  30 #include "libavutil/avstring.h"
  31 #include "libavcodec/avcodec.h"
  32 #include "libavutil/pixdesc.h"
  33 #include "internal.h"
  34 #include <pthread.h>
  35
  36 #if !CONFIG_VT_BT2020
  37 # define kCVImageBufferColorPrimaries_ITU_R_2020   CFSTR("ITU_R_2020")
  38 # define kCVImageBufferTransferFunction_ITU_R_2020 CFSTR("ITU_R_2020")
  39 # define kCVImageBufferYCbCrMatrix_ITU_R_2020      CFSTR("ITU_R_2020")
  40 #endif
  41
  42 typedef enum VT_H264Profile {
  43     H264_PROF_AUTO,
  44     H264_PROF_BASELINE,
  45     H264_PROF_MAIN,
  46     H264_PROF_HIGH,
  47     H264_PROF_COUNT
  48 } VT_H264Profile;
  49
  50 typedef enum VTH264Entropy{
  51     VT_ENTROPY_NOT_SET,
  52     VT_CAVLC,
  53     VT_CABAC
  54 } VTH264Entropy;
  55
  56 static const uint8_t start_code[] = { 0, 0, 0, 1 };
  57
  58 typedef struct BufNode {
  59     CMSampleBufferRef cm_buffer;
  60     struct BufNode* next;
  61     int error;
  62 } BufNode;
  63
  64 typedef struct VTEncContext {
  65     AVClass *class;
  66     VTCompressionSessionRef session;
  67     CFStringRef ycbcr_matrix;
  68     CFStringRef color_primaries;
  69     CFStringRef transfer_function;
  70
  71     pthread_mutex_t lock;
  72     pthread_cond_t  cv_sample_sent;
  73
  74     int async_error;
  75
  76     BufNode *q_head;
  77     BufNode *q_tail;
  78
  79     int64_t frame_ct_out;
  80     int64_t frame_ct_in;
  81
  82     int64_t first_pts;
  83     int64_t dts_delta;
  84
  85     int64_t profile;
  86     int64_t level;
  87     int64_t entropy;
  88     int64_t realtime;
  89     int64_t frames_before;
  90     int64_t frames_after;
  91
  92     int64_t allow_sw;
  93
  94     bool flushing;
  95     bool has_b_frames;
  96     bool warned_color_range;
  97 } VTEncContext;
  98
  99 /**
 100  * NULL-safe release of *refPtr, and sets value to NULL.
 101  */
 102 static void vt_release_num(CFNumberRef* refPtr){
 103     if (!*refPtr) {
 104         return;
 105     }
 106
 107     CFRelease(*refPtr);
 108     *refPtr = NULL;
 109 }
 110
 111 static void set_async_error(VTEncContext *vtctx, int err)
 112 {
 113     BufNode *info;
 114
 115     pthread_mutex_lock(&vtctx->lock);
 116
 117     vtctx->async_error = err;
 118
 119     info = vtctx->q_head;
 120     vtctx->q_head = vtctx->q_tail = NULL;
 121
 122     while (info) {
 123         BufNode *next = info->next;
 124         CFRelease(info->cm_buffer);
 125         av_free(info);
 126         info = next;
 127     }
 128
 129     pthread_mutex_unlock(&vtctx->lock);
 130 }
 131
 132 static int vtenc_q_pop(VTEncContext *vtctx, bool wait, CMSampleBufferRef *buf)
 133 {
 134     BufNode *info;
 135
 136     pthread_mutex_lock(&vtctx->lock);
 137
 138     if (vtctx->async_error) {
 139         pthread_mutex_unlock(&vtctx->lock);
 140         return vtctx->async_error;
 141     }
 142
 143     if (vtctx->flushing && vtctx->frame_ct_in == vtctx->frame_ct_out) {
 144         *buf = NULL;
 145
 146         pthread_mutex_unlock(&vtctx->lock);
 147         return 0;
 148     }
 149
 150     while (!vtctx->q_head && !vtctx->async_error && wait) {
 151         pthread_cond_wait(&vtctx->cv_sample_sent, &vtctx->lock);
 152     }
 153
 154     if (!vtctx->q_head) {
 155         pthread_mutex_unlock(&vtctx->lock);
 156         *buf = NULL;
 157         return 0;
 158     }
 159
 160     info = vtctx->q_head;
 161     vtctx->q_head = vtctx->q_head->next;
 162     if (!vtctx->q_head) {
 163         vtctx->q_tail = NULL;
 164     }
 165
 166     pthread_mutex_unlock(&vtctx->lock);
 167
 168     *buf = info->cm_buffer;
 169     av_free(info);
 170
 171     vtctx->frame_ct_out++;
 172
 173     return 0;
 174 }
 175
 176 static void vtenc_q_push(VTEncContext *vtctx, CMSampleBufferRef buffer)
 177 {
 178     BufNode *info = av_malloc(sizeof(BufNode));
 179     if (!info) {
 180         set_async_error(vtctx, AVERROR(ENOMEM));
 181         return;
 182     }
 183
 184     CFRetain(buffer);
 185     info->cm_buffer = buffer;
 186     info->next = NULL;
 187
 188     pthread_mutex_lock(&vtctx->lock);
 189     pthread_cond_signal(&vtctx->cv_sample_sent);
 190
 191     if (!vtctx->q_head) {
 192         vtctx->q_head = info;
 193     } else {
 194         vtctx->q_tail->next = info;
 195     }
 196
 197     vtctx->q_tail = info;
 198
 199     pthread_mutex_unlock(&vtctx->lock);
 200 }
 201
 202 static CMVideoCodecType get_cm_codec_type(enum AVCodecID id)
 203 {
 204     switch (id) {
 205     case AV_CODEC_ID_H264: return kCMVideoCodecType_H264;
 206     default:               return 0;
 207     }
 208 }
 209
 210 static void vtenc_free_block(void *opaque, uint8_t *data)
 211 {
 212     CMBlockBufferRef block = opaque;
 213     CFRelease(block);
 214 }
 215
 216 /**
 217  * Get the parameter sets from a CMSampleBufferRef.
 218  * @param dst If *dst isn't NULL, the parameters are copied into existing
 219  *            memory. *dst_size must be set accordingly when *dst != NULL.
 220  *            If *dst is NULL, it will be allocated.
 221  *            In all cases, *dst_size is set to the number of bytes used starting
 222  *            at *dst.
 223  */
 224 static int get_params_size(
 225     AVCodecContext              *avctx,
 226     CMVideoFormatDescriptionRef vid_fmt,
 227     size_t                      *size)
 228 {
 229     size_t total_size = 0;
 230     size_t ps_count;
 231     int is_count_bad = 0;
 232     size_t i;
 233     int status;
 234     status = CMVideoFormatDescriptionGetH264ParameterSetAtIndex(vid_fmt,
 235                                                                 0,
 236                                                                 NULL,
 237                                                                 NULL,
 238                                                                 &ps_count,
 239                                                                 NULL);
 240     if (status) {
 241         is_count_bad = 1;
 242         ps_count     = 0;
 243         status       = 0;
 244     }
 245
 246     for (i = 0; i < ps_count || is_count_bad; i++) {
 247         const uint8_t *ps;
 248         size_t ps_size;
 249         status = CMVideoFormatDescriptionGetH264ParameterSetAtIndex(vid_fmt,
 250                                                                     i,
 251                                                                     &ps,
 252                                                                     &ps_size,
 253                                                                     NULL,
 254                                                                     NULL);
 255         if (status) {
 256             /*
 257              * When ps_count is invalid, status != 0 ends the loop normally
 258              * unless we didn't get any parameter sets.
 259              */
 260             if (i > 0 && is_count_bad) status = 0;
 261
 262             break;
 263         }
 264
 265         total_size += ps_size + sizeof(start_code);
 266     }
 267
 268     if (status) {
 269         av_log(avctx, AV_LOG_ERROR, "Error getting parameter set sizes: %d\n", status);
 270         return AVERROR_EXTERNAL;
 271     }
 272
 273     *size = total_size;
 274     return 0;
 275 }
 276
 277 static int copy_param_sets(
 278     AVCodecContext              *avctx,
 279     CMVideoFormatDescriptionRef vid_fmt,
 280     uint8_t                     *dst,
 281     size_t                      dst_size)
 282 {
 283     size_t ps_count;
 284     int is_count_bad = 0;
 285     int status;
 286     size_t offset = 0;
 287     size_t i;
 288
 289     status = CMVideoFormatDescriptionGetH264ParameterSetAtIndex(vid_fmt,
 290                                                                 0,
 291                                                                 NULL,
 292                                                                 NULL,
 293                                                                 &ps_count,
 294                                                                 NULL);
 295     if (status) {
 296         is_count_bad = 1;
 297         ps_count     = 0;
 298         status       = 0;
 299     }
 300
 301
 302     for (i = 0; i < ps_count || is_count_bad; i++) {
 303         const uint8_t *ps;
 304         size_t ps_size;
 305         size_t next_offset;
 306
 307         status = CMVideoFormatDescriptionGetH264ParameterSetAtIndex(vid_fmt,
 308                                                                     i,
 309                                                                     &ps,
 310                                                                     &ps_size,
 311                                                                     NULL,
 312                                                                     NULL);
 313         if (status) {
 314             if (i > 0 && is_count_bad) status = 0;
 315
 316             break;
 317         }
 318
 319         next_offset = offset + sizeof(start_code) + ps_size;
 320         if (dst_size < next_offset) {
 321             av_log(avctx, AV_LOG_ERROR, "Error: buffer too small for parameter sets.\n");
 322             return AVERROR_BUFFER_TOO_SMALL;
 323         }
 324
 325         memcpy(dst + offset, start_code, sizeof(start_code));
 326         offset += sizeof(start_code);
 327
 328         memcpy(dst + offset, ps, ps_size);
 329         offset = next_offset;
 330     }
 331
 332     if (status) {
 333         av_log(avctx, AV_LOG_ERROR, "Error getting parameter set data: %d\n", status);
 334         return AVERROR_EXTERNAL;
 335     }
 336
 337     return 0;
 338 }
 339
 340 static int set_extradata(AVCodecContext *avctx, CMSampleBufferRef sample_buffer)
 341 {
 342     CMVideoFormatDescriptionRef vid_fmt;
 343     size_t total_size;
 344     int status;
 345
 346     vid_fmt = CMSampleBufferGetFormatDescription(sample_buffer);
 347     if (!vid_fmt) {
 348         av_log(avctx, AV_LOG_ERROR, "No video format.\n");
 349         return AVERROR_EXTERNAL;
 350     }
 351
 352     status = get_params_size(avctx, vid_fmt, &total_size);
 353     if (status) {
 354         av_log(avctx, AV_LOG_ERROR, "Could not get parameter sets.\n");
 355         return status;
 356     }
 357
 358     avctx->extradata = av_malloc(total_size);
 359     if (!avctx->extradata) {
 360         return AVERROR(ENOMEM);
 361     }
 362     avctx->extradata_size = total_size;
 363
 364     status = copy_param_sets(avctx, vid_fmt, avctx->extradata, total_size);
 365
 366     if (status) {
 367         av_log(avctx, AV_LOG_ERROR, "Could not copy param sets.\n");
 368         return status;
 369     }
 370
 371     return 0;
 372 }
 373
 374 static void vtenc_output_callback(
 375     void *ctx,
 376     void *sourceFrameCtx,
 377     OSStatus status,
 378     VTEncodeInfoFlags flags,
 379     CMSampleBufferRef sample_buffer)
 380 {
 381     AVCodecContext *avctx = ctx;
 382     VTEncContext   *vtctx = avctx->priv_data;
 383
 384     if (vtctx->async_error) {
 385         if(sample_buffer) CFRelease(sample_buffer);
 386         return;
 387     }
 388
 389     if (status || !sample_buffer) {
 390         av_log(avctx, AV_LOG_ERROR, "Error encoding frame: %d\n", (int)status);
 391         set_async_error(vtctx, AVERROR_EXTERNAL);
 392         return;
 393     }
 394
 395     if (!avctx->extradata && (avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER)) {
 396         int set_status = set_extradata(avctx, sample_buffer);
 397         if (set_status) {
 398             set_async_error(vtctx, set_status);
 399             return;
 400         }
 401     }
 402
 403     vtenc_q_push(vtctx, sample_buffer);
 404 }
 405
 406 static int get_length_code_size(
 407     AVCodecContext    *avctx,
 408     CMSampleBufferRef sample_buffer,
 409     size_t            *size)
 410 {
 411     CMVideoFormatDescriptionRef vid_fmt;
 412     int isize;
 413     int status;
 414
 415     vid_fmt = CMSampleBufferGetFormatDescription(sample_buffer);
 416     if (!vid_fmt) {
 417         av_log(avctx, AV_LOG_ERROR, "Error getting buffer format description.\n");
 418         return AVERROR_EXTERNAL;
 419     }
 420
 421     status = CMVideoFormatDescriptionGetH264ParameterSetAtIndex(vid_fmt,
 422                                                                 0,
 423                                                                 NULL,
 424                                                                 NULL,
 425                                                                 NULL,
 426                                                                 &isize);
 427     if (status) {
 428         av_log(avctx, AV_LOG_ERROR, "Error getting length code size: %d\n", status);
 429         return AVERROR_EXTERNAL;
 430     }
 431
 432     *size = isize;
 433     return 0;
 434 }
 435
 436 /*
 437  * Returns true on success.
 438  *
 439  * If profile_level_val is NULL and this method returns true, don't specify the
 440  * profile/level to the encoder.
 441  */
 442 static bool get_vt_profile_level(AVCodecContext *avctx,
 443                                  CFStringRef    *profile_level_val)
 444 {
 445     VTEncContext *vtctx = avctx->priv_data;
 446     int64_t profile = vtctx->profile;
 447
 448     if (profile == H264_PROF_AUTO && vtctx->level) {
 449         //Need to pick a profile if level is not auto-selected.
 450         profile = vtctx->has_b_frames ? H264_PROF_MAIN : H264_PROF_BASELINE;
 451     }
 452
 453     *profile_level_val = NULL;
 454
 455     switch (profile) {
 456         case H264_PROF_AUTO:
 457             return true;
 458
 459         case H264_PROF_BASELINE:
 460             switch (vtctx->level) {
 461                 case  0: *profile_level_val = kVTProfileLevel_H264_Baseline_AutoLevel; break;
 462                 case 13: *profile_level_val = kVTProfileLevel_H264_Baseline_1_3;       break;
 463                 case 30: *profile_level_val = kVTProfileLevel_H264_Baseline_3_0;       break;
 464                 case 31: *profile_level_val = kVTProfileLevel_H264_Baseline_3_1;       break;
 465                 case 32: *profile_level_val = kVTProfileLevel_H264_Baseline_3_2;       break;
 466                 case 40: *profile_level_val = kVTProfileLevel_H264_Baseline_4_0;       break;
 467                 case 41: *profile_level_val = kVTProfileLevel_H264_Baseline_4_1;       break;
 468                 case 42: *profile_level_val = kVTProfileLevel_H264_Baseline_4_2;       break;
 469                 case 50: *profile_level_val = kVTProfileLevel_H264_Baseline_5_0;       break;
 470                 case 51: *profile_level_val = kVTProfileLevel_H264_Baseline_5_1;       break;
 471                 case 52: *profile_level_val = kVTProfileLevel_H264_Baseline_5_2;       break;
 472             }
 473             break;
 474
 475         case H264_PROF_MAIN:
 476             switch (vtctx->level) {
 477                 case  0: *profile_level_val = kVTProfileLevel_H264_Main_AutoLevel; break;
 478                 case 30: *profile_level_val = kVTProfileLevel_H264_Main_3_0;       break;
 479                 case 31: *profile_level_val = kVTProfileLevel_H264_Main_3_1;       break;
 480                 case 32: *profile_level_val = kVTProfileLevel_H264_Main_3_2;       break;
 481                 case 40: *profile_level_val = kVTProfileLevel_H264_Main_4_0;       break;
 482                 case 41: *profile_level_val = kVTProfileLevel_H264_Main_4_1;       break;
 483                 case 42: *profile_level_val = kVTProfileLevel_H264_Main_4_2;       break;
 484                 case 50: *profile_level_val = kVTProfileLevel_H264_Main_5_0;       break;
 485                 case 51: *profile_level_val = kVTProfileLevel_H264_Main_5_1;       break;
 486                 case 52: *profile_level_val = kVTProfileLevel_H264_Main_5_2;       break;
 487             }
 488             break;
 489
 490         case H264_PROF_HIGH:
 491             switch (vtctx->level) {
 492                 case  0: *profile_level_val = kVTProfileLevel_H264_High_AutoLevel; break;
 493                 case 30: *profile_level_val = kVTProfileLevel_H264_High_3_0;       break;
 494                 case 31: *profile_level_val = kVTProfileLevel_H264_High_3_1;       break;
 495                 case 32: *profile_level_val = kVTProfileLevel_H264_High_3_2;       break;
 496                 case 40: *profile_level_val = kVTProfileLevel_H264_High_4_0;       break;
 497                 case 41: *profile_level_val = kVTProfileLevel_H264_High_4_1;       break;
 498                 case 42: *profile_level_val = kVTProfileLevel_H264_High_4_2;       break;
 499                 case 50: *profile_level_val = kVTProfileLevel_H264_High_5_0;       break;
 500                 case 51: *profile_level_val = kVTProfileLevel_H264_High_5_1;       break;
 501                 case 52: *profile_level_val = kVTProfileLevel_H264_High_5_2;       break;
 502             }
 503             break;
 504     }
 505
 506     if (!*profile_level_val) {
 507         av_log(avctx, AV_LOG_ERROR, "Invalid Profile/Level.\n");
 508         return false;
 509     }
 510
 511     return true;
 512 }
 513
 514 static int get_cv_pixel_format(AVCodecContext* avctx,
 515                                enum AVPixelFormat fmt,
 516                                enum AVColorRange range,
 517                                int* av_pixel_format,
 518                                int* range_guessed)
 519 {
 520     if (range_guessed) *range_guessed = range != AVCOL_RANGE_MPEG &&
 521                                         range != AVCOL_RANGE_JPEG;
 522
 523     //MPEG range is used when no range is set
 524     if (fmt == AV_PIX_FMT_NV12) {
 525         *av_pixel_format = range == AVCOL_RANGE_JPEG ?
 526                                         kCVPixelFormatType_420YpCbCr8BiPlanarFullRange :
 527                                         kCVPixelFormatType_420YpCbCr8BiPlanarVideoRange;
 528     } else if (fmt == AV_PIX_FMT_YUV420P) {
 529         *av_pixel_format = range == AVCOL_RANGE_JPEG ?
 530                                         kCVPixelFormatType_420YpCbCr8PlanarFullRange :
 531                                         kCVPixelFormatType_420YpCbCr8Planar;
 532     } else {
 533         return AVERROR(EINVAL);
 534     }
 535
 536     return 0;
 537 }
 538
 539 static void add_color_attr(AVCodecContext *avctx, CFMutableDictionaryRef dict) {
 540     VTEncContext *vtctx = avctx->priv_data;
 541
 542     if (vtctx->color_primaries) {
 543         CFDictionarySetValue(dict,
 544                              kCVImageBufferColorPrimariesKey,
 545                              vtctx->color_primaries);
 546     }
 547
 548     if (vtctx->transfer_function) {
 549         CFDictionarySetValue(dict,
 550                              kCVImageBufferTransferFunctionKey,
 551                              vtctx->transfer_function);
 552     }
 553
 554     if (vtctx->ycbcr_matrix) {
 555         CFDictionarySetValue(dict,
 556                              kCVImageBufferYCbCrMatrixKey,
 557                              vtctx->ycbcr_matrix);
 558     }
 559 }
 560
 561 static int create_cv_pixel_buffer_info(AVCodecContext* avctx,
 562                                        CFMutableDictionaryRef* dict)
 563 {
 564     CFNumberRef cv_color_format_num = NULL;
 565     CFNumberRef width_num = NULL;
 566     CFNumberRef height_num = NULL;
 567     CFMutableDictionaryRef pixel_buffer_info = NULL;
 568     int cv_color_format;
 569     int status = get_cv_pixel_format(avctx,
 570                                      avctx->pix_fmt,
 571                                      avctx->color_range,
 572                                      &cv_color_format,
 573                                      NULL);
 574     if (status) return status;
 575
 576     pixel_buffer_info = CFDictionaryCreateMutable(
 577                             kCFAllocatorDefault,
 578                             20,
 579                             &kCFCopyStringDictionaryKeyCallBacks,
 580                             &kCFTypeDictionaryValueCallBacks);
 581
 582     if (!pixel_buffer_info) goto pbinfo_nomem;
 583
 584     cv_color_format_num = CFNumberCreate(kCFAllocatorDefault,
 585                                          kCFNumberSInt32Type,
 586                                          &cv_color_format);
 587     if (!cv_color_format_num) goto pbinfo_nomem;
 588
 589     CFDictionarySetValue(pixel_buffer_info,
 590                          kCVPixelBufferPixelFormatTypeKey,
 591                          cv_color_format_num);
 592     vt_release_num(&cv_color_format_num);
 593
 594     width_num = CFNumberCreate(kCFAllocatorDefault,
 595                                kCFNumberSInt32Type,
 596                                &avctx->width);
 597     if (!width_num) return AVERROR(ENOMEM);
 598
 599     CFDictionarySetValue(pixel_buffer_info,
 600                          kCVPixelBufferWidthKey,
 601                          width_num);
 602     vt_release_num(&width_num);
 603
 604     height_num = CFNumberCreate(kCFAllocatorDefault,
 605                                 kCFNumberSInt32Type,
 606                                 &avctx->height);
 607     if (!height_num) goto pbinfo_nomem;
 608
 609     CFDictionarySetValue(pixel_buffer_info,
 610                          kCVPixelBufferHeightKey,
 611                          height_num);
 612     vt_release_num(&height_num);
 613
 614     add_color_attr(avctx, pixel_buffer_info);
 615
 616     *dict = pixel_buffer_info;
 617     return 0;
 618
 619 pbinfo_nomem:
 620     vt_release_num(&cv_color_format_num);
 621     vt_release_num(&width_num);
 622     vt_release_num(&height_num);
 623     if (pixel_buffer_info) CFRelease(pixel_buffer_info);
 624
 625     return AVERROR(ENOMEM);
 626 }
 627
 628 static int get_cv_color_primaries(AVCodecContext *avctx,
 629                                   CFStringRef *primaries)
 630 {
 631     enum AVColorPrimaries pri = avctx->color_primaries;
 632     switch (pri) {
 633         case AVCOL_PRI_UNSPECIFIED:
 634             *primaries = NULL;
 635             break;
 636
 637         case AVCOL_PRI_BT709:
 638             *primaries = kCVImageBufferColorPrimaries_ITU_R_709_2;
 639             break;
 640
 641         case AVCOL_PRI_BT2020:
 642             *primaries = kCVImageBufferColorPrimaries_ITU_R_2020;
 643             break;
 644
 645         default:
 646             av_log(avctx, AV_LOG_ERROR, "Color primaries %s is not supported.\n", av_color_primaries_name(pri));
 647             *primaries = NULL;
 648             return -1;
 649     }
 650
 651     return 0;
 652 }
 653
 654 static int get_cv_transfer_function(AVCodecContext *avctx,
 655                                     CFStringRef *transfer_fnc,
 656                                     CFNumberRef *gamma_level)
 657 {
 658     enum AVColorTransferCharacteristic trc = avctx->color_trc;
 659     Float32 gamma;
 660     *gamma_level = NULL;
 661
 662     switch (trc) {
 663         case AVCOL_TRC_UNSPECIFIED:
 664             *transfer_fnc = NULL;
 665             break;
 666
 667         case AVCOL_TRC_BT709:
 668             *transfer_fnc = kCVImageBufferTransferFunction_ITU_R_709_2;
 669             break;
 670
 671         case AVCOL_TRC_SMPTE240M:
 672             *transfer_fnc = kCVImageBufferTransferFunction_SMPTE_240M_1995;
 673             break;
 674
 675         case AVCOL_TRC_GAMMA22:
 676             gamma = 2.2;
 677             *transfer_fnc = kCVImageBufferTransferFunction_UseGamma;
 678             *gamma_level = CFNumberCreate(NULL, kCFNumberFloat32Type, &gamma);
 679             break;
 680
 681         case AVCOL_TRC_GAMMA28:
 682             gamma = 2.8;
 683             *transfer_fnc = kCVImageBufferTransferFunction_UseGamma;
 684             *gamma_level = CFNumberCreate(NULL, kCFNumberFloat32Type, &gamma);
 685             break;
 686
 687         case AVCOL_TRC_BT2020_10:
 688         case AVCOL_TRC_BT2020_12:
 689             *transfer_fnc = kCVImageBufferTransferFunction_ITU_R_2020;
 690             break;
 691
 692         default:
 693             av_log(avctx, AV_LOG_ERROR, "Transfer function %s is not supported.\n", av_color_transfer_name(trc));
 694             return -1;
 695     }
 696
 697     return 0;
 698 }
 699
 700 static int get_cv_ycbcr_matrix(AVCodecContext *avctx, CFStringRef *matrix) {
 701     switch(avctx->colorspace) {
 702         case AVCOL_SPC_BT709:
 703             *matrix = kCVImageBufferYCbCrMatrix_ITU_R_709_2;
 704             break;
 705
 706         case AVCOL_SPC_UNSPECIFIED:
 707             *matrix = NULL;
 708             break;
 709
 710         case AVCOL_SPC_BT470BG:
 711         case AVCOL_SPC_SMPTE170M:
 712             *matrix = kCVImageBufferYCbCrMatrix_ITU_R_601_4;
 713             break;
 714
 715         case AVCOL_SPC_SMPTE240M:
 716             *matrix = kCVImageBufferYCbCrMatrix_SMPTE_240M_1995;
 717             break;
 718
 719         case AVCOL_SPC_BT2020_NCL:
 720             *matrix = kCVImageBufferYCbCrMatrix_ITU_R_2020;
 721             break;
 722
 723         default:
 724             av_log(avctx, AV_LOG_ERROR, "Color space %s is not supported.\n", av_color_space_name(avctx->colorspace));
 725             return -1;
 726     }
 727
 728     return 0;
 729 }
 730
 731
 732 static av_cold int vtenc_init(AVCodecContext *avctx)
 733 {
 734     CFMutableDictionaryRef enc_info;
 735     CFMutableDictionaryRef pixel_buffer_info;
 736     CMVideoCodecType       codec_type;
 737     VTEncContext           *vtctx = avctx->priv_data;
 738     CFStringRef            profile_level;
 739     SInt32                 bit_rate = avctx->bit_rate;
 740     CFNumberRef            bit_rate_num;
 741     CFBooleanRef           has_b_frames_cfbool;
 742     CFNumberRef            gamma_level;
 743     int                    status;
 744
 745     codec_type = get_cm_codec_type(avctx->codec_id);
 746     if (!codec_type) {
 747         av_log(avctx, AV_LOG_ERROR, "Error: no mapping for AVCodecID %d\n", avctx->codec_id);
 748         return AVERROR(EINVAL);
 749     }
 750
 751     vtctx->has_b_frames = avctx->max_b_frames > 0;
 752     if(vtctx->has_b_frames && vtctx->profile == H264_PROF_BASELINE){
 753         av_log(avctx, AV_LOG_WARNING, "Cannot use B-frames with baseline profile. Output will not contain B-frames.\n");
 754         vtctx->has_b_frames = false;
 755     }
 756
 757     if (vtctx->entropy == VT_CABAC && vtctx->profile == H264_PROF_BASELINE) {
 758         av_log(avctx, AV_LOG_WARNING, "CABAC entropy requires 'main' or 'high' profile, but baseline was requested. Encode will not use CABAC entropy.\n");
 759         vtctx->entropy = VT_ENTROPY_NOT_SET;
 760     }
 761
 762     if (!get_vt_profile_level(avctx, &profile_level)) return AVERROR(EINVAL);
 763
 764     vtctx->session = NULL;
 765
 766     enc_info = CFDictionaryCreateMutable(
 767         kCFAllocatorDefault,
 768         20,
 769         &kCFCopyStringDictionaryKeyCallBacks,
 770         &kCFTypeDictionaryValueCallBacks
 771     );
 772
 773     if (!enc_info) return AVERROR(ENOMEM);
 774
 775 #if !TARGET_OS_IPHONE
 776     if (!vtctx->allow_sw) {
 777         CFDictionarySetValue(enc_info, kVTVideoEncoderSpecification_RequireHardwareAcceleratedVideoEncoder, kCFBooleanTrue);
 778     } else {
 779         CFDictionarySetValue(enc_info, kVTVideoEncoderSpecification_EnableHardwareAcceleratedVideoEncoder,  kCFBooleanTrue);
 780     }
 781 #endif
 782
 783     if (avctx->pix_fmt != AV_PIX_FMT_VIDEOTOOLBOX) {
 784         status = create_cv_pixel_buffer_info(avctx, &pixel_buffer_info);
 785         if (status) {
 786             CFRelease(enc_info);
 787             return status;
 788         }
 789     } else {
 790         pixel_buffer_info = NULL;
 791     }
 792
 793     status = VTCompressionSessionCreate(
 794         kCFAllocatorDefault,
 795         avctx->width,
 796         avctx->height,
 797         codec_type,
 798         enc_info,
 799         pixel_buffer_info,
 800         kCFAllocatorDefault,
 801         vtenc_output_callback,
 802         avctx,
 803         &vtctx->session
 804     );
 805
 806     if (pixel_buffer_info) CFRelease(pixel_buffer_info);
 807     CFRelease(enc_info);
 808
 809     if (status || !vtctx->session) {
 810         av_log(avctx, AV_LOG_ERROR, "Error: cannot create compression session: %d\n", status);
 811
 812 #if !TARGET_OS_IPHONE
 813         if (!vtctx->allow_sw) {
 814             av_log(avctx, AV_LOG_ERROR, "Try -allow_sw 1. The hardware encoder may be busy, or not supported.\n");
 815         }
 816 #endif
 817
 818         return AVERROR_EXTERNAL;
 819     }
 820
 821     bit_rate_num = CFNumberCreate(kCFAllocatorDefault,
 822                                   kCFNumberSInt32Type,
 823                                   &bit_rate);
 824     if (!bit_rate_num) return AVERROR(ENOMEM);
 825
 826     status = VTSessionSetProperty(vtctx->session,
 827                                   kVTCompressionPropertyKey_AverageBitRate,
 828                                   bit_rate_num);
 829     CFRelease(bit_rate_num);
 830
 831     if (status) {
 832         av_log(avctx, AV_LOG_ERROR, "Error setting bitrate property: %d\n", status);
 833         return AVERROR_EXTERNAL;
 834     }
 835
 836     if (profile_level) {
 837         status = VTSessionSetProperty(vtctx->session,
 838                                       kVTCompressionPropertyKey_ProfileLevel,
 839                                       profile_level);
 840         if (status) {
 841             av_log(avctx, AV_LOG_ERROR, "Error setting profile/level property: %d\n", status);
 842             return AVERROR_EXTERNAL;
 843         }
 844     }
 845
 846     if (avctx->gop_size > 0) {
 847         CFNumberRef interval = CFNumberCreate(kCFAllocatorDefault,
 848                                               kCFNumberIntType,
 849                                               &avctx->gop_size);
 850         if (!interval) {
 851             return AVERROR(ENOMEM);
 852         }
 853
 854         status = VTSessionSetProperty(vtctx->session,
 855                                       kVTCompressionPropertyKey_MaxKeyFrameInterval,
 856                                       interval);
 857         CFRelease(interval);
 858
 859         if (status) {
 860             av_log(avctx, AV_LOG_ERROR, "Error setting 'max key-frame interval' property: %d\n", status);
 861             return AVERROR_EXTERNAL;
 862         }
 863     }
 864
 865     if (vtctx->frames_before) {
 866         status = VTSessionSetProperty(vtctx->session,
 867                                       kVTCompressionPropertyKey_MoreFramesBeforeStart,
 868                                       kCFBooleanTrue);
 869
 870         if (status == kVTPropertyNotSupportedErr) {
 871             av_log(avctx, AV_LOG_WARNING, "frames_before property is not supported on this device. Ignoring.\n");
 872         } else if (status) {
 873             av_log(avctx, AV_LOG_ERROR, "Error setting frames_before property: %d\n", status);
 874         }
 875     }
 876
 877     if (vtctx->frames_after) {
 878         status = VTSessionSetProperty(vtctx->session,
 879                                       kVTCompressionPropertyKey_MoreFramesAfterEnd,
 880                                       kCFBooleanTrue);
 881
 882         if (status == kVTPropertyNotSupportedErr) {
 883             av_log(avctx, AV_LOG_WARNING, "frames_after property is not supported on this device. Ignoring.\n");
 884         } else if (status) {
 885             av_log(avctx, AV_LOG_ERROR, "Error setting frames_after property: %d\n", status);
 886         }
 887     }
 888
 889     if (avctx->sample_aspect_ratio.num != 0) {
 890         CFNumberRef num;
 891         CFNumberRef den;
 892         CFMutableDictionaryRef par;
 893         AVRational *avpar = &avctx->sample_aspect_ratio;
 894
 895         av_reduce(&avpar->num, &avpar->den,
 896                    avpar->num,  avpar->den,
 897                   0xFFFFFFFF);
 898
 899         num = CFNumberCreate(kCFAllocatorDefault,
 900                              kCFNumberIntType,
 901                              &avpar->num);
 902
 903         den = CFNumberCreate(kCFAllocatorDefault,
 904                              kCFNumberIntType,
 905                              &avpar->den);
 906
 907
 908
 909         par = CFDictionaryCreateMutable(kCFAllocatorDefault,
 910                                         2,
 911                                         &kCFCopyStringDictionaryKeyCallBacks,
 912                                         &kCFTypeDictionaryValueCallBacks);
 913
 914         if (!par || !num || !den) {
 915             if (par) CFRelease(par);
 916             if (num) CFRelease(num);
 917             if (den) CFRelease(den);
 918
 919             return AVERROR(ENOMEM);
 920         }
 921
 922         CFDictionarySetValue(
 923             par,
 924             kCMFormatDescriptionKey_PixelAspectRatioHorizontalSpacing,
 925             num);
 926
 927         CFDictionarySetValue(
 928             par,
 929             kCMFormatDescriptionKey_PixelAspectRatioVerticalSpacing,
 930             den);
 931
 932         status = VTSessionSetProperty(vtctx->session,
 933                                       kVTCompressionPropertyKey_PixelAspectRatio,
 934                                       par);
 935
 936         CFRelease(par);
 937         CFRelease(num);
 938         CFRelease(den);
 939
 940         if (status) {
 941             av_log(avctx,
 942                    AV_LOG_ERROR,
 943                    "Error setting pixel aspect ratio to %d:%d: %d.\n",
 944                    avctx->sample_aspect_ratio.num,
 945                    avctx->sample_aspect_ratio.den,
 946                    status);
 947
 948             return AVERROR_EXTERNAL;
 949         }
 950     }
 951
 952     status = get_cv_transfer_function(avctx, &vtctx->transfer_function, &gamma_level);
 953     if (!status && vtctx->transfer_function) {
 954         status = VTSessionSetProperty(vtctx->session,
 955                                       kVTCompressionPropertyKey_TransferFunction,
 956                                       vtctx->transfer_function);
 957
 958         if (status) {
 959             av_log(avctx, AV_LOG_WARNING, "Could not set transfer function: %d\n", status);
 960         }
 961     }
 962
 963     status = get_cv_ycbcr_matrix(avctx, &vtctx->ycbcr_matrix);
 964     if (!status && vtctx->ycbcr_matrix) {
 965         status = VTSessionSetProperty(vtctx->session,
 966                                       kVTCompressionPropertyKey_YCbCrMatrix,
 967                                       vtctx->ycbcr_matrix);
 968
 969         if (status) {
 970             av_log(avctx, AV_LOG_WARNING, "Could not set ycbcr matrix: %d\n", status);
 971         }
 972     }
 973
 974     status = get_cv_color_primaries(avctx, &vtctx->color_primaries);
 975     if (!status && vtctx->color_primaries) {
 976         status = VTSessionSetProperty(vtctx->session,
 977                                       kVTCompressionPropertyKey_ColorPrimaries,
 978                                       vtctx->color_primaries);
 979
 980         if (status) {
 981             av_log(avctx, AV_LOG_WARNING, "Could not set color primaries: %d\n", status);
 982         }
 983     }
 984
 985     if (!status && gamma_level) {
 986         status = VTSessionSetProperty(vtctx->session,
 987                                       kCVImageBufferGammaLevelKey,
 988                                       gamma_level);
 989
 990         if (status) {
 991             av_log(avctx, AV_LOG_WARNING, "Could not set gamma level: %d\n", status);
 992         }
 993     }
 994
 995     if (!vtctx->has_b_frames) {
 996         status = VTSessionSetProperty(vtctx->session,
 997                                       kVTCompressionPropertyKey_AllowFrameReordering,
 998                                       kCFBooleanFalse);
 999
1000         if (status) {
1001             av_log(avctx, AV_LOG_ERROR, "Error setting 'allow frame reordering' property: %d\n", status);
1002             return AVERROR_EXTERNAL;
1003         }
1004     }
1005
1006     if (vtctx->entropy != VT_ENTROPY_NOT_SET) {
1007         CFStringRef entropy = vtctx->entropy == VT_CABAC ?
1008                                 kVTH264EntropyMode_CABAC:
1009                                 kVTH264EntropyMode_CAVLC;
1010
1011         status = VTSessionSetProperty(vtctx->session,
1012                                       kVTCompressionPropertyKey_H264EntropyMode,
1013                                       entropy);
1014
1015         if (status) {
1016             av_log(avctx, AV_LOG_ERROR, "Error setting entropy property: %d\n", status);
1017             return AVERROR_EXTERNAL;
1018         }
1019     }
1020
1021     if (vtctx->realtime) {
1022         status = VTSessionSetProperty(vtctx->session,
1023                                       kVTCompressionPropertyKey_RealTime,
1024                                       kCFBooleanTrue);
1025
1026         if (status) {
1027             av_log(avctx, AV_LOG_ERROR, "Error setting realtime property: %d\n", status);
1028         }
1029     }
1030
1031     status = VTCompressionSessionPrepareToEncodeFrames(vtctx->session);
1032     if (status) {
1033         av_log(avctx, AV_LOG_ERROR, "Error: cannot prepare encoder: %d\n", status);
1034         return AVERROR_EXTERNAL;
1035     }
1036
1037     pthread_mutex_init(&vtctx->lock, NULL);
1038     pthread_cond_init(&vtctx->cv_sample_sent, NULL);
1039     vtctx->dts_delta = vtctx->has_b_frames ? -1 : 0;
1040
1041     status = VTSessionCopyProperty(vtctx->session,
1042                                    kVTCompressionPropertyKey_AllowFrameReordering,
1043                                    kCFAllocatorDefault,
1044                                    &has_b_frames_cfbool);
1045
1046     if (!status) {
1047         //Some devices don't output B-frames for main profile, even if requested.
1048         vtctx->has_b_frames = CFBooleanGetValue(has_b_frames_cfbool);
1049         CFRelease(has_b_frames_cfbool);
1050     }
1051     avctx->has_b_frames = vtctx->has_b_frames;
1052
1053     return 0;
1054 }
1055
1056 static void vtenc_get_frame_info(CMSampleBufferRef buffer, bool *is_key_frame)
1057 {
1058     CFArrayRef      attachments;
1059     CFDictionaryRef attachment;
1060     CFBooleanRef    not_sync;
1061     CFIndex         len;
1062
1063     attachments = CMSampleBufferGetSampleAttachmentsArray(buffer, false);
1064     len = !attachments ? 0 : CFArrayGetCount(attachments);
1065
1066     if (!len) {
1067         *is_key_frame = true;
1068         return;
1069     }
1070
1071     attachment = CFArrayGetValueAtIndex(attachments, 0);
1072
1073     if (CFDictionaryGetValueIfPresent(attachment,
1074                                       kCMSampleAttachmentKey_NotSync,
1075                                       (const void **)&not_sync))
1076     {
1077         *is_key_frame = !CFBooleanGetValue(not_sync);
1078     } else {
1079         *is_key_frame = true;
1080     }
1081 }
1082
1083 /**
1084  * Replaces length codes with H.264 Annex B start codes.
1085  * length_code_size must equal sizeof(start_code).
1086  * On failure, the contents of data may have been modified.
1087  *
1088  * @param length_code_size Byte length of each length code
1089  * @param data Call with NAL units prefixed with length codes.
1090  *             On success, the length codes are replace with
1091  *             start codes.
1092  * @param size Length of data, excluding any padding.
1093  * @return 0 on success
1094  *         AVERROR_BUFFER_TOO_SMALL if length code size is smaller
1095  *         than a start code or if a length_code in data specifies
1096  *         data beyond the end of its buffer.
1097  */
1098 static int replace_length_codes(size_t  length_code_size,
1099                                 uint8_t *data,
1100                                 size_t  size)
1101 {
1102     size_t remaining_size = size;
1103
1104     if (length_code_size != sizeof(start_code)) {
1105         av_log(NULL, AV_LOG_ERROR, "Start code size and length code size not equal.\n");
1106         return AVERROR_BUFFER_TOO_SMALL;
1107     }
1108
1109     while (remaining_size > 0) {
1110         size_t box_len = 0;
1111         size_t i;
1112
1113         for (i = 0; i < length_code_size; i++) {
1114             box_len <<= 8;
1115             box_len |= data[i];
1116         }
1117
1118         if (remaining_size < box_len + sizeof(start_code)) {
1119             av_log(NULL, AV_LOG_ERROR, "Length is out of range.\n");
1120             AVERROR_BUFFER_TOO_SMALL;
1121         }
1122
1123         memcpy(data, start_code, sizeof(start_code));
1124         data += box_len + sizeof(start_code);
1125         remaining_size -= box_len + sizeof(start_code);
1126     }
1127
1128     return 0;
1129 }
1130
1131 /**
1132  * Copies NAL units and replaces length codes with
1133  * H.264 Annex B start codes. On failure, the contents of
1134  * dst_data may have been modified.
1135  *
1136  * @param length_code_size Byte length of each length code
1137  * @param src_data NAL units prefixed with length codes.
1138  * @param src_size Length of buffer, excluding any padding.
1139  * @param dst_data Must be zeroed before calling this function.
1140  *                 Contains the copied NAL units prefixed with
1141  *                 start codes when the function returns
1142  *                 successfully.
1143  * @param dst_size Length of dst_data
1144  * @return 0 on success
1145  *         AVERROR_INVALIDDATA if length_code_size is invalid
1146  *         AVERROR_BUFFER_TOO_SMALL if dst_data is too small
1147  *         or if a length_code in src_data specifies data beyond
1148  *         the end of its buffer.
1149  */
1150 static int copy_replace_length_codes(
1151     size_t        length_code_size,
1152     const uint8_t *src_data,
1153     size_t        src_size,
1154     uint8_t       *dst_data,
1155     size_t        dst_size)
1156 {
1157     size_t remaining_src_size = src_size;
1158     size_t remaining_dst_size = dst_size;
1159
1160     if (length_code_size > 4) {
1161         return AVERROR_INVALIDDATA;
1162     }
1163
1164     while (remaining_src_size > 0) {
1165         size_t curr_src_len;
1166         size_t curr_dst_len;
1167         size_t box_len = 0;
1168         size_t i;
1169
1170         uint8_t       *dst_box;
1171         const uint8_t *src_box;
1172
1173         for (i = 0; i < length_code_size; i++) {
1174             box_len <<= 8;
1175             box_len |= src_data[i];
1176         }
1177
1178         curr_src_len = box_len + length_code_size;
1179         curr_dst_len = box_len + sizeof(start_code);
1180
1181         if (remaining_src_size < curr_src_len) {
1182             return AVERROR_BUFFER_TOO_SMALL;
1183         }
1184
1185         if (remaining_dst_size < curr_dst_len) {
1186             return AVERROR_BUFFER_TOO_SMALL;
1187         }
1188
1189         dst_box = dst_data + sizeof(start_code);
1190         src_box = src_data + length_code_size;
1191
1192         memcpy(dst_data, start_code, sizeof(start_code));
1193         memcpy(dst_box,  src_box,    box_len);
1194
1195         src_data += curr_src_len;
1196         dst_data += curr_dst_len;
1197
1198         remaining_src_size -= curr_src_len;
1199         remaining_dst_size -= curr_dst_len;
1200     }
1201
1202     return 0;
1203 }
1204
1205 static int vtenc_cm_to_avpacket(
1206     AVCodecContext    *avctx,
1207     CMSampleBufferRef sample_buffer,
1208     AVPacket          *pkt)
1209 {
1210     VTEncContext *vtctx = avctx->priv_data;
1211
1212     int     status;
1213     bool    is_key_frame;
1214     bool    add_header;
1215     char    *buf_data;
1216     size_t  length_code_size;
1217     size_t  header_size = 0;
1218     size_t  in_buf_size;
1219     int64_t dts_delta;
1220     int64_t time_base_num;
1221     CMTime  pts;
1222     CMTime  dts;
1223
1224     CMBlockBufferRef            block;
1225     CMVideoFormatDescriptionRef vid_fmt;
1226
1227
1228     vtenc_get_frame_info(sample_buffer, &is_key_frame);
1229     status = get_length_code_size(avctx, sample_buffer, &length_code_size);
1230     if (status) return status;
1231
1232     add_header = is_key_frame && !(avctx->flags & AV_CODEC_FLAG_GLOBAL_HEADER);
1233
1234     if (add_header) {
1235         vid_fmt = CMSampleBufferGetFormatDescription(sample_buffer);
1236         if (!vid_fmt) {
1237             av_log(avctx, AV_LOG_ERROR, "Cannot get format description.\n");
1238         }
1239
1240         int status = get_params_size(avctx, vid_fmt, &header_size);
1241         if (status) return status;
1242     }
1243
1244     block = CMSampleBufferGetDataBuffer(sample_buffer);
1245     if (!block) {
1246         av_log(avctx, AV_LOG_ERROR, "Could not get block buffer from sample buffer.\n");
1247         return AVERROR_EXTERNAL;
1248     }
1249
1250
1251     status = CMBlockBufferGetDataPointer(block, 0, &in_buf_size, NULL, &buf_data);
1252     if (status) {
1253         av_log(avctx, AV_LOG_ERROR, "Error: cannot get data pointer: %d\n", status);
1254         return AVERROR_EXTERNAL;
1255     }
1256
1257     size_t out_buf_size = header_size + in_buf_size;
1258     bool can_reuse_cmbuffer = !add_header &&
1259                               !pkt->data  &&
1260                               length_code_size == sizeof(start_code);
1261
1262     av_init_packet(pkt);
1263
1264     if (can_reuse_cmbuffer) {
1265         AVBufferRef* buf_ref = av_buffer_create(
1266             buf_data,
1267             out_buf_size,
1268             vtenc_free_block,
1269             block,
1270             0
1271         );
1272
1273         if (!buf_ref) return AVERROR(ENOMEM);
1274
1275         CFRetain(block);
1276
1277         pkt->buf  = buf_ref;
1278         pkt->data = buf_data;
1279         pkt->size = in_buf_size;
1280
1281         status = replace_length_codes(length_code_size, pkt->data, pkt->size);
1282         if (status) {
1283             av_log(avctx, AV_LOG_ERROR, "Error replacing length codes: %d\n", status);
1284             return status;
1285         }
1286     } else {
1287         if (!pkt->data) {
1288             status = av_new_packet(pkt, out_buf_size);
1289             if(status) return status;
1290         }
1291
1292         if (pkt->size < out_buf_size) {
1293             av_log(avctx, AV_LOG_ERROR, "Error: packet's buffer is too small.\n");
1294             return AVERROR_BUFFER_TOO_SMALL;
1295         }
1296
1297         if (add_header) {
1298             status = copy_param_sets(avctx, vid_fmt, pkt->data, out_buf_size);
1299             if(status) return status;
1300         }
1301
1302         status = copy_replace_length_codes(
1303             length_code_size,
1304             buf_data,
1305             in_buf_size,
1306             pkt->data + header_size,
1307             pkt->size - header_size
1308         );
1309
1310         if (status) {
1311             av_log(avctx, AV_LOG_ERROR, "Error copying packet data: %d", status);
1312             return status;
1313         }
1314     }
1315
1316     if (is_key_frame) {
1317         pkt->flags |= AV_PKT_FLAG_KEY;
1318     }
1319
1320     pts = CMSampleBufferGetPresentationTimeStamp(sample_buffer);
1321     dts = CMSampleBufferGetDecodeTimeStamp      (sample_buffer);
1322
1323     if (CMTIME_IS_INVALID(dts)) {
1324         if (!vtctx->has_b_frames) {
1325             dts = pts;
1326         } else {
1327             av_log(avctx, AV_LOG_ERROR, "DTS is invalid.\n");
1328             return AVERROR_EXTERNAL;
1329         }
1330     }
1331
1332     dts_delta = vtctx->dts_delta >= 0 ? vtctx->dts_delta : 0;
1333     time_base_num = avctx->time_base.num;
1334     pkt->pts = pts.value / time_base_num;
1335     pkt->dts = dts.value / time_base_num - dts_delta;
1336
1337     return 0;
1338 }
1339
1340 /*
1341  * contiguous_buf_size is 0 if not contiguous, and the size of the buffer
1342  * containing all planes if so.
1343  */
1344 static int get_cv_pixel_info(
1345     AVCodecContext *avctx,
1346     const AVFrame  *frame,
1347     int            *color,
1348     int            *plane_count,
1349     size_t         *widths,
1350     size_t         *heights,
1351     size_t         *strides,
1352     size_t         *contiguous_buf_size)
1353 {
1354     VTEncContext *vtctx = avctx->priv_data;
1355     int av_format       = frame->format;
1356     int av_color_range  = av_frame_get_color_range(frame);
1357     int i;
1358     int range_guessed;
1359     int status;
1360
1361     status = get_cv_pixel_format(avctx, av_format, av_color_range, color, &range_guessed);
1362     if (status) {
1363         av_log(avctx,
1364             AV_LOG_ERROR,
1365             "Could not get pixel format for color format '%s' range '%s'.\n",
1366             av_get_pix_fmt_name(av_format),
1367             av_color_range > AVCOL_RANGE_UNSPECIFIED &&
1368             av_color_range < AVCOL_RANGE_NB ?
1369                av_color_range_name(av_color_range) :
1370                "Unknown");
1371
1372         return AVERROR(EINVAL);
1373     }
1374
1375     if (range_guessed) {
1376         if (!vtctx->warned_color_range) {
1377             vtctx->warned_color_range = true;
1378             av_log(avctx,
1379                    AV_LOG_WARNING,
1380                    "Color range not set for %s. Using MPEG range.\n",
1381                    av_get_pix_fmt_name(av_format));
1382         }
1383
1384         av_log(avctx, AV_LOG_WARNING, "");
1385     }
1386
1387     switch (av_format) {
1388     case AV_PIX_FMT_NV12:
1389         *plane_count = 2;
1390
1391         widths [0] = avctx->width;
1392         heights[0] = avctx->height;
1393         strides[0] = frame ? frame->linesize[0] : avctx->width;
1394
1395         widths [1] = (avctx->width  + 1) / 2;
1396         heights[1] = (avctx->height + 1) / 2;
1397         strides[1] = frame ? frame->linesize[1] : (avctx->width + 1) & -2;
1398         break;
1399
1400     case AV_PIX_FMT_YUV420P:
1401         *plane_count = 3;
1402
1403         widths [0] = avctx->width;
1404         heights[0] = avctx->height;
1405         strides[0] = frame ? frame->linesize[0] : avctx->width;
1406
1407         widths [1] = (avctx->width  + 1) / 2;
1408         heights[1] = (avctx->height + 1) / 2;
1409         strides[1] = frame ? frame->linesize[1] : (avctx->width + 1) / 2;
1410
1411         widths [2] = (avctx->width  + 1) / 2;
1412         heights[2] = (avctx->height + 1) / 2;
1413         strides[2] = frame ? frame->linesize[2] : (avctx->width + 1) / 2;
1414         break;
1415
1416     default:
1417         av_log(
1418                avctx,
1419                AV_LOG_ERROR,
1420                "Could not get frame format info for color %d range %d.\n",
1421                av_format,
1422                av_color_range);
1423
1424         return AVERROR(EINVAL);
1425     }
1426
1427     *contiguous_buf_size = 0;
1428     for (i = 0; i < *plane_count; i++) {
1429         if (i < *plane_count - 1 &&
1430             frame->data[i] + strides[i] * heights[i] != frame->data[i + 1]) {
1431             *contiguous_buf_size = 0;
1432             break;
1433         }
1434
1435         *contiguous_buf_size += strides[i] * heights[i];
1436     }
1437
1438     return 0;
1439 }
1440
1441 #if !TARGET_OS_IPHONE
1442 //Not used on iOS - frame is always copied.
1443 static void free_avframe(
1444     void       *release_ctx,
1445     const void *data,
1446     size_t      size,
1447     size_t      plane_count,
1448     const void *plane_addresses[])
1449 {
1450     AVFrame *frame = release_ctx;
1451     av_frame_free(&frame);
1452 }
1453 #else
1454 //Not used on OSX - frame is never copied.
1455 static int copy_avframe_to_pixel_buffer(AVCodecContext   *avctx,
1456                                         const AVFrame    *frame,
1457                                         CVPixelBufferRef cv_img,
1458                                         const size_t     *plane_strides,
1459                                         const size_t     *plane_rows)
1460 {
1461     int i, j;
1462     size_t plane_count;
1463     int status;
1464     int rows;
1465     int src_stride;
1466     int dst_stride;
1467     uint8_t *src_addr;
1468     uint8_t *dst_addr;
1469     size_t copy_bytes;
1470
1471     status = CVPixelBufferLockBaseAddress(cv_img, 0);
1472     if (status) {
1473         av_log(
1474             avctx,
1475             AV_LOG_ERROR,
1476             "Error: Could not lock base address of CVPixelBuffer: %d.\n",
1477             status
1478         );
1479     }
1480
1481     if (CVPixelBufferIsPlanar(cv_img)) {
1482         plane_count = CVPixelBufferGetPlaneCount(cv_img);
1483         for (i = 0; frame->data[i]; i++) {
1484             if (i == plane_count) {
1485                 CVPixelBufferUnlockBaseAddress(cv_img, 0);
1486                 av_log(avctx,
1487                     AV_LOG_ERROR,
1488                     "Error: different number of planes in AVFrame and CVPixelBuffer.\n"
1489                 );
1490
1491                 return AVERROR_EXTERNAL;
1492             }
1493
1494             dst_addr = (uint8_t*)CVPixelBufferGetBaseAddressOfPlane(cv_img, i);
1495             src_addr = (uint8_t*)frame->data[i];
1496             dst_stride = CVPixelBufferGetBytesPerRowOfPlane(cv_img, i);
1497             src_stride = plane_strides[i];
1498             rows = plane_rows[i];
1499
1500             if (dst_stride == src_stride) {
1501                 memcpy(dst_addr, src_addr, src_stride * rows);
1502             } else {
1503                 copy_bytes = dst_stride < src_stride ? dst_stride : src_stride;
1504
1505                 for (j = 0; j < rows; j++) {
1506                     memcpy(dst_addr + j * dst_stride, src_addr + j * src_stride, copy_bytes);
1507                 }
1508             }
1509         }
1510     } else {
1511         if (frame->data[1]) {
1512             CVPixelBufferUnlockBaseAddress(cv_img, 0);
1513             av_log(avctx,
1514                 AV_LOG_ERROR,
1515                 "Error: different number of planes in AVFrame and non-planar CVPixelBuffer.\n"
1516             );
1517
1518             return AVERROR_EXTERNAL;
1519         }
1520
1521         dst_addr = (uint8_t*)CVPixelBufferGetBaseAddress(cv_img);
1522         src_addr = (uint8_t*)frame->data[0];
1523         dst_stride = CVPixelBufferGetBytesPerRow(cv_img);
1524         src_stride = plane_strides[0];
1525         rows = plane_rows[0];
1526
1527         if (dst_stride == src_stride) {
1528             memcpy(dst_addr, src_addr, src_stride * rows);
1529         } else {
1530             copy_bytes = dst_stride < src_stride ? dst_stride : src_stride;
1531
1532             for (j = 0; j < rows; j++) {
1533                 memcpy(dst_addr + j * dst_stride, src_addr + j * src_stride, copy_bytes);
1534             }
1535         }
1536     }
1537
1538     status = CVPixelBufferUnlockBaseAddress(cv_img, 0);
1539     if (status) {
1540         av_log(avctx, AV_LOG_ERROR, "Error: Could not unlock CVPixelBuffer base address: %d.\n", status);
1541         return AVERROR_EXTERNAL;
1542     }
1543
1544     return 0;
1545 }
1546 #endif //!TARGET_OS_IPHONE
1547
1548 static int create_cv_pixel_buffer(AVCodecContext   *avctx,
1549                                   const AVFrame    *frame,
1550                                   CVPixelBufferRef *cv_img)
1551 {
1552     int plane_count;
1553     int color;
1554     size_t widths [AV_NUM_DATA_POINTERS];
1555     size_t heights[AV_NUM_DATA_POINTERS];
1556     size_t strides[AV_NUM_DATA_POINTERS];
1557     int status;
1558     size_t contiguous_buf_size;
1559 #if TARGET_OS_IPHONE
1560     CVPixelBufferPoolRef pix_buf_pool;
1561     VTEncContext* vtctx = avctx->priv_data;
1562 #else
1563     CFMutableDictionaryRef pix_buf_attachments = CFDictionaryCreateMutable(
1564                                                    kCFAllocatorDefault,
1565                                                    10,
1566                                                    &kCFCopyStringDictionaryKeyCallBacks,
1567                                                    &kCFTypeDictionaryValueCallBacks);
1568
1569     if (!pix_buf_attachments) return AVERROR(ENOMEM);
1570 #endif
1571
1572     if (avctx->pix_fmt == AV_PIX_FMT_VIDEOTOOLBOX) {
1573         av_assert0(frame->format == AV_PIX_FMT_VIDEOTOOLBOX);
1574
1575         *cv_img = (CVPixelBufferRef)frame->data[3];
1576         av_assert0(*cv_img);
1577
1578         CFRetain(*cv_img);
1579         return 0;
1580     }
1581
1582     memset(widths,  0, sizeof(widths));
1583     memset(heights, 0, sizeof(heights));
1584     memset(strides, 0, sizeof(strides));
1585
1586     status = get_cv_pixel_info(
1587         avctx,
1588         frame,
1589         &color,
1590         &plane_count,
1591         widths,
1592         heights,
1593         strides,
1594         &contiguous_buf_size
1595     );
1596
1597     if (status) {
1598         av_log(
1599             avctx,
1600             AV_LOG_ERROR,
1601             "Error: Cannot convert format %d color_range %d: %d\n",
1602             frame->format,
1603             av_frame_get_color_range(frame),
1604             status
1605         );
1606
1607         return AVERROR_EXTERNAL;
1608     }
1609
1610 #if TARGET_OS_IPHONE
1611     pix_buf_pool = VTCompressionSessionGetPixelBufferPool(vtctx->session);
1612     if (!pix_buf_pool) {
1613         av_log(avctx, AV_LOG_ERROR, "Could not get pixel buffer pool.\n");
1614         return AVERROR_EXTERNAL;
1615     }
1616
1617     status = CVPixelBufferPoolCreatePixelBuffer(NULL,
1618                                                 pix_buf_pool,
1619                                                 cv_img);
1620
1621
1622     if (status) {
1623         av_log(avctx, AV_LOG_ERROR, "Could not create pixel buffer from pool: %d.\n", status);
1624         return AVERROR_EXTERNAL;
1625     }
1626
1627     status = copy_avframe_to_pixel_buffer(avctx, frame, *cv_img, strides, heights);
1628     if (status) {
1629         CFRelease(*cv_img);
1630         *cv_img = NULL;
1631         return status;
1632     }
1633 #else
1634     AVFrame *enc_frame = av_frame_alloc();
1635     if (!enc_frame) return AVERROR(ENOMEM);
1636
1637     status = av_frame_ref(enc_frame, frame);
1638     if (status) {
1639         av_frame_free(&enc_frame);
1640         return status;
1641     }
1642
1643     status = CVPixelBufferCreateWithPlanarBytes(
1644         kCFAllocatorDefault,
1645         enc_frame->width,
1646         enc_frame->height,
1647         color,
1648         NULL,
1649         contiguous_buf_size,
1650         plane_count,
1651         (void **)enc_frame->data,
1652         widths,
1653         heights,
1654         strides,
1655         free_avframe,
1656         enc_frame,
1657         NULL,
1658         cv_img
1659     );
1660
1661     add_color_attr(avctx, pix_buf_attachments);
1662     CVBufferSetAttachments(*cv_img, pix_buf_attachments, kCVAttachmentMode_ShouldPropagate);
1663     CFRelease(pix_buf_attachments);
1664
1665     if (status) {
1666         av_log(avctx, AV_LOG_ERROR, "Error: Could not create CVPixelBuffer: %d\n", status);
1667         return AVERROR_EXTERNAL;
1668     }
1669 #endif
1670
1671     return 0;
1672 }
1673
1674 static int create_encoder_dict_h264(const AVFrame *frame,
1675                                     CFDictionaryRef* dict_out)
1676 {
1677     CFDictionaryRef dict = NULL;
1678     if (frame->pict_type == AV_PICTURE_TYPE_I) {
1679         const void *keys[] = { kVTEncodeFrameOptionKey_ForceKeyFrame };
1680         const void *vals[] = { kCFBooleanTrue };
1681
1682         dict = CFDictionaryCreate(NULL, keys, vals, 1, NULL, NULL);
1683         if(!dict) return AVERROR(ENOMEM);
1684     }
1685
1686     *dict_out = dict;
1687     return 0;
1688 }
1689
1690 static int vtenc_send_frame(AVCodecContext *avctx,
1691                             VTEncContext   *vtctx,
1692                             const AVFrame  *frame)
1693 {
1694     CMTime time;
1695     CFDictionaryRef frame_dict;
1696     CVPixelBufferRef cv_img = NULL;
1697     int status = create_cv_pixel_buffer(avctx, frame, &cv_img);
1698
1699     if (status) return status;
1700
1701     status = create_encoder_dict_h264(frame, &frame_dict);
1702     if (status) {
1703         CFRelease(cv_img);
1704         return status;
1705     }
1706
1707     time = CMTimeMake(frame->pts * avctx->time_base.num, avctx->time_base.den);
1708     status = VTCompressionSessionEncodeFrame(
1709         vtctx->session,
1710         cv_img,
1711         time,
1712         kCMTimeInvalid,
1713         frame_dict,
1714         NULL,
1715         NULL
1716     );
1717
1718     if (frame_dict) CFRelease(frame_dict);
1719     CFRelease(cv_img);
1720
1721     if (status) {
1722         av_log(avctx, AV_LOG_ERROR, "Error: cannot encode frame: %d\n", status);
1723         return AVERROR_EXTERNAL;
1724     }
1725
1726     return 0;
1727 }
1728
1729 static av_cold int vtenc_frame(
1730     AVCodecContext *avctx,
1731     AVPacket       *pkt,
1732     const AVFrame  *frame,
1733     int            *got_packet)
1734 {
1735     VTEncContext *vtctx = avctx->priv_data;
1736     bool get_frame;
1737     int status;
1738     CMSampleBufferRef buf = NULL;
1739
1740     if (frame) {
1741         status = vtenc_send_frame(avctx, vtctx, frame);
1742
1743         if (status) {
1744             status = AVERROR_EXTERNAL;
1745             goto end_nopkt;
1746         }
1747
1748         if (vtctx->frame_ct_in == 0) {
1749             vtctx->first_pts = frame->pts;
1750         } else if(vtctx->frame_ct_in == 1 && vtctx->has_b_frames) {
1751             vtctx->dts_delta = frame->pts - vtctx->first_pts;
1752         }
1753
1754         vtctx->frame_ct_in++;
1755     } else if(!vtctx->flushing) {
1756         vtctx->flushing = true;
1757
1758         status = VTCompressionSessionCompleteFrames(vtctx->session,
1759                                                     kCMTimeIndefinite);
1760
1761         if (status) {
1762             av_log(avctx, AV_LOG_ERROR, "Error flushing frames: %d\n", status);
1763             status = AVERROR_EXTERNAL;
1764             goto end_nopkt;
1765         }
1766     }
1767
1768     *got_packet = 0;
1769     get_frame = vtctx->dts_delta >= 0 || !frame;
1770     if (!get_frame) {
1771         status = 0;
1772         goto end_nopkt;
1773     }
1774
1775     status = vtenc_q_pop(vtctx, !frame, &buf);
1776     if (status) goto end_nopkt;
1777     if (!buf)   goto end_nopkt;
1778
1779     status = vtenc_cm_to_avpacket(avctx, buf, pkt);
1780     CFRelease(buf);
1781     if (status) goto end_nopkt;
1782
1783     *got_packet = 1;
1784     return 0;
1785
1786 end_nopkt:
1787     av_packet_unref(pkt);
1788     return status;
1789 }
1790
1791 static av_cold int vtenc_close(AVCodecContext *avctx)
1792 {
1793     VTEncContext *vtctx = avctx->priv_data;
1794
1795     if(!vtctx->session) return 0;
1796
1797     pthread_cond_destroy(&vtctx->cv_sample_sent);
1798     pthread_mutex_destroy(&vtctx->lock);
1799     CFRelease(vtctx->session);
1800     vtctx->session = NULL;
1801
1802     if (vtctx->color_primaries) {
1803         CFRelease(vtctx->color_primaries);
1804         vtctx->color_primaries = NULL;
1805     }
1806
1807     if (vtctx->transfer_function) {
1808         CFRelease(vtctx->transfer_function);
1809         vtctx->transfer_function = NULL;
1810     }
1811
1812     if (vtctx->ycbcr_matrix) {
1813         CFRelease(vtctx->ycbcr_matrix);
1814         vtctx->ycbcr_matrix = NULL;
1815     }
1816
1817     return 0;
1818 }
1819
1820 static const enum AVPixelFormat pix_fmts[] = {
1821     AV_PIX_FMT_VIDEOTOOLBOX,
1822     AV_PIX_FMT_NV12,
1823     AV_PIX_FMT_YUV420P,
1824     AV_PIX_FMT_NONE
1825 };
1826
1827 #define OFFSET(x) offsetof(VTEncContext, x)
1828 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
1829 static const AVOption options[] = {
1830     { "profile", "Profile", OFFSET(profile), AV_OPT_TYPE_INT, { .i64 = H264_PROF_AUTO }, H264_PROF_AUTO, H264_PROF_COUNT, VE, "profile" },
1831     { "baseline", "Baseline Profile", 0, AV_OPT_TYPE_CONST, { .i64 = H264_PROF_BASELINE }, INT_MIN, INT_MAX, VE, "profile" },
1832     { "main",     "Main Profile",     0, AV_OPT_TYPE_CONST, { .i64 = H264_PROF_MAIN     }, INT_MIN, INT_MAX, VE, "profile" },
1833     { "high",     "High Profile",     0, AV_OPT_TYPE_CONST, { .i64 = H264_PROF_HIGH     }, INT_MIN, INT_MAX, VE, "profile" },
1834
1835     { "level", "Level", OFFSET(level), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 52, VE, "level" },
1836     { "1.3", "Level 1.3, only available with Baseline Profile", 0, AV_OPT_TYPE_CONST, { .i64 = 13 }, INT_MIN, INT_MAX, VE, "level" },
1837     { "3.0", "Level 3.0", 0, AV_OPT_TYPE_CONST, { .i64 = 30 }, INT_MIN, INT_MAX, VE, "level" },
1838     { "3.1", "Level 3.1", 0, AV_OPT_TYPE_CONST, { .i64 = 31 }, INT_MIN, INT_MAX, VE, "level" },
1839     { "3.2", "Level 3.2", 0, AV_OPT_TYPE_CONST, { .i64 = 32 }, INT_MIN, INT_MAX, VE, "level" },
1840     { "4.0", "Level 4.0", 0, AV_OPT_TYPE_CONST, { .i64 = 40 }, INT_MIN, INT_MAX, VE, "level" },
1841     { "4.1", "Level 4.1", 0, AV_OPT_TYPE_CONST, { .i64 = 41 }, INT_MIN, INT_MAX, VE, "level" },
1842     { "4.2", "Level 4.2", 0, AV_OPT_TYPE_CONST, { .i64 = 42 }, INT_MIN, INT_MAX, VE, "level" },
1843     { "5.0", "Level 5.0", 0, AV_OPT_TYPE_CONST, { .i64 = 50 }, INT_MIN, INT_MAX, VE, "level" },
1844     { "5.1", "Level 5.1", 0, AV_OPT_TYPE_CONST, { .i64 = 51 }, INT_MIN, INT_MAX, VE, "level" },
1845     { "5.2", "Level 5.2", 0, AV_OPT_TYPE_CONST, { .i64 = 52 }, INT_MIN, INT_MAX, VE, "level" },
1846
1847     { "allow_sw", "Allow software encoding", OFFSET(allow_sw), AV_OPT_TYPE_BOOL,
1848         { .i64 = 0 }, 0, 1, VE },
1849
1850     { "coder", "Entropy coding", OFFSET(entropy), AV_OPT_TYPE_INT, { .i64 = VT_ENTROPY_NOT_SET }, VT_ENTROPY_NOT_SET, VT_CABAC, VE, "coder" },
1851     { "cavlc", "CAVLC entropy coding", 0, AV_OPT_TYPE_CONST, { .i64 = VT_CAVLC }, INT_MIN, INT_MAX, VE, "coder" },
1852     { "vlc",   "CAVLC entropy coding", 0, AV_OPT_TYPE_CONST, { .i64 = VT_CAVLC }, INT_MIN, INT_MAX, VE, "coder" },
1853     { "cabac", "CABAC entropy coding", 0, AV_OPT_TYPE_CONST, { .i64 = VT_CABAC }, INT_MIN, INT_MAX, VE, "coder" },
1854     { "ac",    "CABAC entropy coding", 0, AV_OPT_TYPE_CONST, { .i64 = VT_CABAC }, INT_MIN, INT_MAX, VE, "coder" },
1855
1856     { "realtime", "Hint that encoding should happen in real-time if not faster (e.g. capturing from camera).",
1857         OFFSET(realtime), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, VE },
1858
1859     { "frames_before", "Other frames will come before the frames in this session. This helps smooth concatenation issues.",
1860         OFFSET(frames_before), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, VE },
1861     { "frames_after", "Other frames will come after the frames in this session. This helps smooth concatenation issues.",
1862         OFFSET(frames_after), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, VE },
1863
1864     { NULL },
1865 };
1866
1867 static const AVClass h264_videotoolbox_class = {
1868     .class_name = "h264_videotoolbox",
1869     .item_name  = av_default_item_name,
1870     .option     = options,
1871     .version    = LIBAVUTIL_VERSION_INT,
1872 };
1873
1874 AVCodec ff_h264_videotoolbox_encoder = {
1875     .name             = "h264_videotoolbox",
1876     .long_name        = NULL_IF_CONFIG_SMALL("VideoToolbox H.264 Encoder"),
1877     .type             = AVMEDIA_TYPE_VIDEO,
1878     .id               = AV_CODEC_ID_H264,
1879     .priv_data_size   = sizeof(VTEncContext),
1880     .pix_fmts         = pix_fmts,
1881     .init             = vtenc_init,
1882     .encode2          = vtenc_frame,
1883     .close            = vtenc_close,
1884     .capabilities     = AV_CODEC_CAP_DELAY,
1885     .priv_class       = &h264_videotoolbox_class,
1886     .caps_internal    = FF_CODEC_CAP_INIT_THREADSAFE |
1887                         FF_CODEC_CAP_INIT_CLEANUP,
1888 };