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]