diff --git a/overcooked_simulator/gui_2d_vis/drawing.py b/overcooked_simulator/gui_2d_vis/drawing.py index ec72975b4c49d8c0bb29a1036edcc7311eedb7f9..d2f44c42de154e0f98b59bb02c572518422c7a63 100644 --- a/overcooked_simulator/gui_2d_vis/drawing.py +++ b/overcooked_simulator/gui_2d_vis/drawing.py @@ -1,3 +1,4 @@ +import argparse import colorsys import json import math @@ -523,9 +524,32 @@ def save_screenshot(state: dict, config: dict, filename: str | Path) -> None: if __name__ == "__main__": - with open(ROOT_DIR / "gui_2d_vis" / "visualization.yaml", "r") as conf_file: - viz_config = yaml.safe_load(conf_file) - with open(ROOT_DIR / "gui_2d_vis" / "sample_state.json", "r") as jsonf: - sample_state = json.load(jsonf) - screenshot_filename = ROOT_DIR / "gui_2d_vis" / "screenshot.jpg" - save_screenshot(sample_state, viz_config, screenshot_filename) + parser = argparse.ArgumentParser( + prog="Overcooked Simulator Image Generation", + description="Generate images for a state in json.", + epilog="For further information, see https://scs.pages.ub.uni-bielefeld.de/cocosy/overcooked-simulator/overcooked_simulator.html", + ) + parser.add_argument( + "-s", + "--state", + type=argparse.FileType("r", encoding="UTF-8"), + default=ROOT_DIR / "gui_2d_vis" / "sample_state.json", + ) + parser.add_argument( + "-v", + "--visualization_config", + type=argparse.FileType("r", encoding="UTF-8"), + default=ROOT_DIR / "gui_2d_vis" / "visualization.yaml", + ) + parser.add_argument( + "-o", + "--output_filename", + type=str, + default="screenshot.jpg", + ) + args = parser.parse_args() + with open(args.visualization_config, "r") as f: + viz_config = yaml.safe_load(f) + with open(args.state, "r") as f: + state = json.load(f) + save_screenshot(state, viz_config, args.output_filename)