Spaces:
Sleeping
Sleeping
File size: 1,590 Bytes
55ecc31 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
"""
Visualisation utils.
"""
import chess
import chess.svg
import matplotlib
import matplotlib.pyplot as plt
COLOR_MAP = matplotlib.colormaps["RdYlBu_r"].resampled(1000)
ALPHA = 1.0
def render_heatmap(
board,
heatmap,
square=None,
vmin=None,
vmax=None,
arrows=None,
normalise="none",
):
"""
Render a heatmap on the board.
"""
if normalise == "abs":
a_max = heatmap.abs().max()
if a_max != 0:
heatmap = heatmap / a_max
vmin = -1
vmax = 1
if vmin is None:
vmin = heatmap.min()
if vmax is None:
vmax = heatmap.max()
norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax, clip=False)
color_dict = {}
for square_index in range(64):
color = COLOR_MAP(norm(heatmap[square_index]))
color = (*color[:3], ALPHA)
color_dict[square_index] = matplotlib.colors.to_hex(color, keep_alpha=True)
fig = plt.figure(figsize=(6, 0.6))
ax = plt.gca()
ax.axis("off")
fig.colorbar(
matplotlib.cm.ScalarMappable(norm=norm, cmap=COLOR_MAP),
ax=ax,
orientation="horizontal",
fraction=1.0,
)
if square is not None:
try:
check = chess.parse_square(square)
except ValueError:
check = None
else:
check = None
if arrows is None:
arrows = []
plt.close()
return (
chess.svg.board(
board,
check=check,
fill=color_dict,
size=350,
arrows=arrows,
),
fig,
)
|