1#!/usr/bin/env python
2
3import operator
4
5from optparse import OptionGroup
6
7import sys
8
9from time import time
10
11from digress.cli import Dispatcher as _Dispatcher
12from digress.errors import ComparisonError, FailedTestError, DisabledTestError
13from digress.testing import depends, comparer, Fixture, Case
14from digress.comparers import compare_pass
15from digress.scm import git as x264git
16
17from subprocess import Popen, PIPE, STDOUT
18
19import os
20import re
21import shlex
22import inspect
23
24from random import randrange, seed
25from math import ceil
26
27from itertools import imap, izip
28
29os.chdir(os.path.join(os.path.dirname(__file__), ".."))
30
31# options
32
33OPTIONS = [
34    [ "--tune %s" % t for t in ("film", "zerolatency") ],
35    ("", "--intra-refresh"),
36    ("", "--no-cabac"),
37    ("", "--interlaced"),
38    ("", "--slice-max-size 1000"),
39    ("", "--frame-packing 5"),
40    [ "--preset %s" % p for p in ("ultrafast",
41                                  "superfast",
42                                  "veryfast",
43                                  "faster",
44                                  "fast",
45                                  "medium",
46                                  "slow",
47                                  "slower",
48                                  "veryslow",
49                                  "placebo") ]
50]
51
52# end options
53
54def compare_yuv_output(width, height):
55    def _compare_yuv_output(file_a, file_b):
56        size_a = os.path.getsize(file_a)
57        size_b = os.path.getsize(file_b)
58
59        if size_a != size_b:
60            raise ComparisonError("%s is not the same size as %s" % (
61                file_a,
62                file_b
63            ))
64
65        BUFFER_SIZE = 8196
66
67        offset = 0
68
69        with open(file_a) as f_a:
70            with open(file_b) as f_b:
71                for chunk_a, chunk_b in izip(
72                    imap(
73                        lambda i: f_a.read(BUFFER_SIZE),
74                        xrange(size_a // BUFFER_SIZE + 1)
75                    ),
76                    imap(
77                        lambda i: f_b.read(BUFFER_SIZE),
78                        xrange(size_b // BUFFER_SIZE + 1)
79                    )
80                ):
81                    chunk_size = len(chunk_a)
82
83                    if chunk_a != chunk_b:
84                        for i in xrange(chunk_size):
85                            if chunk_a[i] != chunk_b[i]:
86                                # calculate the macroblock, plane and frame from the offset
87                                offs = offset + i
88
89                                y_plane_area = width * height
90                                u_plane_area = y_plane_area + y_plane_area * 0.25
91                                v_plane_area = u_plane_area + y_plane_area * 0.25
92
93                                pixel = offs % v_plane_area
94                                frame = offs // v_plane_area
95
96                                if pixel < y_plane_area:
97                                    plane = "Y"
98
99                                    pixel_x = pixel % width
100                                    pixel_y = pixel // width
101
102                                    macroblock = (ceil(pixel_x / 16.0), ceil(pixel_y / 16.0))
103                                elif pixel < u_plane_area:
104                                    plane = "U"
105
106                                    pixel -= y_plane_area
107
108                                    pixel_x = pixel % width
109                                    pixel_y = pixel // width
110
111                                    macroblock = (ceil(pixel_x / 8.0), ceil(pixel_y / 8.0))
112                                else:
113                                    plane = "V"
114
115                                    pixel -= u_plane_area
116
117                                    pixel_x = pixel % width
118                                    pixel_y = pixel // width
119
120                                    macroblock = (ceil(pixel_x / 8.0), ceil(pixel_y / 8.0))
121
122                                macroblock = tuple([ int(x) for x in macroblock ])
123
124                                raise ComparisonError("%s differs from %s at frame %d, " \
125                                                      "macroblock %s on the %s plane (offset %d)" % (
126                                    file_a,
127                                    file_b,
128                                    frame,
129                                    macroblock,
130                                    plane,
131                                    offs)
132                                )
133
134                    offset += chunk_size
135
136    return _compare_yuv_output
137
138def program_exists(program):
139    def is_exe(fpath):
140        return os.path.exists(fpath) and os.access(fpath, os.X_OK)
141
142    fpath, fname = os.path.split(program)
143
144    if fpath:
145        if is_exe(program):
146            return program
147    else:
148        for path in os.environ["PATH"].split(os.pathsep):
149            exe_file = os.path.join(path, program)
150            if is_exe(exe_file):
151                return exe_file
152
153    return None
154
155class x264(Fixture):
156    scm = x264git
157
158class Compile(Case):
159    @comparer(compare_pass)
160    def test_configure(self):
161        Popen([
162            "make",
163            "distclean"
164        ], stdout=PIPE, stderr=STDOUT).communicate()
165
166        configure_proc = Popen([
167            "./configure"
168        ] + self.fixture.dispatcher.configure, stdout=PIPE, stderr=STDOUT)
169
170        output = configure_proc.communicate()[0]
171        if configure_proc.returncode != 0:
172            raise FailedTestError("configure failed: %s" % output.replace("\n", " "))
173
174    @depends("configure")
175    @comparer(compare_pass)
176    def test_make(self):
177        make_proc = Popen([
178            "make",
179            "-j5"
180        ], stdout=PIPE, stderr=STDOUT)
181
182        output = make_proc.communicate()[0]
183        if make_proc.returncode != 0:
184            raise FailedTestError("make failed: %s" % output.replace("\n", " "))
185
186_dimension_pattern = re.compile(r"\w+ [[]info[]]: (\d+)x(\d+)[pi] \d+:\d+ @ \d+/\d+ fps [(][vc]fr[)]")
187
188def _YUVOutputComparisonFactory():
189    class YUVOutputComparison(Case):
190        _dimension_pattern = _dimension_pattern
191
192        depends = [ Compile ]
193        options = []
194
195        def __init__(self):
196            for name, meth in inspect.getmembers(self):
197                if name[:5] == "test_" and name[5:] not in self.fixture.dispatcher.yuv_tests:
198                    delattr(self.__class__, name)
199
200        def _run_x264(self):
201            x264_proc = Popen([
202                "./x264",
203                "-o",
204                "%s.264" % self.fixture.dispatcher.video,
205                "--dump-yuv",
206                "x264-output.yuv"
207            ] + self.options + [
208                self.fixture.dispatcher.video
209            ], stdout=PIPE, stderr=STDOUT)
210
211            output = x264_proc.communicate()[0]
212            if x264_proc.returncode != 0:
213                raise FailedTestError("x264 did not complete properly: %s" % output.replace("\n", " "))
214
215            matches = _dimension_pattern.match(output)
216
217            return (int(matches.group(1)), int(matches.group(2)))
218
219        @comparer(compare_pass)
220        def test_jm(self):
221            if not program_exists("ldecod"): raise DisabledTestError("jm unavailable")
222
223            try:
224                runres = self._run_x264()
225
226                jm_proc = Popen([
227                    "ldecod",
228                    "-i",
229                    "%s.264" % self.fixture.dispatcher.video,
230                    "-o",
231                    "jm-output.yuv"
232                ], stdout=PIPE, stderr=STDOUT)
233
234                output = jm_proc.communicate()[0]
235                if jm_proc.returncode != 0:
236                    raise FailedTestError("jm did not complete properly: %s" % output.replace("\n", " "))
237
238                try:
239                    compare_yuv_output(*runres)("x264-output.yuv", "jm-output.yuv")
240                except ComparisonError, e:
241                    raise FailedTestError(e)
242            finally:
243                try: os.remove("x264-output.yuv")
244                except: pass
245
246                try: os.remove("%s.264" % self.fixture.dispatcher.video)
247                except: pass
248
249                try: os.remove("jm-output.yuv")
250                except: pass
251
252                try: os.remove("log.dec")
253                except: pass
254
255                try: os.remove("dataDec.txt")
256                except: pass
257
258        @comparer(compare_pass)
259        def test_ffmpeg(self):
260            if not program_exists("ffmpeg"): raise DisabledTestError("ffmpeg unavailable")
261            try:
262                runres = self._run_x264()
263
264                ffmpeg_proc = Popen([
265                    "ffmpeg",
266                    "-vsync 0",
267                    "-i",
268                    "%s.264" % self.fixture.dispatcher.video,
269                    "ffmpeg-output.yuv"
270                ], stdout=PIPE, stderr=STDOUT)
271
272                output = ffmpeg_proc.communicate()[0]
273                if ffmpeg_proc.returncode != 0:
274                    raise FailedTestError("ffmpeg did not complete properly: %s" % output.replace("\n", " "))
275
276                try:
277                    compare_yuv_output(*runres)("x264-output.yuv", "ffmpeg-output.yuv")
278                except ComparisonError, e:
279                    raise FailedTestError(e)
280            finally:
281                try: os.remove("x264-output.yuv")
282                except: pass
283
284                try: os.remove("%s.264" % self.fixture.dispatcher.video)
285                except: pass
286
287                try: os.remove("ffmpeg-output.yuv")
288                except: pass
289
290    return YUVOutputComparison
291
292class Regression(Case):
293    depends = [ Compile ]
294
295    _psnr_pattern = re.compile(r"x264 [[]info[]]: PSNR Mean Y:\d+[.]\d+ U:\d+[.]\d+ V:\d+[.]\d+ Avg:\d+[.]\d+ Global:(\d+[.]\d+) kb/s:\d+[.]\d+")
296    _ssim_pattern = re.compile(r"x264 [[]info[]]: SSIM Mean Y:(\d+[.]\d+) [(]\d+[.]\d+db[)]")
297
298    def __init__(self):
299        if self.fixture.dispatcher.x264:
300            self.__class__.__name__ += " %s" % " ".join(self.fixture.dispatcher.x264)
301
302    def test_psnr(self):
303        try:
304            x264_proc = Popen([
305                "./x264",
306                "-o",
307                "%s.264" % self.fixture.dispatcher.video,
308                "--psnr"
309            ] + self.fixture.dispatcher.x264 + [
310                self.fixture.dispatcher.video
311            ], stdout=PIPE, stderr=STDOUT)
312
313            output = x264_proc.communicate()[0]
314
315            if x264_proc.returncode != 0:
316                raise FailedTestError("x264 did not complete properly: %s" % output.replace("\n", " "))
317
318            for line in output.split("\n"):
319                if line.startswith("x264 [info]: PSNR Mean"):
320                    return float(self._psnr_pattern.match(line).group(1))
321
322            raise FailedTestError("no PSNR output caught from x264")
323        finally:
324            try: os.remove("%s.264" % self.fixture.dispatcher.video)
325            except: pass
326
327    def test_ssim(self):
328        try:
329            x264_proc = Popen([
330                "./x264",
331                "-o",
332                "%s.264" % self.fixture.dispatcher.video,
333                "--ssim"
334            ] + self.fixture.dispatcher.x264 + [
335                self.fixture.dispatcher.video
336            ], stdout=PIPE, stderr=STDOUT)
337
338            output = x264_proc.communicate()[0]
339
340            if x264_proc.returncode != 0:
341                raise FailedTestError("x264 did not complete properly: %s" % output.replace("\n", " "))
342
343            for line in output.split("\n"):
344                if line.startswith("x264 [info]: SSIM Mean"):
345                    return float(self._ssim_pattern.match(line).group(1))
346
347            raise FailedTestError("no PSNR output caught from x264")
348        finally:
349            try: os.remove("%s.264" % self.fixture.dispatcher.video)
350            except: pass
351
352def _generate_random_commandline():
353    commandline = []
354
355    for suboptions in OPTIONS:
356        commandline.append(suboptions[randrange(0, len(suboptions))])
357
358    return filter(None, reduce(operator.add, [ shlex.split(opt) for opt in commandline ]))
359
360_generated = []
361
362fixture = x264()
363fixture.register_case(Compile)
364
365fixture.register_case(Regression)
366
367class Dispatcher(_Dispatcher):
368    video = "akiyo_qcif.y4m"
369    products = 50
370    configure = []
371    x264 = []
372    yuv_tests = [ "jm" ]
373
374    def _populate_parser(self):
375        super(Dispatcher, self)._populate_parser()
376
377        # don't do a whole lot with this
378        tcase = _YUVOutputComparisonFactory()
379
380        yuv_tests = [ name[5:] for name, meth in filter(lambda pair: pair[0][:5] == "test_", inspect.getmembers(tcase)) ]
381
382        group = OptionGroup(self.optparse, "x264 testing-specific options")
383
384        group.add_option(
385            "-v",
386            "--video",
387            metavar="FILENAME",
388            action="callback",
389            dest="video",
390            type=str,
391            callback=lambda option, opt, value, parser: setattr(self, "video", value),
392            help="yuv video to perform testing on (default: %s)" % self.video
393        )
394
395        group.add_option(
396            "-s",
397            "--seed",
398            metavar="SEED",
399            action="callback",
400            dest="seed",
401            type=int,
402            callback=lambda option, opt, value, parser: setattr(self, "seed", value),
403            help="seed for the random number generator (default: unix timestamp)"
404        )
405
406        group.add_option(
407            "-p",
408            "--product-tests",
409            metavar="NUM",
410            action="callback",
411            dest="video",
412            type=int,
413            callback=lambda option, opt, value, parser: setattr(self, "products", value),
414            help="number of cartesian products to generate for yuv comparison testing (default: %d)" % self.products
415        )
416
417        group.add_option(
418            "--configure-with",
419            metavar="FLAGS",
420            action="callback",
421            dest="configure",
422            type=str,
423            callback=lambda option, opt, value, parser: setattr(self, "configure", shlex.split(value)),
424            help="options to run ./configure with"
425        )
426
427        group.add_option(
428            "--yuv-tests",
429            action="callback",
430            dest="yuv_tests",
431            type=str,
432            callback=lambda option, opt, value, parser: setattr(self, "yuv_tests", [
433                val.strip() for val in value.split(",")
434            ]),
435            help="select tests to run with yuv comparisons (default: %s, available: %s)" % (
436                ", ".join(self.yuv_tests),
437                ", ".join(yuv_tests)
438            )
439        )
440
441        group.add_option(
442            "--x264-with",
443            metavar="FLAGS",
444            action="callback",
445            dest="x264",
446            type=str,
447            callback=lambda option, opt, value, parser: setattr(self, "x264", shlex.split(value)),
448            help="additional options to run ./x264 with"
449        )
450
451        self.optparse.add_option_group(group)
452
453    def pre_dispatch(self):
454        if not hasattr(self, "seed"):
455            self.seed = int(time())
456
457        print "Using seed: %d" % self.seed
458        seed(self.seed)
459
460        for i in xrange(self.products):
461            YUVOutputComparison = _YUVOutputComparisonFactory()
462
463            commandline = _generate_random_commandline()
464
465            counter = 0
466
467            while commandline in _generated:
468                counter += 1
469                commandline = _generate_random_commandline()
470
471                if counter > 100:
472                    print >>sys.stderr, "Maximum command-line regeneration exceeded. "  \
473                                        "Try a different seed or specify fewer products to generate."
474                    sys.exit(1)
475
476            commandline += self.x264
477
478            _generated.append(commandline)
479
480            YUVOutputComparison.options = commandline
481            YUVOutputComparison.__name__ = ("%s %s" % (YUVOutputComparison.__name__, " ".join(commandline)))
482
483            fixture.register_case(YUVOutputComparison)
484
485Dispatcher(fixture).dispatch()
486