diff --git a/rtc_tools/compare_videos.py b/rtc_tools/compare_videos.py index d9cb670aba..bdb0bea014 100755 --- a/rtc_tools/compare_videos.py +++ b/rtc_tools/compare_videos.py @@ -38,6 +38,12 @@ def _ParseArgs(): 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 @@ def _ParseArgs(): 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 @@ def DecodeBarcodesInVideo(options, path_to_decoder, video, stat_file): 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 @@ def main(): --test_video= --frame_analyzer= + 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 @@ def main(): 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__': diff --git a/rtc_tools/frame_analyzer/frame_analyzer.cc b/rtc_tools/frame_analyzer/frame_analyzer.cc index 49d381fe7e..56dc1f47db 100644 --- a/rtc_tools/frame_analyzer/frame_analyzer.cc +++ b/rtc_tools/frame_analyzer/frame_analyzer.cc @@ -60,6 +60,9 @@ int main(int argc, char* argv[]) { " Default: None\n" " - aligned_output_file: Where to write aligned YUV/Y4M output file." " If not present, no file will be written." + " Default: None\n" + " - yuv_directory: Where to write aligned YUV ref+test output files." + " If not present, no files will be written." " Default: None\n"; webrtc::test::CommandLineParser parser; @@ -74,6 +77,7 @@ int main(int argc, char* argv[]) { parser.SetFlag("reference_file", "ref.yuv"); parser.SetFlag("test_file", "test.yuv"); parser.SetFlag("aligned_output_file", ""); + parser.SetFlag("yuv_directory", ""); parser.SetFlag("chartjson_result_file", ""); parser.SetFlag("help", "false"); @@ -133,14 +137,21 @@ int main(int argc, char* argv[]) { if (!chartjson_result_file.empty()) { webrtc::test::WritePerfResults(chartjson_result_file); } + rtc::scoped_refptr reordered_video = + webrtc::test::GenerateAlignedReferenceVideo(reference_video, + matching_indices); std::string aligned_output_file = parser.GetFlag("aligned_output_file"); if (!aligned_output_file.empty()) { - rtc::scoped_refptr reordered_video = - webrtc::test::GenerateAlignedReferenceVideo(reference_video, - matching_indices); webrtc::test::WriteVideoToFile(reordered_video, aligned_output_file, /*fps=*/30); } + std::string yuv_directory = parser.GetFlag("yuv_directory"); + if (!yuv_directory.empty()) { + webrtc::test::WriteVideoToFile(reordered_video, yuv_directory + "ref.yuv", + /*fps=*/30); + webrtc::test::WriteVideoToFile(test_video, yuv_directory + "test.yuv", + /*fps=*/30); + } return 0; } diff --git a/rtc_tools/video_file_writer.cc b/rtc_tools/video_file_writer.cc index 87c6a1e72c..99f3bf412d 100644 --- a/rtc_tools/video_file_writer.cc +++ b/rtc_tools/video_file_writer.cc @@ -10,7 +10,6 @@ #include "rtc_tools/video_file_writer.h" -#include #include #include "api/video/i420_buffer.h"