1#!/usr/bin/env python
2# Copyright (c) 2013 The WebRTC project authors. All Rights Reserved.
3#
4# Use of this source code is governed by a BSD-style license
5# that can be found in the LICENSE file in the root of the source
6# tree. An additional intellectual property rights grant can be found
7# in the file PATENTS.  All contributing project authors may
8# be found in the AUTHORS file in the root of the source tree.
9
10from __future__ import absolute_import
11from __future__ import division
12from __future__ import print_function
13import json
14import optparse
15import os
16import shutil
17import subprocess
18import sys
19import tempfile
20
21
22SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
23
24# Chrome browsertests will throw away stderr; avoid that output gets lost.
25sys.stderr = sys.stdout
26
27
28def _ParseArgs():
29  """Registers the command-line options."""
30  usage = 'usage: %prog [options]'
31  parser = optparse.OptionParser(usage=usage)
32
33  parser.add_option('--label', type='string', default='MY_TEST',
34                    help=('Label of the test, used to identify different '
35                          'tests. Default: %default'))
36  parser.add_option('--ref_video', type='string',
37                    help='Reference video to compare with (YUV).')
38  parser.add_option('--test_video', type='string',
39                    help=('Test video to be compared with the reference '
40                          'video (YUV).'))
41  parser.add_option('--frame_analyzer', type='string',
42                    help='Path to the frame analyzer executable.')
43  parser.add_option('--aligned_output_file', type='string',
44                    help='Path for output aligned YUV or Y4M file.')
45  parser.add_option('--vmaf', type='string',
46                    help='Path to VMAF executable.')
47  parser.add_option('--vmaf_model', type='string',
48                    help='Path to VMAF model.')
49  parser.add_option('--vmaf_phone_model', action='store_true',
50                    help='Whether to use phone model in VMAF.')
51  parser.add_option('--yuv_frame_width', type='int', default=640,
52                    help='Width of the YUV file\'s frames. Default: %default')
53  parser.add_option('--yuv_frame_height', type='int', default=480,
54                    help='Height of the YUV file\'s frames. Default: %default')
55  parser.add_option('--chartjson_result_file', type='str', default=None,
56                    help='Where to store perf results in chartjson format.')
57  options, _ = parser.parse_args()
58
59  if not options.ref_video:
60    parser.error('You must provide a path to the reference video!')
61  if not os.path.exists(options.ref_video):
62    parser.error('Cannot find the reference video at %s' % options.ref_video)
63
64  if not options.test_video:
65    parser.error('You must provide a path to the test video!')
66  if not os.path.exists(options.test_video):
67    parser.error('Cannot find the test video at %s' % options.test_video)
68
69  if not options.frame_analyzer:
70    parser.error('You must provide the path to the frame analyzer executable!')
71  if not os.path.exists(options.frame_analyzer):
72    parser.error('Cannot find frame analyzer executable at %s!' %
73                 options.frame_analyzer)
74
75  if options.vmaf and not options.vmaf_model:
76    parser.error('You must provide a path to a VMAF model to use VMAF.')
77
78  return options
79
80def _DevNull():
81  """On Windows, sometimes the inherited stdin handle from the parent process
82  fails. Workaround this by passing null to stdin to the subprocesses commands.
83  This function can be used to create the null file handler.
84  """
85  return open(os.devnull, 'r')
86
87
88def _RunFrameAnalyzer(options, yuv_directory=None):
89  """Run frame analyzer to compare the videos and print output."""
90  cmd = [
91    options.frame_analyzer,
92    '--label=%s' % options.label,
93    '--reference_file=%s' % options.ref_video,
94    '--test_file=%s' % options.test_video,
95    '--width=%d' % options.yuv_frame_width,
96    '--height=%d' % options.yuv_frame_height,
97  ]
98  if options.chartjson_result_file:
99    cmd.append('--chartjson_result_file=%s' % options.chartjson_result_file)
100  if options.aligned_output_file:
101    cmd.append('--aligned_output_file=%s' % options.aligned_output_file)
102  if yuv_directory:
103    cmd.append('--yuv_directory=%s' % yuv_directory)
104  frame_analyzer = subprocess.Popen(cmd, stdin=_DevNull(),
105                                    stdout=sys.stdout, stderr=sys.stderr)
106  frame_analyzer.wait()
107  if frame_analyzer.returncode != 0:
108    print('Failed to run frame analyzer.')
109  return frame_analyzer.returncode
110
111
112def _RunVmaf(options, yuv_directory, logfile):
113  """ Run VMAF to compare videos and print output.
114
115  The yuv_directory is assumed to have been populated with a reference and test
116  video in .yuv format, with names according to the label.
117  """
118  cmd = [
119    options.vmaf,
120    'yuv420p',
121    str(options.yuv_frame_width),
122    str(options.yuv_frame_height),
123    os.path.join(yuv_directory, "ref.yuv"),
124    os.path.join(yuv_directory, "test.yuv"),
125    options.vmaf_model,
126    '--log',
127    logfile,
128    '--log-fmt',
129    'json',
130  ]
131  if options.vmaf_phone_model:
132    cmd.append('--phone-model')
133
134  vmaf = subprocess.Popen(cmd, stdin=_DevNull(),
135                          stdout=sys.stdout, stderr=sys.stderr)
136  vmaf.wait()
137  if vmaf.returncode != 0:
138    print('Failed to run VMAF.')
139    return 1
140
141  # Read per-frame scores from VMAF output and print.
142  with open(logfile) as f:
143    vmaf_data = json.load(f)
144    vmaf_scores = []
145    for frame in vmaf_data['frames']:
146      vmaf_scores.append(frame['metrics']['vmaf'])
147    print('RESULT VMAF: %s=' % options.label, vmaf_scores)
148
149  return 0
150
151
152def main():
153  """The main function.
154
155  A simple invocation is:
156  ./webrtc/rtc_tools/compare_videos.py
157  --ref_video=<path_and_name_of_reference_video>
158  --test_video=<path_and_name_of_test_video>
159  --frame_analyzer=<path_and_name_of_the_frame_analyzer_executable>
160
161  Running vmaf requires the following arguments:
162  --vmaf, --vmaf_model, --yuv_frame_width, --yuv_frame_height
163  """
164  options = _ParseArgs()
165
166  if options.vmaf:
167    try:
168      # Directory to save temporary YUV files for VMAF in frame_analyzer.
169      yuv_directory = tempfile.mkdtemp()
170      _, vmaf_logfile = tempfile.mkstemp()
171
172      # Run frame analyzer to compare the videos and print output.
173      if _RunFrameAnalyzer(options, yuv_directory=yuv_directory) != 0:
174        return 1
175
176      # Run VMAF for further video comparison and print output.
177      return _RunVmaf(options, yuv_directory, vmaf_logfile)
178    finally:
179      shutil.rmtree(yuv_directory)
180      os.remove(vmaf_logfile)
181  else:
182    return _RunFrameAnalyzer(options)
183
184  return 0
185
186if __name__ == '__main__':
187  sys.exit(main())
188