tor-browser

The Tor Browser
git clone https://git.dasho.dev/tor-browser.git
Log | Files | Refs | README | LICENSE

vmaf.c (6512B)

      1 /*
      2 * Copyright (c) 2019, Alliance for Open Media. All rights reserved.
      3 *
      4 * This source code is subject to the terms of the BSD 2 Clause License and
      5 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
      6 * was not distributed with this source code in the LICENSE file, you can
      7 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
      8 * Media Patent License 1.0 was not distributed with this source code in the
      9 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
     10 */
     11 
     12 #include "aom_dsp/vmaf.h"
     13 
     14 #include <assert.h>
     15 #include <stdio.h>
     16 #include <stdlib.h>
     17 #include <string.h>
     18 #ifdef _WIN32
     19 #include <process.h>
     20 #else
     21 #include <unistd.h>
     22 #endif
     23 
     24 #include "aom_dsp/blend.h"
     25 
     26 static void vmaf_fatal_error(const char *message) {
     27  fprintf(stderr, "Fatal error: %s\n", message);
     28  exit(EXIT_FAILURE);
     29 }
     30 
     31 void aom_init_vmaf_model(VmafModel **vmaf_model, const char *model_path) {
     32  if (*vmaf_model != NULL) return;
     33  VmafModelConfig model_cfg;
     34  model_cfg.flags = VMAF_MODEL_FLAG_DISABLE_CLIP;
     35  model_cfg.name = "vmaf";
     36 
     37  if (vmaf_model_load_from_path(vmaf_model, &model_cfg, model_path)) {
     38    vmaf_fatal_error("Failed to load VMAF model.");
     39  }
     40 }
     41 
     42 void aom_close_vmaf_model(VmafModel *vmaf_model) {
     43  vmaf_model_destroy(vmaf_model);
     44 }
     45 
     46 static void copy_picture(const int bit_depth, const YV12_BUFFER_CONFIG *src,
     47                         VmafPicture *dst) {
     48  const int width = src->y_width;
     49  const int height = src->y_height;
     50 
     51  if (bit_depth > 8) {
     52    uint16_t *src_ptr = CONVERT_TO_SHORTPTR(src->y_buffer);
     53    uint16_t *dst_ptr = dst->data[0];
     54 
     55    for (int row = 0; row < height; ++row) {
     56      memcpy(dst_ptr, src_ptr, width * sizeof(dst_ptr[0]));
     57      src_ptr += src->y_stride;
     58      dst_ptr += dst->stride[0] / 2;
     59    }
     60  } else {
     61    uint8_t *src_ptr = src->y_buffer;
     62    uint8_t *dst_ptr = (uint8_t *)dst->data[0];
     63 
     64    for (int row = 0; row < height; ++row) {
     65      memcpy(dst_ptr, src_ptr, width * sizeof(dst_ptr[0]));
     66      src_ptr += src->y_stride;
     67      dst_ptr += dst->stride[0];
     68    }
     69  }
     70 }
     71 
     72 void aom_init_vmaf_context(VmafContext **vmaf_context, VmafModel *vmaf_model,
     73                           bool cal_vmaf_neg) {
     74  // TODO(sdeng): make them CLI arguments.
     75  VmafConfiguration cfg;
     76  cfg.log_level = VMAF_LOG_LEVEL_NONE;
     77  cfg.n_threads = 0;
     78  cfg.n_subsample = 0;
     79  cfg.cpumask = 0;
     80 
     81  if (vmaf_init(vmaf_context, cfg)) {
     82    vmaf_fatal_error("Failed to init VMAF context.");
     83  }
     84 
     85  if (cal_vmaf_neg) {
     86    VmafFeatureDictionary *vif_feature = NULL;
     87    if (vmaf_feature_dictionary_set(&vif_feature, "vif_enhn_gain_limit",
     88                                    "1.0")) {
     89      vmaf_fatal_error("Failed to set vif_enhn_gain_limit.");
     90    }
     91    if (vmaf_model_feature_overload(vmaf_model, "float_vif", vif_feature)) {
     92      vmaf_fatal_error("Failed to use feature float_vif.");
     93    }
     94 
     95    VmafFeatureDictionary *adm_feature = NULL;
     96    if (vmaf_feature_dictionary_set(&adm_feature, "adm_enhn_gain_limit",
     97                                    "1.0")) {
     98      vmaf_fatal_error("Failed to set adm_enhn_gain_limit.");
     99    }
    100    if (vmaf_model_feature_overload(vmaf_model, "adm", adm_feature)) {
    101      vmaf_fatal_error("Failed to use feature float_adm.");
    102    }
    103  }
    104 
    105  VmafFeatureDictionary *motion_force_zero = NULL;
    106  if (vmaf_feature_dictionary_set(&motion_force_zero, "motion_force_zero",
    107                                  "1")) {
    108    vmaf_fatal_error("Failed to set motion_force_zero.");
    109  }
    110  if (vmaf_model_feature_overload(vmaf_model, "float_motion",
    111                                  motion_force_zero)) {
    112    vmaf_fatal_error("Failed to use feature float_motion.");
    113  }
    114 
    115  if (vmaf_use_features_from_model(*vmaf_context, vmaf_model)) {
    116    vmaf_fatal_error("Failed to load feature extractors from VMAF model.");
    117  }
    118 }
    119 
    120 void aom_close_vmaf_context(VmafContext *vmaf_context) {
    121  if (vmaf_close(vmaf_context)) {
    122    vmaf_fatal_error("Failed to close VMAF context.");
    123  }
    124 }
    125 
    126 void aom_calc_vmaf(VmafModel *vmaf_model, const YV12_BUFFER_CONFIG *source,
    127                   const YV12_BUFFER_CONFIG *distorted, int bit_depth,
    128                   bool cal_vmaf_neg, double *vmaf) {
    129  VmafContext *vmaf_context;
    130  aom_init_vmaf_context(&vmaf_context, vmaf_model, cal_vmaf_neg);
    131  const int frame_index = 0;
    132  VmafPicture ref, dist;
    133  if (vmaf_picture_alloc(&ref, VMAF_PIX_FMT_YUV420P, bit_depth, source->y_width,
    134                         source->y_height) ||
    135      vmaf_picture_alloc(&dist, VMAF_PIX_FMT_YUV420P, bit_depth,
    136                         source->y_width, source->y_height)) {
    137    vmaf_fatal_error("Failed to alloc VMAF pictures.");
    138  }
    139  copy_picture(bit_depth, source, &ref);
    140  copy_picture(bit_depth, distorted, &dist);
    141  if (vmaf_read_pictures(vmaf_context, &ref, &dist,
    142                         /*picture index=*/frame_index)) {
    143    vmaf_fatal_error("Failed to read VMAF pictures.");
    144  }
    145 
    146  if (vmaf_read_pictures(vmaf_context, NULL, NULL, 0)) {
    147    vmaf_fatal_error("Failed to flush context.");
    148  }
    149 
    150  vmaf_picture_unref(&ref);
    151  vmaf_picture_unref(&dist);
    152 
    153  vmaf_score_at_index(vmaf_context, vmaf_model, vmaf, frame_index);
    154  aom_close_vmaf_context(vmaf_context);
    155 }
    156 
    157 void aom_read_vmaf_image(VmafContext *vmaf_context,
    158                         const YV12_BUFFER_CONFIG *source,
    159                         const YV12_BUFFER_CONFIG *distorted, int bit_depth,
    160                         int frame_index) {
    161  VmafPicture ref, dist;
    162  if (vmaf_picture_alloc(&ref, VMAF_PIX_FMT_YUV420P, bit_depth, source->y_width,
    163                         source->y_height) ||
    164      vmaf_picture_alloc(&dist, VMAF_PIX_FMT_YUV420P, bit_depth,
    165                         source->y_width, source->y_height)) {
    166    vmaf_fatal_error("Failed to alloc VMAF pictures.");
    167  }
    168  copy_picture(bit_depth, source, &ref);
    169  copy_picture(bit_depth, distorted, &dist);
    170  if (vmaf_read_pictures(vmaf_context, &ref, &dist,
    171                         /*picture index=*/frame_index)) {
    172    vmaf_fatal_error("Failed to read VMAF pictures.");
    173  }
    174 
    175  vmaf_picture_unref(&ref);
    176  vmaf_picture_unref(&dist);
    177 }
    178 
    179 double aom_calc_vmaf_at_index(VmafContext *vmaf_context, VmafModel *vmaf_model,
    180                              int frame_index) {
    181  double vmaf;
    182  if (vmaf_score_at_index(vmaf_context, vmaf_model, &vmaf, frame_index)) {
    183    vmaf_fatal_error("Failed to calc VMAF scores.");
    184  }
    185  return vmaf;
    186 }
    187 
    188 void aom_flush_vmaf_context(VmafContext *vmaf_context) {
    189  if (vmaf_read_pictures(vmaf_context, NULL, NULL, 0)) {
    190    vmaf_fatal_error("Failed to flush context.");
    191  }
    192 }