def after_step(self, batch_index, results):
step = retrieve('global_step', results)
eval_dir = os.path.join(self.root, '{:6>0}'.format(step))
os.makedirs(eval_dir, exist_ok=True)
for key in self.keys:
if self.names is not None:
name = self.names[i]
else:
name = key.split('/')[-1]
batched_sequences = retrieve(key, results)
if key.split('/')[-1] == 'images':
batched_sequences = batched_sequences[1:]
iterator = zip(self.indices,
self.boxes,
self.paths,
batched_sequences)
logger.info('seq {}'.format(np.shape(batched_sequences)))
for frame, [idx, b, fid, image] in enumerate(iterator):
logger.info('fid {}'.format(fid))
logger.info('b {}'.format(b))
logger.info('seq {}'.format(image.shape))
savename = os.path.join(eval_dir,
'{:0>7}_{}-{:0>3}.png'.format(idx, name, frame))
save_image(image, savename)
np.save(savename.replace('.png', '-box.npy'), b[frame])
with open(savename.replace('.png', '-org.txt'), 'w+') as f:
f.write(fid[0])
def write_slice(t, out_folder="."):
"""
t = (i, image_slice)
image_slice is [H, W, C]
"""
i, image_slice = t
fname = "{:06d}.png".format(i)
batches.save_image(image_slice, os.path.join(out_folder, fname))
def make_clusters(data: dict, root: str, config: dict,
global_step: int) -> None:
matching_app_features_0 = data["outputs"]["matching_app_features_0"]
N, C, P, F = matching_app_features_0.shape
matching_app_features_0 = np.reshape(
np.rollaxis(matching_app_features_0, 2, 1), (N * P, C, F))
app_features_list = list(
map(np.squeeze, np.split(matching_app_features_0, C, 1)))
k = config.get("num_clusters", 2)
n_vis = config.get("n_vis", 20)
func_ = functools.partial(cluster_features, **{"k": k})
centroids_and_labels = list(map(func_, app_features_list))
new_data = {}
new_data["clusters"] = np.stack([c[0] for c in centroids_and_labels])
new_data["labels"] = np.stack([c[1] for c in centroids_and_labels])
out_path = os.path.join(root, "clusters.p")
with open(out_path, "wb") as f:
pickle.dump(new_data, f)
out_path = os.path.join(root, "figure_02")
os.makedirs(out_path, exist_ok=True)
decoding_mask4 = data["outputs"]["decoding_mask4"][:n_vis]
view0 = data["inputs"]["view0"][:n_vis]
for i in range(C):
_, labels = centroids_and_labels[i]
labels = labels[:n_vis]
j = 0
for canvas in yield_visualize_clusters(
np.squeeze(decoding_mask4[..., i]), view0, labels, n_vis):
save_image(
canvas,
os.path.join(out_path, "{:06d}_cluster{:02d}.png".format(i,
j)))
j += 1