finish image rendering

Author: BartekCupial

balrog/environments/battleships/base.py

@@ -4,6 +4,7 @@
 import gym
 import numpy as np
 import pandas as pd
+from PIL import Image, ImageDraw, ImageFont
 from scipy.ndimage import label
 
 
@@ -38,17 +39,16 @@ def get_feedback(self, reward, old_reward):
             return "MISS! Your missile splashed into empty water."
 
     def battleships_process_obsv(self, obs, reward, old_reward):
-        dataframe = self.get_dataframe(obs)
+        text_observation = self.get_text_observation(obs)
+        image = self.get_image_observation(obs)
 
-        text_observation = self.get_text_observation(dataframe)
         feedback = self.get_feedback(reward, old_reward)
 
         prompt = f"{text_observation}\n{feedback}" if feedback else f"Objects on the map:\n{text_observation}"
 
         obs = defaultdict(lambda: None)
 
         obs["text"] = {"long_term_context": prompt, "short_term_context": ""}
-        image = self.get_image_observation(dataframe)
         obs["image"] = image
 
         return obs
@@ -102,7 +102,7 @@ def step(self, action):
     def get_stats(self):
         return {"progression": self.progression}
 
-    def get_dataframe(self, obs):
+    def get_text_observation(self, obs):
         board = np.empty(self.env.board_size, dtype=str)
 
         # Create a mask for sunk ships
@@ -121,69 +121,97 @@ def get_dataframe(self, obs):
         dataframe = pd.DataFrame(board, columns=columns, index=index)
         dataframe = dataframe.replace([""], " ")
 
-        return dataframe
-
-    def get_text_observation(self, dataframe):
         obsv = str(dataframe)
 
         return obsv
 
