traffic-sign-classifier-robustness-testing / rt_search_based / main.py
main.py
Raw
import argparse
import logging
import sys

from rt_search_based.gui.gui import create_app
from rt_search_based.database.database import Database
from rt_search_based.examples import (
    bruteforce_all_example,
    fitness_function_visualization_example,
    jmetalpy_all_example,
    pymoo_all_example,
)


def run_example(strategy_name: str):
    """Runs a premade example strategy"""

    EXAMPLE_STRATEGIES = {
        "bruteforce": bruteforce_all_example,
        "jmetalpy": jmetalpy_all_example,
        "pymoo": pymoo_all_example,
        "fitness-visualization": fitness_function_visualization_example,
    }

    if strategy_name == "all":
        for strategy in EXAMPLE_STRATEGIES.values():
            strategy.run_example()
    else:
        try:
            EXAMPLE_STRATEGIES[strategy_name].run_example()
        except KeyError as key:
            logging.error(
                f"Invalid example strategy chosen: '{', '.join(key.args)}'",
                f"Available examples: {', '.join(EXAMPLE_STRATEGIES.keys())}",
            )


if __name__ == "__main__":
    logging.basicConfig(
        level=logging.INFO,
        handlers=[logging.FileHandler("info.log"), logging.StreamHandler(sys.stdout)],
    )

    parser = argparse.ArgumentParser(
        description=(
            "Tool used for comparing and evaluating different search "
            "strategies and fitness functions on their ability to fool "
            "classifiers with sticker attacks."
        )
    )
    parser.add_argument(
        "--gui",
        type=str,
        nargs="?",
        default=False,
        const=True,
        help="'--gui' Starts the web GUI, '--gui debug' starts the GUI in debug mode",
    )
    parser.add_argument(
        "--example",
        type=str,
        nargs="?",
        help=(
            "'--example <insert_name>' to run an example."
            "Currently available: 'bruteforce', 'jmetalpy', 'pymoo', 'fitness-visualization' or run "
            "all strategies with 'all'"
        ),
    )

    # If main.py is run without arguments print help message
    if len(sys.argv) == 1:
        parser.print_help(sys.stderr)
        sys.exit(1)

    args = parser.parse_args()

    if args.example:
        run_example(args.example)

    if args.gui:
        # start the gui
        with Database() as db:
            gui_app = create_app(db)
            gui_app.run(debug=(args.gui == "debug"))