From 67fed26e3e632265473f417b87565c39f2cd15ff Mon Sep 17 00:00:00 2001
From: fheinrich <fheinrich@techfak.uni-bielefeld.de>
Date: Wed, 7 Feb 2024 14:54:27 +0100
Subject: [PATCH] Get state image as np array
---
overcooked_simulator/gui_2d_vis/drawing.py | 19 ++++++++++++++++++-
1 file changed, 18 insertions(+), 1 deletion(-)
diff --git a/overcooked_simulator/gui_2d_vis/drawing.py b/overcooked_simulator/gui_2d_vis/drawing.py
index d9e5e8e2..62878e72 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
--
GitLab