4
4
# This source code is licensed under the BSD-style license found in the
5
5
# LICENSE file in the root directory of this source tree.
6
6
7
+ import argparse
7
8
import json
8
9
import os
9
10
import platform
26
27
def main () -> None :
27
28
"""Benchmarks the performance of a few video decoders on synthetic videos"""
28
29
30
+ parser = argparse .ArgumentParser ()
31
+ parser .add_argument (
32
+ "--test_run" ,
33
+ help = "Test run only; use small values for experiments to ensure everything works. Does not overwrite the data file." ,
34
+ action = "store_true" ,
35
+ )
36
+ args = parser .parse_args ()
37
+
38
+ # The logic is clearer internally if we invert the boolean. However, we want to
39
+ # maintain the external default that a test run is off by default.
40
+ data_generation_run = not args .test_run
41
+
42
+ if data_generation_run :
43
+ resolutions = ["1280x720" ]
44
+ encodings = ["libx264" ]
45
+ patterns = ["mandelbrot" ]
46
+ fpses = [60 ]
47
+ gop_sizes = [600 ]
48
+ durations = [120 ]
49
+ pix_fmts = ["yuv420p" ]
50
+ ffmpeg_path = "ffmpeg"
51
+ min_runtime_seconds = 30
52
+
53
+ # These are the number of uniform seeks we do in the seek+decode benchmark.
54
+ num_samples = 10
55
+ else :
56
+ resolutions = ["640x480" ]
57
+ encodings = ["libx264" ]
58
+ patterns = ["mandelbrot" ]
59
+ fpses = [30 ]
60
+ gop_sizes = [20 ]
61
+ durations = [10 ] # if this goes too low, we hit EOF errors in some decoders
62
+ pix_fmts = ["yuv420p" ]
63
+ ffmpeg_path = "ffmpeg"
64
+ min_runtime_seconds = 1
65
+
66
+ num_samples = 4
67
+
29
68
videos_dir_path = "/tmp/torchcodec_benchmarking_videos"
30
69
shutil .rmtree (videos_dir_path , ignore_errors = True )
31
70
os .makedirs (videos_dir_path )
32
71
33
- resolutions = ["1280x720" ]
34
- encodings = ["libx264" ]
35
- patterns = ["mandelbrot" ]
36
- fpses = [60 ]
37
- gop_sizes = [600 ]
38
- durations = [120 ]
39
- pix_fmts = ["yuv420p" ]
40
- ffmpeg_path = "ffmpeg"
41
72
generate_videos (
42
73
resolutions ,
43
74
encodings ,
@@ -61,15 +92,13 @@ def main() -> None:
61
92
decoder_dict ["TorchAudio" ] = TorchAudioDecoder ()
62
93
decoder_dict ["Decord" ] = DecordAccurateBatch ()
63
94
64
- # These are the number of uniform seeks we do in the seek+decode benchmark.
65
- num_samples = 10
66
95
video_files_paths = list (Path (videos_dir_path ).glob ("*.mp4" ))
67
96
df_data = run_benchmarks (
68
97
decoder_dict ,
69
98
video_files_paths ,
70
99
num_samples ,
71
100
num_sequential_frames_from_start = [100 ],
72
- min_runtime_seconds = 30 ,
101
+ min_runtime_seconds = min_runtime_seconds ,
73
102
benchmark_video_creation = False ,
74
103
)
75
104
df_data .append (
@@ -82,9 +111,10 @@ def main() -> None:
82
111
}
83
112
)
84
113
85
- data_json = Path (__file__ ).parent / "benchmark_readme_data.json"
86
- with open (data_json , "w" ) as write_file :
87
- json .dump (df_data , write_file , sort_keys = True , indent = 4 )
114
+ if data_generation_run :
115
+ data_json = Path (__file__ ).parent / "benchmark_readme_data.json"
116
+ with open (data_json , "w" ) as write_file :
117
+ json .dump (df_data , write_file , sort_keys = True , indent = 4 )
88
118
89
119
90
120
if __name__ == "__main__" :
0 commit comments