diff --git a/README.md b/README.md index de175c7..2980a27 100644 --- a/README.md +++ b/README.md @@ -114,6 +114,21 @@ When drawing tasks, arckit will intelligently resize all of the grids such that ![Example of arckit output](./images/arcsave_example.png) +Alternatively, the `print_grid` function outputs a grid directly to the terminal without creating any files. Note that it takes a numpy array as input, which is the standard grid format for a *Task*. + +```python +>>> task = train_set[0] # Get the first task +>>> for i, (input_array, output_array) in enumerate(task.train): # Loop through the training examples +>>> print(f"Training Example {i+1}") +>>> print("Input:") +>>> vis.print_grid(input_array) +>>> print("Output:") +>>> vis.print_grid(output_array) +>>> print() +``` + +![Example of print grid in terminal](./images/arc_terminal_print_grid_example.png) + ## 💻 Command-line tools `arcshow` draws a visualisation of a specific task straight to the console: diff --git a/arckit/vis.py b/arckit/vis.py index 1c3d622..798be69 100644 --- a/arckit/vis.py +++ b/arckit/vis.py @@ -1,6 +1,7 @@ import drawsvg import numpy as np import io +import rich cmap = [ '#252525', # black @@ -262,4 +263,45 @@ def output_drawing(d: drawsvg.Drawing, filename: str, context=None): import cairosvg cairosvg.svg2pdf(bytestring=buffer.getvalue(), write_to=filename) else: - raise ValueError(f'Unknown file extension for {filename}') \ No newline at end of file + raise ValueError(f'Unknown file extension for {filename}') + +def print_grid(grid: np.ndarray): + """ + Print a grid to the terminal using rich library + + Parameters + ---------- + grid : np.ndarray + the standard grid format for Task + """ + CELL_WIDTH = 2 + + def get_color(color_str): + color_str = color_str.strip('#') + return rich.color.Color.from_rgb(*bytes.fromhex(color_str)) + + # Translate 'cmap' to rich style with respective color as background + rich_scmap = { + i: rich.style.Style(bgcolor=get_color(color)) + for i, color in enumerate(cmap) + } + + # Create and populate rich table + table = rich.table.Table.grid(expand=False) + + height, width = grid.shape + + # Add columns + for x in range(width): + table.add_column() + table.columns[x]._cells = [''] * height + + table.rows = [rich.table.Row()] * height + + # Populate cells + for y in range(height): + for x in range(width): + color_idx = grid[y, x] + table.columns[x]._cells[y] = rich.text.Text(' ' * CELL_WIDTH, style=rich_scmap[color_idx]) + + rich.print(table) \ No newline at end of file diff --git a/images/arc_terminal_print_grid_example.png b/images/arc_terminal_print_grid_example.png new file mode 100644 index 0000000..eea54c7 Binary files /dev/null and b/images/arc_terminal_print_grid_example.png differ