Add libvmaf_rc support
libvmaf_rc is the new release candidate API that support VMAF_NEG
calculations. This CL enables VMAF_NEG calculations in libaom.
Change-Id: I430eaa390858f8f6b6d91da8a09bbcab5a898361
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6f7bf29..7e1f730 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -432,6 +432,10 @@
list(APPEND AOM_APP_TARGETS ${AOM_ENCODER_EXAMPLE_TARGETS}
${AOM_ENCODER_TOOL_TARGETS})
+ if(CONFIG_USE_VMAF_RC AND NOT CONFIG_TUNE_VMAF)
+ message(FATAL_ERROR "Turn on CONFIG_TUNE_VMAF to use CONFIG_USE_VMAF_RC.")
+ endif()
+
if(CONFIG_TUNE_VMAF)
find_package(PkgConfig)
if(PKG_CONFIG_FOUND)
@@ -446,6 +450,21 @@
message(FATAL_ERROR "VMAF library not found.")
endif()
endif()
+
+ if(CONFIG_USE_VMAF_RC)
+ if(PKG_CONFIG_FOUND)
+ pkg_check_modules(VMAF_RC REQUIRED libvmaf_rc)
+ target_link_libraries(aom
+ PRIVATE ${VMAF_RC_LDFLAGS} ${VMAF_RC_LIBRARIES})
+ target_include_directories(aom PRIVATE ${VMAF_RC_INCLUDE_DIRS})
+ if(VMAF_RC_CFLAGS)
+ append_compiler_flag("${VMAF_RC_CFLAGS}")
+ endif()
+ else()
+ message(FATAL_ERROR "CONFIG_USE_VMAF_RC error: pkg-config not found.")
+ endif()
+ endif()
+
set_target_properties(aom PROPERTIES LINKER_LANGUAGE CXX)
if(BUILD_SHARED_LIBS)
set_target_properties(aom_static PROPERTIES LINKER_LANGUAGE CXX)
diff --git a/aom_dsp/vmaf.c b/aom_dsp/vmaf.c
index 4241f93..b0be482 100644
--- a/aom_dsp/vmaf.c
+++ b/aom_dsp/vmaf.c
@@ -9,6 +9,8 @@
* PATENTS file, you can obtain it at www.aomedia.org/license/patent.
*/
+#include "aom_dsp/vmaf.h"
+
#include <assert.h>
#include <libvmaf.h>
#include <stdio.h>
@@ -20,8 +22,11 @@
#include <unistd.h>
#endif
+#if CONFIG_USE_VMAF_RC
+#include <libvmaf.rc.h>
+#endif
+
#include "aom_dsp/blend.h"
-#include "aom_dsp/vmaf.h"
#include "aom_ports/system_state.h"
typedef struct FrameData {
@@ -168,3 +173,103 @@
aom_clear_system_state();
}
+
+#if CONFIG_USE_VMAF_RC
+void aom_init_vmaf_rc(VmafModel **vmaf_model, const char *model_path) {
+ if (*vmaf_model != NULL) return;
+ VmafModelConfig model_cfg;
+ model_cfg.flags = VMAF_MODEL_FLAG_DISABLE_CLIP;
+ model_cfg.name = "vmaf";
+ model_cfg.path = (char *)model_path;
+
+ if (vmaf_model_load_from_path(vmaf_model, &model_cfg)) {
+ vmaf_fatal_error("Failed to load VMAF model.");
+ }
+}
+
+void aom_close_vmaf_rc(VmafModel *vmaf_model) {
+ vmaf_model_destroy(vmaf_model);
+}
+
+static void copy_picture(const int bit_depth, const YV12_BUFFER_CONFIG *src,
+ VmafPicture *dst) {
+ const int width = src->y_width;
+ const int height = src->y_height;
+
+ if (bit_depth > 8) {
+ uint16_t *src_ptr = CONVERT_TO_SHORTPTR(src->y_buffer);
+ uint16_t *dst_ptr = dst->data[0];
+
+ for (int row = 0; row < height; ++row) {
+ memcpy(dst_ptr, src_ptr, width * sizeof(dst_ptr[0]));
+ src_ptr += src->y_stride;
+ dst_ptr += dst->stride[0] / 2;
+ }
+ } else {
+ uint8_t *src_ptr = src->y_buffer;
+ uint8_t *dst_ptr = (uint8_t *)dst->data[0];
+
+ for (int row = 0; row < height; ++row) {
+ memcpy(dst_ptr, src_ptr, width * sizeof(dst_ptr[0]));
+ src_ptr += src->y_stride;
+ dst_ptr += dst->stride[0];
+ }
+ }
+}
+
+void aom_calc_vmaf_rc(VmafModel *vmaf_model, const YV12_BUFFER_CONFIG *source,
+ const YV12_BUFFER_CONFIG *distorted, int bit_depth,
+ int cal_vmaf_neg, double *vmaf) {
+ VmafConfiguration cfg;
+ cfg.log_level = VMAF_LOG_LEVEL_NONE;
+ cfg.n_threads = 0;
+ cfg.n_subsample = 0;
+ cfg.cpumask = 0;
+
+ VmafContext *vmaf_context;
+ if (vmaf_init(&vmaf_context, cfg)) {
+ vmaf_fatal_error("Failed to init VMAF context.");
+ }
+
+ if (vmaf_use_features_from_model(vmaf_context, vmaf_model)) {
+ vmaf_fatal_error("Failed to load feature extractors from VMAF model.");
+ }
+
+ if (cal_vmaf_neg) {
+ VmafFeatureDictionary *vif_feature = NULL;
+ vmaf_feature_dictionary_set(&vif_feature, "vif_enhn_gain_limit", "1.0");
+ if (vmaf_use_feature(vmaf_context, "float_vif", vif_feature)) {
+ vmaf_fatal_error("Failed to use feature float_vif.");
+ }
+
+ VmafFeatureDictionary *adm_feature = NULL;
+ vmaf_feature_dictionary_set(&adm_feature, "adm_enhn_gain_limit", "1.0");
+ if (vmaf_use_feature(vmaf_context, "float_adm", adm_feature)) {
+ vmaf_fatal_error("Failed to use feature float_adm.");
+ }
+ }
+
+ VmafPicture ref, dist;
+ if (vmaf_picture_alloc(&ref, VMAF_PIX_FMT_YUV420P, bit_depth, source->y_width,
+ source->y_height) ||
+ vmaf_picture_alloc(&dist, VMAF_PIX_FMT_YUV420P, bit_depth,
+ source->y_width, source->y_height)) {
+ vmaf_fatal_error("Failed to alloc VMAF pictures.");
+ }
+ copy_picture(bit_depth, source, &ref);
+ copy_picture(bit_depth, distorted, &dist);
+ if (vmaf_read_pictures(vmaf_context, &ref, &dist, /*picture index=*/0)) {
+ vmaf_fatal_error("Failed to read VMAF pictures.");
+ }
+
+ vmaf_picture_unref(&ref);
+ vmaf_picture_unref(&dist);
+
+ vmaf_score_at_index(vmaf_context, vmaf_model, vmaf, 0);
+
+ if (vmaf_close(vmaf_context)) {
+ vmaf_fatal_error("Failed to close VMAF context.");
+ }
+}
+
+#endif // CONFIG_USE_VMAF_RC
diff --git a/aom_dsp/vmaf.h b/aom_dsp/vmaf.h
index 6079944..02e59ed 100644
--- a/aom_dsp/vmaf.h
+++ b/aom_dsp/vmaf.h
@@ -14,6 +14,11 @@
#include "aom_scale/yv12config.h"
+#if CONFIG_USE_VMAF_RC
+typedef struct VmafContext VmafContext;
+typedef struct VmafModel VmafModel;
+#endif
+
typedef struct {
// Stores the scaling factors for rdmult when tuning for VMAF.
// rdmult_scaling_factors[row * num_cols + col] stores the scaling factors for
@@ -28,6 +33,11 @@
// Stores the filter strength of the last frame.
double last_frame_unsharp_amount;
+
+#if CONFIG_USE_VMAF_RC
+ // VMAF model used in VMAF caculations.
+ VmafModel *vmaf_model;
+#endif
} TuneVMAFInfo;
void aom_calc_vmaf(const char *model_path, const YV12_BUFFER_CONFIG *source,
@@ -40,4 +50,14 @@
int stride_byte, void *user_data),
int frame_width, int frame_height, int bit_depth, double *vmaf);
+#if CONFIG_USE_VMAF_RC
+void aom_init_vmaf_rc(VmafModel **vmaf_model, const char *model_path);
+
+void aom_calc_vmaf_rc(VmafModel *vmaf_model, const YV12_BUFFER_CONFIG *source,
+ const YV12_BUFFER_CONFIG *distorted, int bit_depth,
+ int cal_vmaf_neg, double *vmaf);
+
+void aom_close_vmaf_rc(VmafModel *vmaf_model);
+#endif // CONFIG_USE_VMAF_RC
+
#endif // AOM_AOM_DSP_VMAF_H_
diff --git a/av1/av1_cx_iface.c b/av1/av1_cx_iface.c
index 5f204a7..6463ad6 100644
--- a/av1/av1_cx_iface.c
+++ b/av1/av1_cx_iface.c
@@ -2162,6 +2162,11 @@
av1_apply_encoding_flags(cpi_lap, flags);
}
+#if CONFIG_USE_VMAF_RC
+ aom_init_vmaf_rc(&cpi->vmaf_info.vmaf_model,
+ cpi->oxcf.tune_cfg.vmaf_model_path);
+#endif
+
// Handle fixed keyframe intervals
if (is_stat_generation_stage(cpi)) {
if (ctx->cfg.kf_mode == AOM_KF_AUTO &&
diff --git a/av1/encoder/encoder_alloc.h b/av1/encoder/encoder_alloc.h
index 176dbae..0c3d210 100644
--- a/av1/encoder/encoder_alloc.h
+++ b/av1/encoder/encoder_alloc.h
@@ -241,6 +241,10 @@
#if CONFIG_TUNE_VMAF
aom_free(cpi->vmaf_info.rdmult_scaling_factors);
cpi->vmaf_info.rdmult_scaling_factors = NULL;
+
+#if CONFIG_USE_VMAF_RC
+ aom_close_vmaf_rc(cpi->vmaf_info.vmaf_model);
+#endif
#endif
release_obmc_buffers(&cpi->td.mb.obmc_buffer);
diff --git a/build/cmake/aom_config_defaults.cmake b/build/cmake/aom_config_defaults.cmake
index 298ed28..ccb720c 100644
--- a/build/cmake/aom_config_defaults.cmake
+++ b/build/cmake/aom_config_defaults.cmake
@@ -107,6 +107,7 @@
set_aom_config_var(DECODE_HEIGHT_LIMIT 0 "Set limit for decode height.")
set_aom_config_var(DECODE_WIDTH_LIMIT 0 "Set limit for decode width.")
set_aom_config_var(CONFIG_TUNE_VMAF 0 "Enable encoding tuning for VMAF.")
+set_aom_config_var(CONFIG_USE_VMAF_RC 0 "Use libvmaf_rc tune for VMAF_NEG.")
# AV1 experiment flags.
set_aom_config_var(CONFIG_SPEED_STATS 0 "AV1 experiment flag.")
