feat: add color utils

This commit is contained in:
Jensun Ravichandran 2022-03-10 14:45:55 +01:00
parent 695559fd4a
commit 236cbbc4d2
No known key found for this signature in database
GPG Key ID: 4E9348239810B51F
2 changed files with 52 additions and 2 deletions

View File

@ -1,6 +1,11 @@
"""ProtoFlow utils module""" """ProtoTorch utils module"""
from .colors import hex_to_rgb, rgb_to_hex from .colors import (
get_colors,
get_legend_handles,
hex_to_rgb,
rgb_to_hex,
)
from .utils import ( from .utils import (
mesh2d, mesh2d,
parse_data_arg, parse_data_arg,

View File

@ -1,5 +1,14 @@
"""ProtoTorch color utilities""" """ProtoTorch color utilities"""
import matplotlib.lines as mlines
import torch
from matplotlib import cm
from matplotlib.colors import (
Normalize,
to_hex,
to_rgb,
)
def hex_to_rgb(hex_values): def hex_to_rgb(hex_values):
for v in hex_values: for v in hex_values:
@ -13,3 +22,39 @@ def rgb_to_hex(rgb_values):
for v in rgb_values: for v in rgb_values:
c = "%02x%02x%02x" % tuple(v) c = "%02x%02x%02x" % tuple(v)
yield c yield c
def get_colors(vmax, vmin=0, cmap="viridis"):
cmap = cm.get_cmap(cmap)
colornorm = Normalize(vmin=vmin, vmax=vmax)
colors = dict()
for c in range(vmin, vmax + 1):
colors[c] = to_hex(cmap(colornorm(c)))
return colors
def get_legend_handles(colors, labels, marker="dots", zero_indexed=False):
handles = list()
for color, label in zip(colors.values(), labels):
if marker == "dots":
handle = mlines.Line2D(
xdata=[],
ydata=[],
label=label,
color="white",
markerfacecolor=color,
marker="o",
markersize=10,
markeredgecolor="k",
)
else:
handle = mlines.Line2D(
xdata=[],
ydata=[],
label=label,
color=color,
marker="",
markersize=15,
)
handles.append(handle)
return handles