diff --git a/overcooked_simulator/gui_2d_vis/drawing.py b/overcooked_simulator/gui_2d_vis/drawing.py index d9e5e8e22b0c3ad16df665e496c759178cd26262..62878e72a6dacb9fe63d3e3edc4c3c2eb0ff85f5 100644 --- a/overcooked_simulator/gui_2d_vis/drawing.py +++ b/overcooked_simulator/gui_2d_vis/drawing.py @@ -399,6 +399,8 @@ class Visualizer: if "center_offset" in part: d = pygame.math.Vector2(part["center_offset"]) * grid_size d.rotate_ip(angle_offset) + d[0] = -d[0] + draw_pos += np.array(d) height = part["height"] * grid_size width = part["width"] * grid_size @@ -805,9 +807,24 @@ class Visualizer: flags = pygame.HIDDEN screen = pygame.display.set_mode((width, height), flags=flags) - self.draw_gamescreen(screen, state, grid_size) + self.draw_gamescreen(screen, state, grid_size, [0 for _ in state["players"]]) pygame.image.save(screen, filename) + def get_state_image(self, grid_size: int, state: dict) -> npt.NDArray[np.uint8]: + width = int(np.ceil(state["kitchen"]["width"] * grid_size)) + height = int(np.ceil(state["kitchen"]["height"] * grid_size)) + + flags = pygame.HIDDEN + screen = pygame.display.set_mode((width, height), flags=flags) + + self.draw_gamescreen(screen, state, grid_size, [0 for _ in state["players"]]) + + red = pygame.surfarray.array_red(screen) + green = pygame.surfarray.array_green(screen) + blue = pygame.surfarray.array_blue(screen) + res = np.stack([red, green, blue], axis=2) + return res + def save_screenshot(state: dict, config: dict, filename: str | Path) -> None: """Standalone function to save a screenshot. Creates a visualizer from the config and visualizes