stbt-match

#!/usr/bin/env python

"""
Copyright 2013 YouView TV Ltd.
License: LGPL v2.1 or (at your option) any later version (see
https://github.com/stb-tester/stb-tester/blob/master/LICENSE for details).
"""

import argparse
import os
import sys
from contextlib import contextmanager

import cv2

import _stbt.logging
import stbt


def error(s):
    sys.stderr.write("stbt match: error: %s\n" % s)
    sys.exit(1)


parser = argparse.ArgumentParser()
parser.prog = "stbt match"
parser.description = """Run stbt's image-matching algorithm against a single
    frame (which you can capture using `stbt screenshot`)."""
parser.add_argument(
    "-v", "--verbose", action="store_true",
    help="Dump image processing debug images to ./stbt-debug directory")
parser.add_argument(
    "source_file", help="""The screenshot to compare against (you can capture it
        using 'stbt screenshot')""")
parser.add_argument(
    "template_file", help="The image to search for")
parser.add_argument(
    "match_parameters", nargs="*",
    help="""Parameters for the image processing algorithm. See
        'MatchParameters' in the stbt API documentation. For example:
        'confirm_threshold=0.20')""")
args = parser.parse_args(sys.argv[1:])

mp = stbt.MatchParameters()
try:
    for p in args.match_parameters:
        name, value = p.split("=")
        if name == "match_method":
            mp.match_method = value
        elif name == "match_threshold":
            mp.match_threshold = float(value)
        elif name == "confirm_method":
            mp.confirm_method = value
        elif name == "confirm_threshold":
            mp.confirm_threshold = float(value)
        elif name == "erode_passes":
            mp.erode_passes = int(value)
        else:
            raise Exception("Unknown match_parameter argument '%s'" % p)
except Exception:  # pylint: disable=W0703
    error("Invalid argument '%s'" % p)

source_image = cv2.imread(args.source_file)
if source_image is None:
    error("Invalid image '%s'" % args.source_file)


@contextmanager
def noop_contextmanager():
    yield

with (_stbt.logging.scoped_debug_level(2) if args.verbose
      else noop_contextmanager()):
    try:
        result = stbt.match(
            os.path.abspath(args.template_file), frame=source_image,
            match_parameters=mp)
        print "%s: %s" % (
            "Match found" if result else "No match found. Closest match",
            result)
        sys.exit(0 if result else 1)
    except stbt.UITestError as e:
        error(e.message)