def setUp(self):
self.class_config = ClassConfig(names=['a', 'b', 'c'],
colors=['red', 'green', 'blue'])
self.class_config.ensure_null_class()
self.transformer = SegmentationClassTransformer(self.class_config)
self.rgb_image = np.zeros((1, 3, 3))
self.rgb_image[0, 0, :] = color_to_triple('red')
self.rgb_image[0, 1, :] = color_to_triple('green')
self.rgb_image[0, 2, :] = color_to_triple('blue')
self.class_image = np.array([[0, 1, 2]])
def plot_xyz(self, ax, x, y, z=None):
x = x.permute(1, 2, 0)
if x.shape[2] == 1:
x = torch.cat([x for _ in range(3)], dim=2)
ax.imshow(x)
ax.axis('off')
labels = z if z is not None else y
colors = [color_to_triple(c) for c in self.cfg.data.class_colors]
colors = [tuple([_c / 255 for _c in c]) for c in colors]
cmap = matplotlib.colors.ListedColormap(colors)
labels = labels.numpy()
ax.imshow(labels, alpha=0.4, vmin=0, vmax=len(colors), cmap=cmap)
def __init__(self, class_config):
color_to_class = class_config.get_color_to_class_id()
color_int_to_class = dict(
zip([color_to_integer(c) for c in color_to_class.keys()],
color_to_class.values()))
null_class_id = class_config.get_null_class_id()
def color_int_to_class_fn(color: int) -> int:
# Convert unspecified colors to null class
return color_int_to_class.get(color, null_class_id)
self.transform_color_int_to_class = \
np.vectorize(color_int_to_class_fn, otypes=[np.uint8])
# class to color triple
class_to_color_triple = dict(
zip(color_to_class.values(),
[color_to_triple(c) for c in color_to_class.keys()]))
def class_to_channel_color(channel: int, class_id: int) -> int:
"""Given a channel (red, green, or blue) and a class, return the
intensity of that channel.
Args:
channel: An integer with value 0, 1, or 2
representing the channel.
class_id: The class id represented as an integer.
Returns:
The intensity of the channel for the color associated
with the given class.
"""
default_triple = (0x00, 0x00, 0x00)
return class_to_color_triple.get(class_id, default_triple)[channel]
class_to_r = np.vectorize(
lambda c: class_to_channel_color(0, c), otypes=[np.uint8])
class_to_g = np.vectorize(
lambda c: class_to_channel_color(1, c), otypes=[np.uint8])
class_to_b = np.vectorize(
lambda c: class_to_channel_color(2, c), otypes=[np.uint8])
self.transform_class_to_color = [class_to_r, class_to_g, class_to_b]
def update(self, pipeline=None):
if not self.colors:
self.colors = [color_to_triple() for _ in self.names]