Upload 10 files

app.py

@@ -1,26 +1,92 @@
 import gradio as gr
 import numpy as np
-import matplotlib.pyplot as plt
+import plotly.graph_objects as go
 
-from src.boards import GenerateBoard
+from src.utils import select_board
 
 TITLE = """<h1 align="center" id="space-title"> Pento-LLaVA 🤖🎯🎮</h1>"""
 
-initial_board_image, target_positions, info = GenerateBoard('easy', 18).setup_initial_board() 
+pento_llava_app = gr.Blocks()
+fig, targets, info = select_board('easy', 18, 0)
 
-# Convert initial_board_image to a matplotlib figure
-fig, ax = plt.subplots()
-ax.imshow(initial_board_image)
-ax.axis('off')
+# Update the figure to include a boundary/frame
+fig.update_layout(
+    xaxis=dict(showline=True, linecolor='black', linewidth=2, showgrid=True, zeroline=True),
+    yaxis=dict(showline=True, linecolor='black', linewidth=2, showgrid=True, zeroline=True),
+    margin=dict(l=0, r=0, t=0, b=0),
+    height=512,
+    width=512
+)
+
+target_strs = []
+for t in targets:
+    target_strs.append(t['target_str'])
+
+def gen_new_board(value):
+    value = int(value)
+    fig, _, _ = select_board('easy', 18, value)
+
+    # Update the figure to include a boundary/frame
+    fig.update_layout(
+        xaxis=dict(showline=True, linecolor='black', linewidth=2, showgrid=True, zeroline=True),
+        yaxis=dict(showline=True, linecolor='black', linewidth=2, showgrid=True, zeroline=True),
+        margin=dict(l=0, r=0, t=0, b=0),
+        height=512,
+        width=512
+    )
+
+    return fig
+
+def gen_info(value):
+    value = int(value)
+    _, _, info = select_board('easy', 18, value)
+
+    return info
+
+def gen_target_str(value):
+    value = int(value)
+    _, _, info = select_board('easy', 18, value)
+    target_info = info[0]
+    target_str = f"<span>Target piece for this episode is <span style='color: {target_info['piece_colour']};'>{target_info['piece_colour']}</span> <span style='color: {target_info['piece_colour']};'>{target_info['piece_shape']}</span> located at <span style='color: {target_info['piece_colour']};'>{target_info['piece_region']}</span></span>"
+
+    return target_str
 
-pento_llava_app = gr.Blocks()
 
 with pento_llava_app:
 
     gr.HTML(TITLE)
-    gr.Plot(fig)
+
+    with gr.Row():
+        with gr.Column():    
+            main_board = gr.Plot(fig)
+
+        with gr.Column():
+            with gr.Row():
+                select_board_items = gr.Dropdown(
+                    choices=range(512),
+                    interactive=True
+                )
+
+            with gr.Row():
+                 display_string = gr.HTML(value=gen_target_str(0)) 
+        
+
+            select_board_items.change(
+                fn=gen_new_board,
+                inputs=[select_board_items],
+                outputs=[main_board],
+                queue=True
+            )
+
+            select_board_items.change(
+                fn=gen_target_str,
+                inputs=[select_board_items],
+                outputs=[display_string],
+                queue=True
+            )
+
 
     pento_llava_app.load()
 
 pento_llava_app.queue()
-pento_llava_app.launch()
+pento_llava_app.launch()

requirements.txt

@@ -14,4 +14,5 @@ accelerate==0.31.0
 bitsandbytes
 datasets==3.0.1
 gradio==5.5.0
-matplotlib==3.9.2
+matplotlib==3.9.2
+plotly==5.24.1

src/boards.py

@@ -7,7 +7,7 @@ import numpy as np
 
 class GenerateBoard():
 
-    def __init__(self, level: str, board_size: int):
+    def __init__(self, level: str, board_size: int, board_number: int):
         self.level = level
         self.board_size = board_size
 
@@ -16,9 +16,8 @@ class GenerateBoard():
         with open(metadata_path, 'r') as f:
             metadata = json.load(f)
 
-        num_boards = len(metadata)
-        random_board_num = np.random.randint(0, num_boards)
-        self.board_data = metadata[random_board_num]
+        self.board_num = board_number
+        self.board_data = metadata[self.board_num]
     
     def setup_initial_board(self):
 
@@ -29,10 +28,11 @@ class GenerateBoard():
         info = metadata_obj['info']
         target_options = []
         for piece in info:
+            target_info_dict = {}
             target = f"{piece['piece_colour']} {piece['piece_shape']} at {piece['piece_region']}" 
-            target_options.append(target)
-
-
+            target_info_dict['target_str'] = target
+            target_info_dict['piece_info'] = piece
+            target_options.append(target_info_dict)
         
         env = GridWorldEnv(render_mode="rgb_array", size=self.board_size, grid_info=info, agent_pos=default_start_pos, target_pos=default_target_pos)
         env.reset()

src/utils.py

@@ -0,0 +1,12 @@
+from src.boards import GenerateBoard
+import plotly.express as px  # Import Plotly Express
+
+def select_board(level: str = 'easy', size: int = 18, board_number: int = 0):
+    initial_board_image, target_positions, info = GenerateBoard(level, size, board_number).setup_initial_board() 
+
+    # Convert initial_board_image to a Plotly figure
+    fig = px.imshow(initial_board_image)  # Use Plotly's imshow
+    fig.update_xaxes(showticklabels=False)  # Hide x-axis ticks
+    fig.update_yaxes(showticklabels=False)  # Hide y-axis ticks
+
+    return fig, target_positions, info