1# Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Output the test video parameters for VDA tests.
6
7Output the test parameters of h264/vp8 videos for running VDA tests:
8filename:width:height:frames:fragments:minFPSwithRender:minFPSnoRender:profile
9(chromium content/common/gpu/media/video_decode_accelerator_unittest.cc)
10
11Executing this script with a single h264 or vp8 video:
12'output_test_video_params.py video.h264|video.vp8' will directly print the test
13parameters for that video on screen.
14
15Executing this script with a directory containing h264/vp8 videos:
16'output_test_video_params.py video-dir/' will output
17__test_video_list_[timestamp] file that contains a list of test parameters for
18all h264/vp8 videos under video-dir. Only valid test parameters will be written
19into the output file; unsupported videos/files under the video-dir will be
20ignored.
21"""
22
23import mmap
24import os
25import re
26import struct
27import subprocess
28import sys
29import time
30
31INVALID_PARAM = '##'
32
33def h264_fragments(path):
34    with open(path, "rb") as f:
35        mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
36
37    # Count NAL by searching for the 0x00000001 start code prefix.
38    pattern = '00000001'.decode('hex')
39    regex = re.compile(pattern, re.MULTILINE)
40    count = sum(1 for _ in regex.finditer(mm))
41    mm.close()
42
43    return str(count)
44
45def vp8_fragments(path):
46    # Read IVF 32-byte header and parse the frame number in the header.
47    # IVF header definition: http://wiki.multimedia.cx/index.php?title=IVF
48    with open(path, "rb") as f:
49        mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
50
51    mm.seek(24, 0)
52    header_frame_num = struct.unpack('i', mm.read(4))[0]
53    mm.seek(32, 0)
54
55    # Count the following frames and check if the count matches the frame number
56    # in the header.
57    count = 0
58    size = mm.size()
59    while mm.tell() + 4 <= size:
60        (frame,) = struct.unpack('i', mm.read(4))
61        offset = 8 + frame
62        if (mm.tell() + offset <= size):
63            mm.seek(offset, 1)
64            count = count + 1
65
66    mm.close()
67    if header_frame_num != count:
68        return INVALID_PARAM
69
70    return str(count)
71
72def full_test_parameters(path):
73    try:
74        ffmpeg_cmd = ["ffmpeg", "-i", path, "-vcodec", "copy", "-an", "-f",
75                      "null", "/dev/null"]
76        content = subprocess.check_output(ffmpeg_cmd, stderr=subprocess.STDOUT)
77    except subprocess.CalledProcessError as e:
78        content = e.output
79
80    # Get video format, dimension, and frame number from the ffmpeg output.
81    # Sample of ffmpeg output:
82    # Input #0, h264, from 'video.h264':
83    # Stream #0.0: Video: h264 (High), yuv420p, 640x360, 25 fps, ...
84    # frame=   82 fps=  0 ...
85    results = re.findall('Input #0, (\S+),', content)
86    profile = INVALID_PARAM
87    frag = INVALID_PARAM
88    if (results):
89        video_format = results[0].lower()
90        if (video_format == 'h264'):
91            profile = '1'
92            frag = h264_fragments(path)
93        elif (video_format == 'ivf'):
94            profile = '11'
95            frag = vp8_fragments(path)
96
97    dimen = [INVALID_PARAM, INVALID_PARAM]
98    fps = [INVALID_PARAM, INVALID_PARAM]
99    results = re.findall('Stream #0.*Video:.* (\d+)x(\d+)', content)
100    if (results):
101        dimen = results[0]
102        fps = ['30', '30']
103
104    results = re.findall('frame= *(\d+)', content)
105    frame = results[0] if results else INVALID_PARAM
106
107    filename = os.path.basename(path)
108    return '%s:%s:%s:%s:%s:%s:%s:%s' % (
109            filename, dimen[0], dimen[1], frame, frag, fps[0], fps[1], profile)
110
111def check_before_output(line):
112    if INVALID_PARAM in line:
113        print 'Warning: %s' % line
114        return False
115    return True
116
117def main(argv):
118    if len(argv) != 1:
119        print 'Please provide a h264/vp8 directory or file.'
120        sys.exit(1)
121
122    if os.path.isdir(argv[0]):
123        name = '__test_video_list_%s' % time.strftime("%Y%m%d_%H%M%S")
124        with open(name, 'w') as output:
125            output.write('[\n')
126            for root, _, files in os.walk(argv[0]):
127                for f in files:
128                    path = os.path.join(root, f)
129                    line = full_test_parameters(path)
130                    if check_before_output(line):
131                        # Output in json format (no trailing comma in the list.)
132                        sep = '' if f is files[-1] else ','
133                        output.write('\"%s\"%s\n' % (line, sep))
134
135            output.write(']\n')
136    elif os.path.isfile(argv[0]):
137        line = full_test_parameters(argv[0])
138        if check_before_output(line):
139            print line
140    else:
141        print 'Invalid input.'
142
143main(sys.argv[1:])
144