def after_step(self, batch_index, last_results):
if batch_index % self.interval == 0:
step = last_results["global_step"]
# TODO: fix hard-coded font type
# TODO: add option to log overview to tensorboard
batches = []
try:
fnt = ImageFont.truetype("LiberationMono-Regular.ttf", 20)
except OSError:
fnt = ImageFont.load_default()
last_results = last_results["logging"]
for name, im in sorted(last_results["images"].items()):
canvas = batch_to_canvas(im)
canvas = (canvas + 1.0) / 2.0
canvas = np.clip(255 * canvas, 0, 255)
canvas = np.array(canvas, dtype="uint8")
im = Image.fromarray(canvas)
im.thumbnail((512, 512), Image.ANTIALIAS)
d = ImageDraw.Draw(im)
d.text((10, 10), name, fill=(255, 0, 0), font=fnt)
batches.append(im)
im = Image.new("RGB", batches[0].size, color=(0, 0, 0))
try:
fnt = ImageFont.truetype("LiberationMono-Regular.ttf", 50)
except OSError:
fnt = ImageFont.load_default()
d = ImageDraw.Draw(im)
d.text((10, 10), "epoch\n{:07d}".format(step), fill=(255, 0, 0), font=fnt)
batches.append(im)
batch = np.stack(batches, axis=0) / 255.0 * 2 - 1.0
out_path = os.path.join(self.root, "overview_{:07d}.png".format(step))
plot_batch(batch, out_path)
def test_batch_to_canvas():
x = np.ones((9, 100, 100, 3))
canvas = batches.batch_to_canvas(x)
assert canvas.shape == (300, 300, 3)
canvas = batches.batch_to_canvas(x, cols=5)
assert canvas.shape == (200, 500, 3)
canvas = batches.batch_to_canvas(x, cols=1)
assert canvas.shape == (900, 100, 3)
canvas = batches.batch_to_canvas(x, cols=0)
assert canvas.shape == (900, 100, 3)
canvas = batches.batch_to_canvas(x, cols=None)
assert canvas.shape == (300, 300, 3)
def yield_visualize_clusters(mask: np.ndarray,
view: np.ndarray,
labels: np.ndarray,
N_vis: int = 10) -> np.ndarray:
"""
Takes masks, views and assigned cluster labels and creates an image per cluster with the
N_vis assigned images per cluster.
Parameters
----------
mask : np.ndarray
[N, H, W]-shaped array where each element is an integer giving the assigned connected component
view : np.ndarray
[N, H, W, 3]-shaped array representing the RGB image
labels : np.array
[N]-shaped array with assignments for each item in features to one of the k-clusters.
Returns
-------
canvas_list : list
a list of canvas plots. Each canvas represents one clusters.
"""
uniques = sorted(np.unique(labels))
for u in uniques:
m = mask[labels == u]
v = view[labels == u]
mask_one_hot = tf.one_hot(m, 5, axis=-1)[:, :, :, 1:] # [N, H, W, 5]
components = np.expand_dims(np.array(mask_one_hot),
-1) * np.expand_dims(v, 3)
N, H, W, P, _ = components.shape
components = np.rollaxis(components, 3, 1)
components = np.reshape(components, (N * P, H, W, 3))
canvas = batch_to_canvas(components)
yield canvas
def log_tensorboard_images(writer, results, step, path):
results = dict((path + "/" + k, v) for k, v in results.items())
for k, v in results.items():
v = batch_to_canvas(v)
v = ((v + 1) * 127.5).astype(np.uint8)
writer.add_image(tag=k,
img_tensor=v,
global_step=step,
dataformats="HWC")
writer.flush()
def make_legend_image(legend_labels, colors, text_colors, legend_size, n_cols):
"""
legend_labels : list of str
list of semantic labels (str) for each label
colors : ndarray or list
list of same length as legend_labels with colors for each label OR ndarray of shape [n_labels, 3]. Colors in range [0, 1]
text_colors : ndarray or list
list of same length as legend_labels with colors for each text OR ndarray of shape [n_labels, 3]. Colors in range [0, 1]
legend_size : tuple of ints
the size of the legend image. the width should match the label_image because it is stacked vertically
n_cols : int
into how many colums the legend labels should be splitted for visualization
# TODO: example usage
"""
N_labels = len(legend_labels)
box_h = 512 // len(legend_labels)
box_w = 512 // n_cols
legend = [
colors[np.ones((box_h, box_w), dtype=np.int16) * i]
for i in range(N_labels)
]
legend = [
imageutils.put_text(patch,
t,
loc="center",
font_scale=1.0,
thickness=1,
color=tc)
for patch, t, tc in zip(legend, legend_labels, text_colors)
]
legend = np.stack(legend, axis=0)
legend = batches.batch_to_canvas(legend, cols=n_cols)
legend = cv2.resize(legend, (legend_size[1], legend_size[0]),
cv2.INTER_LANCZOS4)
return legend
def log_wandb_images(results, step, path):
results = dict((
path + "/" + k,
wandb.Image(((batch_to_canvas(v) + 1) * 127.5).astype(np.uint8)),
) for k, v in results.items())
wandb.log(results, step=step)