Get state image as np array

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