solve_set.py

#!/usr/bin/env python
"""Solve SET from a game image. Prints the sets found, and optionally
writes and/or displays the game image with the sets indicated.
"""

import argparse
import sys
from SetGame import SetGame
from common import game_img_filename


def get_args():
    parser = argparse.ArgumentParser(description="Solve SET from a game image.")
    parser.add_argument(
        "filename", metavar="filename", type=str, nargs="?", help="Game image filename"
    )
    parser.add_argument(
        "--game",
        dest="game_num",
        type=int,
        help="use a test image from image-data/set-games/setgame<GAME_NUM>.jpg",
    )
    parser.add_argument(
        "--write",
        dest="write",
        action="store_true",
        help="Write the solved image to solve-out/solved.jpg",
    )
    parser.add_argument(
        "--display",
        dest="display",
        action="store_true",
        help="Display the solved image with cv2.display()",
    )
    return parser.parse_args()


def main():
    args = get_args()

    if args.game_num:
        filename = game_img_filename(args.game_num)
    elif args.filename:
        filename = args.filename
    else:
        print("Error: must pass either a filename or --game GAME_NUM")
        sys.exit(1)

    game = SetGame(filename)
    game.solve()
    game.print_sets()
    game.draw_sets()

    if args.write:
        game.write_im("solved.jpg")
    if args.display:
        game.display()


if __name__ == "__main__":
    main()