-    def get_image_observation(self, dataframe):
-        # import matplotlib.pyplot as plt
-        # from matplotlib.colors import LinearSegmentedColormap
-        # from matplotlib.figure import Figure
-        # from matplotlib.backends.backend_agg import FigureCanvasAgg
-        # import io
-        # from PIL import Image
-
-        # # Define colors for each cell type
-        # color_map = {
-        #     "⬜": [0.9, 0.9, 1.0],    # Light blue for empty water
-        #     "❌": [1.0, 0.0, 0.0],    # Red for hits
-        #     "⚫": [0.3, 0.3, 0.3],    # Dark gray for misses
-        #     "💥": [1.0, 0.6, 0.0]     # Orange for sunk ships
-        # }
-
-        # # Create a numerical representation for colormapping
-        # numeric_board = np.zeros(dataframe.shape + (3,), dtype=float)
-
-        # for i in range(dataframe.shape[0]):
-        #     for j in range(dataframe.shape[1]):
-        #         cell_value = dataframe.iloc[i, j]
-        #         numeric_board[i, j] = color_map.get(cell_value, [1, 1, 1])
-
-        # # Create a figure with the right dimensions and no padding
-        # fig_width = dataframe.shape[1] + 1  # +1 for row labels
-        # fig_height = dataframe.shape[0] + 1  # +1 for column labels
-        # fig = Figure(figsize=(fig_width, fig_height), dpi=72)
-        # canvas = FigureCanvasAgg(fig)
-        # ax = fig.add_subplot(111)
-
-        # # Plot the board
-        # ax.imshow(numeric_board, aspect='equal')
-
-        # # Add grid lines
-        # ax.set_xticks(np.arange(-0.5, dataframe.shape[1], 1), minor=True)
-        # ax.set_yticks(np.arange(-0.5, dataframe.shape[0], 1), minor=True)
-        # ax.grid(which='minor', color='black', linestyle='-', linewidth=1)
-
-        # # Add column labels (A, B, C, ...)
-        # ax.set_xticks(np.arange(dataframe.shape[1]))
-        # ax.set_xticklabels(dataframe.columns)
-
-        # # Add row labels (1, 2, 3, ...)
-        # ax.set_yticks(np.arange(dataframe.shape[0]))
-        # ax.set_yticklabels(dataframe.index)
-
-        # # Remove axis padding
-        # ax.set_xlim(-0.5, dataframe.shape[1] - 0.5)
-        # ax.set_ylim(-0.5, dataframe.shape[0] - 0.5)
-
-        # # Render the figure to a numpy array
-        # canvas.draw()
-        # buf = io.BytesIO()
-        # fig.savefig(buf, format='png', bbox_inches='tight', pad_inches=0.1)
-        # buf.seek(0)
-
-        # return Image.open(buf)
-        return None
+    def get_image_observation(self, obs, symbol=" ", cell_size=80, font_size=40, border_width=2):
+        board = np.empty(self.env.board_size, dtype=str)
+        board[obs[0] != 0] = "X"
+        board[obs[1] != 0] = "O"
+
+        num_rows, num_columns = board.shape
+        columns = [chr(i) for i in range(ord("A"), ord("A") + num_columns)]
+        index = [str(i + 1) for i in range(num_rows)]
+
+        # Calculate image dimensions with space for row/column labels
+        header_size = cell_size // 2
+        width = (num_columns * cell_size) + header_size
+        height = (num_rows * cell_size) + header_size
+
+        # Create image with white background
+        image = Image.new("RGB", (width, height), color="white")
+        draw = ImageDraw.Draw(image)
+
+        try:
+            # Try to load a font that supports Unicode symbols
+            font = ImageFont.truetype("Arial Unicode MS", font_size)
+        except IOError:
+            try:
+                # Try another common font
+                font = ImageFont.truetype("DejaVuSans.ttf", font_size)
+            except IOError:
+                # Fallback to default font
+                font = ImageFont.load_default()
+
+        # Draw column headers (A, B, C, ...)
+        for col_idx, col in enumerate(columns):
+            x = header_size + (col_idx * cell_size) + (cell_size // 2)
+            y = header_size // 2
+            draw.text((x, y), col, fill="black", font=font, anchor="mm")
+
+        # Draw row headers (1, 2, 3, ...)
+        for row_idx, row in enumerate(index):
+            x = header_size // 2
+            y = header_size + (row_idx * cell_size) + (cell_size // 2)
+            draw.text((x, y), row, fill="black", font=font, anchor="mm")
+
+        # Draw grid
+        for row_idx in range(num_rows + 1):
+            y = header_size + (row_idx * cell_size)
+            draw.line([(header_size, y), (width, y)], fill="black", width=border_width)
+
+        for col_idx in range(num_columns + 1):
+            x = header_size + (col_idx * cell_size)
+            draw.line([(x, header_size), (x, height)], fill="black", width=border_width)
+
+        # Draw cell contents
+        for row_idx in range(num_rows):
+            for col_idx in range(num_columns):
+                cell_content = board[row_idx, col_idx]
+                x0 = header_size + (col_idx * cell_size) + 5
+                y0 = header_size + (row_idx * cell_size) + 5
+                x1 = header_size + ((col_idx + 1) * cell_size) - 5
+                y1 = header_size + ((row_idx + 1) * cell_size) - 5
+
+                # Center of the cell for drawing shapes
+                center_x = header_size + (col_idx * cell_size) + (cell_size // 2)
+                center_y = header_size + (row_idx * cell_size) + (cell_size // 2)
+                radius = (cell_size // 2) - 10
+                radius = int(radius * 0.7)
+                # Draw based on cell content
+                if cell_content == "X":  # Hit
+                    # Draw a red X
+                    draw.line(
+                        [(center_x - radius, center_y - radius), (center_x + radius, center_y + radius)],
+                        fill="red",
+                        width=10,
+                    )
+                    draw.line(
+                        [(center_x + radius, center_y - radius), (center_x - radius, center_y + radius)],
+                        fill="red",
+                        width=10,
+                    )
+                elif cell_content == "O":  # Miss
+                    # Draw a blue circle
+                    radius = 5
+                    draw.ellipse(
+                        [(center_x - radius, center_y - radius), (center_x + radius, center_y + radius)],
+                        outline="black",
+                        width=5,
+                    )
+                else:  # Empty
+                    # Draw a light white square
+                    draw.rectangle([x0, y0, x1, y1], fill="white")
+
+        return image