PuristanLabs1's picture
Update app.py
d11ee3f verified
import json
import os
import numpy as np
import matplotlib.pyplot as plt
import gradio as gr
# Update this to the correct path of your folder
task_folder = "./evaluation" # Use absolute path if necessary
def plot_arc_problem(task_file_path, data_type="train"):
"""
Visualize input-output pairs for a given ARC task file.
Args:
task_file_path (str): Path to the ARC task JSON file.
data_type (str): 'train' or 'test' to visualize respective examples.
Returns:
matplotlib.figure.Figure or str: The plotted figure or an error message.
"""
try:
# Load the JSON file
with open(task_file_path, 'r') as f:
task = json.load(f)
# Get the data type (train or test) pairs
pairs = task.get(data_type, [])
if not pairs: # Check if the section exists and has data
return f"No '{data_type}' data found in the selected file."
# Create a figure with subplots for each pair
fig, axes = plt.subplots(len(pairs), 2, figsize=(10, 5 * len(pairs)))
fig.suptitle(f"ARC Task: {os.path.basename(task_file_path)} ({data_type.capitalize()})", fontsize=16)
# Handle case where there is only one pair
if len(pairs) == 1:
axes = [axes]
for idx, pair in enumerate(pairs):
input_grid = np.array(pair['input'])
output_grid = np.array(pair['output'])
# Plot input grid
axes[idx][0].imshow(input_grid, cmap='tab20', interpolation='none')
axes[idx][0].set_title(f"Input Pair {idx + 1} ({data_type.capitalize()})")
axes[idx][0].axis('off')
# Plot output grid
axes[idx][1].imshow(output_grid, cmap='tab20', interpolation='none')
axes[idx][1].set_title(f"Output Pair {idx + 1} ({data_type.capitalize()})")
axes[idx][1].axis('off')
plt.tight_layout(rect=[0, 0, 1, 0.96])
return fig
except Exception as e:
print("Error in plot_arc_problem:", e)
return f"An error occurred while plotting the {data_type} data: {e}"
def visualize_task(file_name, data_type):
"""
Load and visualize the ARC task for a given file and data type.
Args:
file_name (str): Name of the JSON file in the folder.
data_type (str): 'train' or 'test'.
Returns:
matplotlib.figure.Figure or str: Figure of the visualized task or an error message.
"""
try:
print(f"Selected file: {file_name}, Data type: {data_type}") # Debugging
task_file_path = os.path.join(task_folder, file_name)
result = plot_arc_problem(task_file_path, data_type)
return result
except Exception as e:
print("Error in visualize_task:", e)
return f"An error occurred while visualizing the task: {e}"
# Gradio Interface
task_files = [f for f in os.listdir(task_folder) if f.endswith('.json')]
interface = gr.Interface(
fn=visualize_task,
inputs=[
gr.Dropdown(choices=task_files, label="Select ARC Task File"),
gr.Radio(choices=["train", "test"], label="Select Data Type to Visualize", value="train")
],
outputs="plot",
title="ARC Task Visualizer",
description="Select a task file and data type (train or test) to visualize its input-output grids."
)
if __name__ == "__main__":
print("Task files:", task_files) # Debugging
interface.launch()