Add option to call VMAF in compare_videos.py.

VMAF compares videos on several metrics and produces a unified score.

Calling it from compare_videos required passing in a path to a VMAF
directory, where there should be a C++ wrapper executable and a model.
For now, the relative paths to those are constant.

VMAF needs to compare aligned videos in YUV format, so two videos
(ref and test) will be saved by frame_analyzer after it has aligned
them.

Bug: webrtc:9642
Change-Id: Idddfcf6b1b235e7f925696ffc38938fb84c4ff9e
Reviewed-on: https://webrtc-review.googlesource.com/102140
Reviewed-by: Patrik Höglund <phoglund@webrtc.org>
Reviewed-by: Sami Kalliomäki <sakal@webrtc.org>
Commit-Queue: Paulina Hensman <phensman@webrtc.org>
Cr-Commit-Position: refs/heads/master@{#24876}
diff --git a/rtc_tools/compare_videos.py b/rtc_tools/compare_videos.py
index d9cb670..bdb0bea 100755
--- a/rtc_tools/compare_videos.py
+++ b/rtc_tools/compare_videos.py
@@ -38,6 +38,12 @@
                     help='Path to the frame analyzer executable.')
   parser.add_option('--aligned_output_file', type='string',
                     help='Path for output aligned YUV or Y4M file.')
+  parser.add_option('--vmaf', type='string',
+                    help='Path to VMAF executable.')
+  parser.add_option('--vmaf_model', type='string',
+                    help='Path to VMAF model.')
+  parser.add_option('--vmaf_phone_model', action='store_true',
+                    help='Whether to use phone model in VMAF.')
   parser.add_option('--barcode_decoder', type='string',
                     help=('Path to the barcode decoder script. By default, we '
                           'will assume we can find it in barcode_tools/'
@@ -88,6 +94,10 @@
   if not os.path.exists(options.frame_analyzer):
     parser.error('Cannot find frame analyzer executable at %s!' %
                  options.frame_analyzer)
+
+  if options.vmaf and not options.vmaf_model:
+    parser.error('You must provide a path to a VMAF model to use VMAF.')
+
   return options
 
 def _DevNull():
@@ -125,6 +135,65 @@
     return 1
   return 0
 
+
+def _RunFrameAnalyzer(options):
+  """Run frame analyzer to compare the videos and print output."""
+  cmd = [
+    options.frame_analyzer,
+    '--label=%s' % options.label,
+    '--reference_file=%s' % options.ref_video,
+    '--test_file=%s' % options.test_video,
+    '--stats_file_ref=%s' % options.stats_file_ref,
+    '--stats_file_test=%s' % options.stats_file_test,
+    '--width=%d' % options.yuv_frame_width,
+    '--height=%d' % options.yuv_frame_height,
+  ]
+  if options.chartjson_result_file:
+    cmd.append('--chartjson_result_file=%s' % options.chartjson_result_file)
+  if options.aligned_output_file:
+    cmd.append('--aligned_output_file=%s' % options.aligned_output_file)
+  if options.yuv_directory:
+    cmd.append('--yuv_directory=%s' % options.yuv_directory)
+  frame_analyzer = subprocess.Popen(cmd, stdin=_DevNull(),
+                                    stdout=sys.stdout, stderr=sys.stderr)
+  frame_analyzer.wait()
+  return frame_analyzer.returncode
+
+
+def _RunVmaf(options):
+  """ Run VMAF to compare videos and print output.
+
+  The provided vmaf directory is assumed to contain a c++ wrapper executable
+  and a model.
+
+  The yuv_directory is assumed to have been populated with a reference and test
+  video in .yuv format, with names according to the label.
+  """
+  cmd = [
+    options.vmaf,
+    'yuv420p',
+    str(options.yuv_frame_width),
+    str(options.yuv_frame_height),
+    os.path.join(options.yuv_directory, "ref.yuv"),
+    os.path.join(options.yuv_directory, "test.yuv"),
+    options.vmaf_model,
+  ]
+  if options.vmaf_phone_model:
+    cmd.append('--phone-model')
+
+  vmaf = subprocess.Popen(cmd, stdin=_DevNull(),
+                                  stdout=subprocess.PIPE, stderr=sys.stderr)
+  vmaf.wait()
+  if vmaf.returncode != 0:
+    print 'Failed to run VMAF.'
+    return 1
+  output = vmaf.stdout.read()
+  # Extract score from VMAF output.
+  score = float(output.split('\n')[2].split()[3])
+  print 'RESULT Vmaf: %s= %f' % (options.label, score)
+  return 0
+
+
 def main():
   """The main function.
 
@@ -134,6 +203,9 @@
   --test_video=<path_and_name_of_test_video>
   --frame_analyzer=<path_and_name_of_the_frame_analyzer_executable>
 
+  Running vmaf requires the following arguments:
+  --vmaf, --vmaf_model, --yuv_frame_width, --yuv_frame_height
+
   Notice that the prerequisites for barcode_decoder.py also applies to this
   script. The means the following executables have to be available in the PATH:
   * zxing
@@ -154,28 +226,21 @@
                            options.test_video, options.stats_file_test) != 0:
     return 1
 
-  # Run frame analyzer to compare the videos and print output.
-  cmd = [
-    options.frame_analyzer,
-    '--label=%s' % options.label,
-    '--reference_file=%s' % options.ref_video,
-    '--test_file=%s' % options.test_video,
-    '--stats_file_ref=%s' % options.stats_file_ref,
-    '--stats_file_test=%s' % options.stats_file_test,
-    '--width=%d' % options.yuv_frame_width,
-    '--height=%d' % options.yuv_frame_height,
-  ]
-  if options.chartjson_result_file:
-    cmd.append('--chartjson_result_file=%s' % options.chartjson_result_file)
-  if options.aligned_output_file:
-    cmd.append('--aligned_output_file=%s' % options.aligned_output_file)
-  frame_analyzer = subprocess.Popen(cmd, stdin=_DevNull(),
-                                    stdout=sys.stdout, stderr=sys.stderr)
-  frame_analyzer.wait()
-  if frame_analyzer.returncode != 0:
+  # Create a directory to save temporary YUV files for VMAF in frame_analyzer.
+  if options.vmaf:
+    options.yuv_directory = tempfile.mkdtemp() + '/'
+
+  # Run frame_analyzer to compare the videos and print output.
+  if _RunFrameAnalyzer(options) != 0:
     print 'Failed to run frame analyzer.'
     return 1
 
+  # Run VMAF for further video comparison and print output.
+  if options.vmaf:
+    return_code = _RunVmaf(options)
+    shutil.rmtree(options.yuv_directory)
+    return return_code
+
   return 0
 
 if __name__ == '__main__':