def main()

in model/utils/visualize_bboxes.py [0:0]

• 27 lines of code
• 8 McCabe index (conditional complexity)

def main(args):
    if args["scene_names"] == "all":
        # Get all the scene names in the folder.
        scene_names = [
            ii.rsplit("_", 1)[0] for ii in os.listdir(args["scene_json_root"])
        ]
        scene_names = list(set(scene_names))
        # Remove explore.py.
        scene_names.remove("explore.py")
    else:
        scene_names = args["scene_names"]

    print(f"""Reading scene JSONS: {args["scene_json_root"]}""")
    print(f"""Reading scene screenshots: {args["screenshot_root"]}""")
    for scene in progressbar(scene_names):
        json_path = os.path.join(args["scene_json_root"], f"{scene}_scene.json")
        # Check if file exists, else try with "m_"
        if not os.path.exists(json_path):
            json_path = os.path.join(args["scene_json_root"], f"m_{scene}_scene.json")
            assert os.path.exists(json_path), f"{json_path} not found!"
        with open(json_path, "r") as file_id:
            scene_json = json.load(file_id)
        object_bboxes = scene_json["scenes"][0]["objects"]

        # Image load and save paths.
        trimmed_scene_name = scene[2:] if scene[:2] == "m_" else scene
        screenshot_load_path = os.path.join(
            args["screenshot_root"], f"{trimmed_scene_name}.png"
        )
        screenshot_save_path = os.path.join(
            args["save_root"], f"{trimmed_scene_name}_bbox.png"
        )
        draw_bboxes(object_bboxes, screenshot_load_path, screenshot_save_path)