avcodec/videotoolboxenc: add hevc_videotoolbox encoder

Signed-off-by: Aman Gupta <aman@tmm1.net>
Reviewed-by: Rodger Combs <rodger.combs@gmail.com>

1 files changed, 143 insertions, 25 deletions

diff --git a/libavcodec/videotoolboxenc.c b/libavcodec/videotoolboxenc.c
index eba6cc672f..7b50edbb8f 100644
--- a/libavcodec/videotoolboxenc.c
+++ b/libavcodec/videotoolboxenc.c
@@ -35,6 +35,17 @@
 #include "h264_sei.h"
 #include <dlfcn.h>
 
+#if !HAVE_KCMVIDEOCODECTYPE_HEVC
+enum { kCMVideoCodecType_HEVC = 'hvc1' };
+#endif
+
+typedef OSStatus (*getParameterSetAtIndex)(CMFormatDescriptionRef videoDesc,
+                                           size_t parameterSetIndex,
+                                           const uint8_t * _Nullable *parameterSetPointerOut,
+                                           size_t *parameterSetSizeOut,
+                                           size_t *parameterSetCountOut,
+                                           int *NALUnitHeaderLengthOut);
+
 //These symbols may not be present
 static struct{
     CFStringRef kCVImageBufferColorPrimaries_ITU_R_2020;
@@ -65,10 +76,15 @@ static struct{
     CFStringRef kVTProfileLevel_H264_High_5_2;
     CFStringRef kVTProfileLevel_H264_High_AutoLevel;
 
+    CFStringRef kVTProfileLevel_HEVC_Main_AutoLevel;
+    CFStringRef kVTProfileLevel_HEVC_Main10_AutoLevel;
+
     CFStringRef kVTCompressionPropertyKey_RealTime;
 
     CFStringRef kVTVideoEncoderSpecification_EnableHardwareAcceleratedVideoEncoder;
     CFStringRef kVTVideoEncoderSpecification_RequireHardwareAcceleratedVideoEncoder;
+
+    getParameterSetAtIndex CMVideoFormatDescriptionGetHEVCParameterSetAtIndex;
 } compat_keys;
 
 #define GET_SYM(symbol, defaultVal)                                     \
@@ -83,6 +99,12 @@ do{                                                                     \
 static pthread_once_t once_ctrl = PTHREAD_ONCE_INIT;
 
 static void loadVTEncSymbols(){
+    compat_keys.CMVideoFormatDescriptionGetHEVCParameterSetAtIndex =
+        (getParameterSetAtIndex)dlsym(
+            RTLD_DEFAULT,
+            "CMVideoFormatDescriptionGetHEVCParameterSetAtIndex"
+        );
+
     GET_SYM(kCVImageBufferColorPrimaries_ITU_R_2020,   "ITU_R_2020");
     GET_SYM(kCVImageBufferTransferFunction_ITU_R_2020, "ITU_R_2020");
     GET_SYM(kCVImageBufferYCbCrMatrix_ITU_R_2020,      "ITU_R_2020");
@@ -111,6 +133,9 @@ static void loadVTEncSymbols(){
     GET_SYM(kVTProfileLevel_H264_High_5_2,           "H264_High_5_2");
     GET_SYM(kVTProfileLevel_H264_High_AutoLevel,     "H264_High_AutoLevel");
 
+    GET_SYM(kVTProfileLevel_HEVC_Main_AutoLevel,     "HEVC_Main_AutoLevel");
+    GET_SYM(kVTProfileLevel_HEVC_Main10_AutoLevel,   "HEVC_Main10_AutoLevel");
+
     GET_SYM(kVTCompressionPropertyKey_RealTime, "RealTime");
 
     GET_SYM(kVTVideoEncoderSpecification_EnableHardwareAcceleratedVideoEncoder,
@@ -133,6 +158,13 @@ typedef enum VTH264Entropy{
     VT_CABAC
 } VTH264Entropy;
 
+typedef enum VT_HEVCProfile {
+    HEVC_PROF_AUTO,
+    HEVC_PROF_MAIN,
+    HEVC_PROF_MAIN10,
+    HEVC_PROF_COUNT
+} VT_HEVCProfile;
+
 static const uint8_t start_code[] = { 0, 0, 0, 1 };
 
 typedef struct ExtraSEI {
@@ -149,10 +181,12 @@ typedef struct BufNode {
 
 typedef struct VTEncContext {
     AVClass *class;
+    enum AVCodecID codec_id;
     VTCompressionSessionRef session;
     CFStringRef ycbcr_matrix;
     CFStringRef color_primaries;
     CFStringRef transfer_function;
+    getParameterSetAtIndex get_param_set_func;
 
     pthread_mutex_t lock;
     pthread_cond_t  cv_sample_sent;
@@ -348,6 +382,7 @@ static CMVideoCodecType get_cm_codec_type(enum AVCodecID id)
 {
     switch (id) {
     case AV_CODEC_ID_H264: return kCMVideoCodecType_H264;
+    case AV_CODEC_ID_HEVC: return kCMVideoCodecType_HEVC;
     default:               return 0;
     }
 }
@@ -365,12 +400,13 @@ static int get_params_size(
     CMVideoFormatDescriptionRef vid_fmt,
     size_t                      *size)
 {
+    VTEncContext *vtctx = avctx->priv_data;
     size_t total_size = 0;
     size_t ps_count;
     int is_count_bad = 0;
     size_t i;
     int status;
-    status = CMVideoFormatDescriptionGetH264ParameterSetAtIndex(vid_fmt,
+    status = vtctx->get_param_set_func(vid_fmt,
                                                                 0,
                                                                 NULL,
                                                                 NULL,
@@ -385,7 +421,7 @@ static int get_params_size(
     for (i = 0; i < ps_count || is_count_bad; i++) {
         const uint8_t *ps;
         size_t ps_size;
-        status = CMVideoFormatDescriptionGetH264ParameterSetAtIndex(vid_fmt,
+        status = vtctx->get_param_set_func(vid_fmt,
                                                                     i,
                                                                     &ps,
                                                                     &ps_size,
@@ -419,13 +455,14 @@ static int copy_param_sets(
     uint8_t                     *dst,
     size_t                      dst_size)
 {
+    VTEncContext *vtctx = avctx->priv_data;
     size_t ps_count;
     int is_count_bad = 0;
     int status;
     size_t offset = 0;
     size_t i;
 
-    status = CMVideoFormatDescriptionGetH264ParameterSetAtIndex(vid_fmt,
+    status = vtctx->get_param_set_func(vid_fmt,
                                                                 0,
                                                                 NULL,
                                                                 NULL,
@@ -443,7 +480,7 @@ static int copy_param_sets(
         size_t ps_size;
         size_t next_offset;
 
-        status = CMVideoFormatDescriptionGetH264ParameterSetAtIndex(vid_fmt,
+        status = vtctx->get_param_set_func(vid_fmt,
                                                                     i,
                                                                     &ps,
                                                                     &ps_size,
@@ -548,6 +585,7 @@ static int get_length_code_size(
     CMSampleBufferRef sample_buffer,
     size_t            *size)
 {
+    VTEncContext *vtctx = avctx->priv_data;
     CMVideoFormatDescriptionRef vid_fmt;
     int isize;
     int status;
@@ -558,7 +596,7 @@ static int get_length_code_size(
         return AVERROR_EXTERNAL;
     }
 
-    status = CMVideoFormatDescriptionGetH264ParameterSetAtIndex(vid_fmt,
+    status = vtctx->get_param_set_func(vid_fmt,
                                                                 0,
                                                                 NULL,
                                                                 NULL,
@@ -579,8 +617,8 @@ static int get_length_code_size(
  * If profile_level_val is NULL and this method returns true, don't specify the
  * profile/level to the encoder.
  */
-static bool get_vt_profile_level(AVCodecContext *avctx,
-                                 CFStringRef    *profile_level_val)
+static bool get_vt_h264_profile_level(AVCodecContext *avctx,
+                                      CFStringRef    *profile_level_val)
 {
     VTEncContext *vtctx = avctx->priv_data;
     int64_t profile = vtctx->profile;
@@ -670,6 +708,41 @@ static bool get_vt_profile_level(AVCodecContext *avctx,
     return true;
 }
 
+/*
+ * Returns true on success.
+ *
+ * If profile_level_val is NULL and this method returns true, don't specify the
+ * profile/level to the encoder.
+ */
+static bool get_vt_hevc_profile_level(AVCodecContext *avctx,
+                                      CFStringRef    *profile_level_val)
+{
+    VTEncContext *vtctx = avctx->priv_data;
+    int64_t profile = vtctx->profile;
+
+    *profile_level_val = NULL;
+
+    switch (profile) {
+        case HEVC_PROF_AUTO:
+            return true;
+        case HEVC_PROF_MAIN:
+            *profile_level_val =
+                compat_keys.kVTProfileLevel_HEVC_Main_AutoLevel;
+            break;
+        case HEVC_PROF_MAIN10:
+            *profile_level_val =
+                compat_keys.kVTProfileLevel_HEVC_Main10_AutoLevel;
+            break;
+    }
+
+    if (!*profile_level_val) {
+        av_log(avctx, AV_LOG_ERROR, "Invalid Profile/Level.\n");
+        return false;
+    }
+
+    return true;
+}
+
 static int get_cv_pixel_format(AVCodecContext* avctx,
                                enum AVPixelFormat fmt,
                                enum AVColorRange range,
@@ -944,6 +1017,8 @@ static int vtenc_create_encoder(AVCodecContext   *avctx,
         return AVERROR_EXTERNAL;
     }
 
+    if (vtctx->codec_id == AV_CODEC_ID_H264) {
+    // kVTCompressionPropertyKey_DataRateLimits is not available for HEVC
     bytes_per_second_value = max_rate >> 3;
     bytes_per_second = CFNumberCreate(kCFAllocatorDefault,
                                       kCFNumberSInt64Type,
@@ -959,10 +1034,10 @@ static int vtenc_create_encoder(AVCodecContext   *avctx,
         CFRelease(bytes_per_second);
         return AVERROR(ENOMEM);
     }
-    nums[0] = bytes_per_second;
-    nums[1] = one_second;
+    nums[0] = (void *)bytes_per_second;
+    nums[1] = (void *)one_second;
     data_rate_limits = CFArrayCreate(kCFAllocatorDefault,
-                                     nums,
+                                     (const void **)nums,
                                      2,
                                      &kCFTypeArrayCallBacks);
 
@@ -992,6 +1067,7 @@ static int vtenc_create_encoder(AVCodecContext   *avctx,
             av_log(avctx, AV_LOG_ERROR, "Error setting profile/level property: %d\n", status);
         }
     }
+    }
 
     if (avctx->gop_size > 0) {
         CFNumberRef interval = CFNumberCreate(kCFAllocatorDefault,
@@ -1205,6 +1281,11 @@ static av_cold int vtenc_init(AVCodecContext *avctx)
         return AVERROR(EINVAL);
     }
 
+    vtctx->codec_id = avctx->codec_id;
+
+    if (vtctx->codec_id == AV_CODEC_ID_H264) {
+    vtctx->get_param_set_func = CMVideoFormatDescriptionGetH264ParameterSetAtIndex;
+
     vtctx->has_b_frames = avctx->max_b_frames > 0;
     if(vtctx->has_b_frames && vtctx->profile == H264_PROF_BASELINE){
         av_log(avctx, AV_LOG_WARNING, "Cannot use B-frames with baseline profile. Output will not contain B-frames.\n");
@@ -1216,7 +1297,12 @@ static av_cold int vtenc_init(AVCodecContext *avctx)
         vtctx->entropy = VT_ENTROPY_NOT_SET;
     }
 
-    if (!get_vt_profile_level(avctx, &profile_level)) return AVERROR(EINVAL);
+    if (!get_vt_h264_profile_level(avctx, &profile_level)) return AVERROR(EINVAL);
+    } else {
+        vtctx->get_param_set_func = compat_keys.CMVideoFormatDescriptionGetHEVCParameterSetAtIndex;
+        if (!vtctx->get_param_set_func) return AVERROR(EINVAL);
+        if (!get_vt_hevc_profile_level(avctx, &profile_level)) return AVERROR(EINVAL);
+    }
 
     vtctx->session = NULL;
 
@@ -2424,9 +2510,19 @@ static const enum AVPixelFormat pix_fmts[] = {
     AV_PIX_FMT_NONE
 };
 
-#define OFFSET(x) offsetof(VTEncContext, x)
 #define VE AV_OPT_FLAG_VIDEO_PARAM | AV_OPT_FLAG_ENCODING_PARAM
-static const AVOption options[] = {
+#define COMMON_OPTIONS \
+    { "allow_sw", "Allow software encoding", OFFSET(allow_sw), AV_OPT_TYPE_BOOL, \
+        { .i64 = 0 }, 0, 1, VE }, \
+    { "realtime", "Hint that encoding should happen in real-time if not faster (e.g. capturing from camera).", \
+        OFFSET(realtime), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, VE }, \
+    { "frames_before", "Other frames will come before the frames in this session. This helps smooth concatenation issues.", \
+        OFFSET(frames_before), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, VE }, \
+    { "frames_after", "Other frames will come after the frames in this session. This helps smooth concatenation issues.", \
+        OFFSET(frames_after), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, VE },
+
+#define OFFSET(x) offsetof(VTEncContext, x)
+static const AVOption h264_options[] = {
     { "profile", "Profile", OFFSET(profile), AV_OPT_TYPE_INT, { .i64 = H264_PROF_AUTO }, H264_PROF_AUTO, H264_PROF_COUNT, VE, "profile" },
     { "baseline", "Baseline Profile", 0, AV_OPT_TYPE_CONST, { .i64 = H264_PROF_BASELINE }, INT_MIN, INT_MAX, VE, "profile" },
     { "main",     "Main Profile",     0, AV_OPT_TYPE_CONST, { .i64 = H264_PROF_MAIN     }, INT_MIN, INT_MAX, VE, "profile" },
@@ -2444,32 +2540,22 @@ static const AVOption options[] = {
     { "5.1", "Level 5.1", 0, AV_OPT_TYPE_CONST, { .i64 = 51 }, INT_MIN, INT_MAX, VE, "level" },
     { "5.2", "Level 5.2", 0, AV_OPT_TYPE_CONST, { .i64 = 52 }, INT_MIN, INT_MAX, VE, "level" },
 
-    { "allow_sw", "Allow software encoding", OFFSET(allow_sw), AV_OPT_TYPE_BOOL,
-        { .i64 = 0 }, 0, 1, VE },
-
     { "coder", "Entropy coding", OFFSET(entropy), AV_OPT_TYPE_INT, { .i64 = VT_ENTROPY_NOT_SET }, VT_ENTROPY_NOT_SET, VT_CABAC, VE, "coder" },
     { "cavlc", "CAVLC entropy coding", 0, AV_OPT_TYPE_CONST, { .i64 = VT_CAVLC }, INT_MIN, INT_MAX, VE, "coder" },
     { "vlc",   "CAVLC entropy coding", 0, AV_OPT_TYPE_CONST, { .i64 = VT_CAVLC }, INT_MIN, INT_MAX, VE, "coder" },
     { "cabac", "CABAC entropy coding", 0, AV_OPT_TYPE_CONST, { .i64 = VT_CABAC }, INT_MIN, INT_MAX, VE, "coder" },
     { "ac",    "CABAC entropy coding", 0, AV_OPT_TYPE_CONST, { .i64 = VT_CABAC }, INT_MIN, INT_MAX, VE, "coder" },
 
-    { "realtime", "Hint that encoding should happen in real-time if not faster (e.g. capturing from camera).",
-        OFFSET(realtime), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, VE },
-
-    { "frames_before", "Other frames will come before the frames in this session. This helps smooth concatenation issues.",
-        OFFSET(frames_before), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, VE },
-    { "frames_after", "Other frames will come after the frames in this session. This helps smooth concatenation issues.",
-        OFFSET(frames_after), AV_OPT_TYPE_BOOL, { .i64 = 0 }, 0, 1, VE },
-
     { "a53cc", "Use A53 Closed Captions (if available)", OFFSET(a53_cc), AV_OPT_TYPE_BOOL, {.i64 = 1}, 0, 1, VE },
 
+    COMMON_OPTIONS
     { NULL },
 };
 
 static const AVClass h264_videotoolbox_class = {
     .class_name = "h264_videotoolbox",
     .item_name  = av_default_item_name,
-    .option     = options,
+    .option     = h264_options,
     .version    = LIBAVUTIL_VERSION_INT,
 };
 
@@ -2488,3 +2574,35 @@ AVCodec ff_h264_videotoolbox_encoder = {
     .caps_internal    = FF_CODEC_CAP_INIT_THREADSAFE |
                         FF_CODEC_CAP_INIT_CLEANUP,
 };
+
+static const AVOption hevc_options[] = {
+    { "profile", "Profile", OFFSET(profile), AV_OPT_TYPE_INT, { .i64 = HEVC_PROF_AUTO }, HEVC_PROF_AUTO, HEVC_PROF_COUNT, VE, "profile" },
+    { "main",     "Main Profile",     0, AV_OPT_TYPE_CONST, { .i64 = HEVC_PROF_MAIN   }, INT_MIN, INT_MAX, VE, "profile" },
+    { "main10",   "Main10 Profile",   0, AV_OPT_TYPE_CONST, { .i64 = HEVC_PROF_MAIN10 }, INT_MIN, INT_MAX, VE, "profile" },
+
+    COMMON_OPTIONS
+    { NULL },
+};
+
+static const AVClass hevc_videotoolbox_class = {
+    .class_name = "hevc_videotoolbox",
+    .item_name  = av_default_item_name,
+    .option     = hevc_options,
+    .version    = LIBAVUTIL_VERSION_INT,
+};
+
+AVCodec ff_hevc_videotoolbox_encoder = {
+    .name             = "hevc_videotoolbox",
+    .long_name        = NULL_IF_CONFIG_SMALL("VideoToolbox H.265 Encoder"),
+    .type             = AVMEDIA_TYPE_VIDEO,
+    .id               = AV_CODEC_ID_HEVC,
+    .priv_data_size   = sizeof(VTEncContext),
+    .pix_fmts         = pix_fmts,
+    .init             = vtenc_init,
+    .encode2          = vtenc_frame,
+    .close            = vtenc_close,
+    .capabilities     = AV_CODEC_CAP_DELAY,
+    .priv_class       = &hevc_videotoolbox_class,
+    .caps_internal    = FF_CODEC_CAP_INIT_THREADSAFE |
+                        FF_CODEC_CAP_INIT_CLEANUP,
+};