#!/usr/bin/env python
# coding=utf-8
from __future__ import (print_function, division, absolute_import, unicode_literals)
import os
import numpy as np
import tensorflow as tf
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import hsv_to_rgb
from sklearn.metrics import adjusted_mutual_info_score
def save_image(filename, image_array):
import scipy.misc
if image_array.shape[2] == 1:
if np.min(image_array) >= 0.:
scipy.misc.toimage(image_array[:, :, 0], cmin=0.0, cmax=1.0).save(filename)
else:
scipy.misc.toimage(image_array[:, :, 0], cmin=-1.0, cmax=1.0).save(filename)
else:
scipy.misc.toimage(255*image_array).save(filename)
def delete_files(folder, recursive=False):
for the_file in os.listdir(folder):
file_path = os.path.join(folder, the_file)
try:
if os.path.isfile(file_path):
os.unlink(file_path)
elif recursive and os.path.isdir(file_path):
delete_files(file_path, recursive)
os.unlink(file_path)
except Exception as e:
print(e)
def create_directory(directory):
if not os.path.exists(directory):
os.makedirs(directory)
def print_vars(_vars):
total_n_vars = 0
for var in _vars:
sh = var.get_shape().as_list()
total_n_vars += np.prod(sh)
print(var.name, sh)
print(total_n_vars, 'total variables')
return total_n_vars
ACTIVATION_FUNCTIONS = {
'sigmoid': tf.nn.sigmoid,
'tanh': tf.nn.tanh,
'relu': tf.nn.relu,
'elu': tf.nn.elu,
'linear': lambda x: x,
'exp': lambda x: tf.exp(x),
'softplus': tf.nn.softplus,
'clip': lambda x: tf.clip_by_value(x, -1., 1.),
'clip_low': lambda x: tf.clip_by_value(x, -1., 1e6)
}
def parse_activation_function(name_list):
return [ACTIVATION_FUNCTIONS[name] for name in name_list]
def evaluate_groups_seq(true_groups, predicted, weights):
""" Compute the weighted AMI score and corresponding mean confidence for given gammas.
:param true_groups: (T, B, 1, W, H, 1)
:param predicted: (T, B, K, W, H, 1)
:param weights: (T)
:return: scores, confidences (B,)
"""
w_scores, w_confidences = 0., 0.
assert true_groups.ndim == predicted.ndim == 6, true_groups.shape
for t in range(true_groups.shape[0]):
scores, confidences = evaluate_groups(true_groups[t], predicted[t])
w_scores += weights[t] * np.array(scores)
w_confidences += weights[t] * np.array(confidences)
norm = np.sum(weights)
return w_scores/norm, w_confidences/norm
def evaluate_groups(true_groups, predicted):
""" Compute the AMI score and corresponding mean confidence for given gammas.
:param true_groups: (B, 1, W, H, 1)
:param predicted: (B, K, W, H, 1)
:return: scores, confidences (B,)
"""
scores, confidences = [], []
assert true_groups.ndim == predicted.ndim == 5, true_groups.shape
batch_size, K = predicted.shape[:2]
true_groups = true_groups.reshape(batch_size, -1)
predicted = predicted.reshape(batch_size, K, -1)
predicted_groups = predicted.argmax(1)
predicted_conf = predicted.max(1)
for i in range(batch_size):
true_group = true_groups[i]
idxs = np.where(true_group != 0.0)[0]
scores.append(adjusted_mutual_info_score(true_group[idxs], predicted_groups[i, idxs]))
confidences.append(np.mean(predicted_conf[i, idxs]))
return scores, confidences
def tf_adjusted_rand_index(groups, gammas, iter_weights):
"""
Inputs:
groups: shape=(T, B, 1, W, H, 1)
These are the masks as stored in the hdf5 files
gammas: shape=(T, B, K, W, H, 1)
These are the gammas as predicted by the network
"""
with tf.name_scope('ARI'):
# ignore first iteration
groups = groups[1:]
gammas = gammas[1:]
# reshape gammas and convert to one-hot
yshape = tf.shape(gammas)
gammas = tf.reshape(gammas, shape=tf.stack([yshape[0] * yshape[1], yshape[2], yshape[3] * yshape[4] * yshape[5]]))
Y = tf.one_hot(tf.argmax(gammas, axis=1), depth=yshape[2], axis=1)
# reshape masks
gshape = tf.shape(groups)
groups = tf.reshape(groups, shape=tf.stack([gshape[0] * gshape[1], 1, gshape[3] * gshape[4] * gshape[5]]))
G = tf.one_hot(tf.cast(groups[:, 0], tf.int32), depth=tf.cast(tf.reduce_max(groups) + 1, tf.int32), axis=1)
# now Y and G both have dim (B*T, K, N) where N=W*H*C
# mask entries with group 0
M = tf.cast(tf.greater(groups, 0.5), tf.float32)
n = tf.cast(tf.reduce_sum(M, axis=[1, 2]), tf.float32)
DM = G * M
YM = Y * M
# contingency table for overlap between G and Y
nij = tf.einsum('bij,bkj->bki', YM, DM)
a = tf.reduce_sum(nij, axis=1)
b = tf.reduce_sum(nij, axis=2)
# rand index
rindex = tf.cast(tf.reduce_sum(nij * (nij-1), axis=[1, 2]), tf.float32)
aindex = tf.cast(tf.reduce_sum(a * (a-1), axis=1), tf.float32)
bindex = tf.cast(tf.reduce_sum(b * (b-1), axis=1), tf.float32)
expected_rindex = aindex * bindex / (n*(n-1) + 1e-6)
max_rindex = (aindex + bindex) / 2
ARI = (rindex - expected_rindex)/tf.clip_by_value(max_rindex - expected_rindex, 1e-6, 1e6)
ARI = tf.reshape(ARI, shape=(yshape[0], yshape[1]))
iter_weigths= tf.constant(np.array(iter_weights)[:, None], dtype=tf.float32)
sum_iter_weights = tf.constant(np.sum(iter_weights), dtype=tf.float32)
seq_ARI = tf.reduce_mean(tf.reduce_sum(ARI * iter_weigths, axis=0) / sum_iter_weights)
last_ARI = tf.reduce_mean(ARI[-1])
confidences = tf.reduce_sum(tf.reduce_max(gammas, axis=1, keep_dims=True) * M, axis=[1, 2]) / n
confidences = tf.reshape(confidences, shape=(yshape[0], yshape[1]))
seq_conf = tf.reduce_mean(tf.reduce_sum(confidences * iter_weigths, axis=0) / sum_iter_weights)
last_conf = tf.reduce_mean(confidences[-1])
return seq_ARI, last_ARI, seq_conf, last_conf
def color_spines(ax, color, lw=2):
for sn in ['top', 'bottom', 'left', 'right']:
ax.spines[sn].set_linewidth(lw)
ax.spines[sn].set_color(color)
ax.spines[sn].set_visible(True)
def color_half_spines(ax, color1, color2, lw=2):
for sn in ['top', 'left']:
ax.spines[sn].set_linewidth(lw)
ax.spines[sn].set_color(color1)
ax.spines[sn].set_visible(True)
for sn in ['bottom', 'right']:
ax.spines[sn].set_linewidth(lw)
ax.spines[sn].set_color(color2)
ax.spines[sn].set_visible(True)
def get_gamma_colors(nr_colors):
hsv_colors = np.ones((nr_colors, 3))
hsv_colors[:, 0] = (np.linspace(0, 1, nr_colors, endpoint=False) + 2/3) % 1.0
color_conv = hsv_to_rgb(hsv_colors)
return color_conv
def overview_plot(i, gammas, preds, inputs, corrupted=None, **kwargs):
attentions = np.array(kwargs['attentions']) if 'attentions' in kwargs else None
T, B, K, W, H, C = gammas.shape
T -= 1 # the initialization doesn't count as iteration
corrupted = corrupted if corrupted is not None else inputs
gamma_colors = get_gamma_colors(K)
# restrict to sample i and get rid of useless dims
inputs = inputs[:, i, 0]
gammas = gammas[:, i, :, :, :, 0]
if preds.shape[1] != B:
preds = preds[:, 0]
preds = preds[:, i]
corrupted = corrupted[:, i, 0]
inputs = np.clip(inputs, 0., 1.)
preds = np.clip(preds, 0., 1.)
corrupted = np.clip(corrupted, 0., 1.)
def plot_img(ax, data, cmap='Greys_r', xlabel=None, ylabel=None, border_color=None):
if data.shape[-1] == 1:
ax.matshow(data[:, :, 0], cmap=cmap, vmin=0., vmax=1., interpolation='nearest')
else:
ax.imshow(data, interpolation='nearest')
ax.set_xticks([]); ax.set_yticks([])
ax.set_xlabel(xlabel, color=border_color or 'k') if xlabel else None
ax.set_ylabel(ylabel, color=border_color or 'k') if ylabel else None
if border_color:
color_spines(ax, color=border_color)
def plot_attention_summary_img(ax, attention, k_excluded, preds, cmap='Greys_r', xlabel=None, ylabel=None, border_color=None):
# copy so we don't mutate
preds = np.copy(preds)
attention = np.copy(attention)
# get focus object as rgb version of black-white
focus_pred = np.tile(np.copy(preds[k_excluded]), [1, 1, 3])
# we are safe to do what we want to do to the k_excluded row
preds[k_excluded] = 0
attention = np.insert(attention, k_excluded, 0) # zero out the focus object
# mask preds by attention
preds = np.transpose(preds[:, :, :, 0], [1, 2, 0]) # (28, 28, K)
preds *= attention
# color the preds
preds = preds.reshape(-1, preds.shape[-1]).dot(gamma_colors).reshape(preds.shape[:-1] + (3,))
# add in the focus object
preds += focus_pred
# plot
ax.imshow(preds, interpolation='nearest')
ax.set_xticks([]); ax.set_yticks([])
ax.set_xlabel(xlabel, color=border_color or 'k') if xlabel else None
ax.set_ylabel(ylabel, color=border_color or 'k') if ylabel else None
if border_color:
color_spines(ax, color=border_color)
def plot_gamma(ax, gamma, xlabel=None, ylabel=None):
gamma = np.transpose(gamma, [1, 2, 0])
gamma = gamma.reshape(-1, gamma.shape[-1]).dot(gamma_colors).reshape(gamma.shape[:-1] + (3,))
ax.imshow(gamma, interpolation='nearest')
ax.set_xticks([])
ax.set_yticks([])
ax.set_xlabel(xlabel) if xlabel else None
ax.set_ylabel(ylabel) if ylabel else None
nrows, ncols = (K + 4, T + 1)
fig, axes = plt.subplots(nrows=nrows, ncols=ncols,
figsize=(2 * ncols, 2 * nrows))
axes[0, 0].set_visible(False)
axes[1, 0].set_visible(False)
plot_gamma(axes[2, 0], gammas[0], ylabel='Gammas')
for k in range(K + 1):
axes[k + 3, 0].set_visible(False)
for t in range(1, T + 1):
g = gammas[t]
p = preds[t]
reconst = np.sum(g[:, :, :, None] * p, axis=0)
plot_img(axes[0, t], inputs[t])
plot_img(axes[1, t], reconst)
plot_gamma(axes[2, t], g)
for k in range(K):
if attentions is not None:
plot_attention_summary_img(axes[k + 3, t], attentions[t-1, i, k], k, p,
border_color=tuple(gamma_colors[k]),
ylabel=('contexts {} for {}'.format(k-1, k) if t == 1 else None))
else:
plot_img(axes[k + 3, t], p[k], border_color=tuple(gamma_colors[k]), ylabel=('mu_{}'.format(k) if t == 1 else None))
plot_img(axes[K + 3, t], corrupted[t - 1])
plt.subplots_adjust(hspace=0.1, wspace=0.1)
return fig
def curve_plot(values_dict, coarse_range, fine_range):
if fine_range is not None:
fig, ax = plt.subplots(1, 2, figsize=(40, 10))
else:
fig, ax = plt.subplots(1, 1, figsize=(20, 10))
ax = [ax]
for key, values in values_dict.items():
# coarse
ax[0].plot(values, label=key)
ax[0].set_xlabel('epochs')
ax[0].axis([0, len(values), coarse_range[0], coarse_range[1]])
ax[0].set_title("coarse range")
ax[0].legend()
# fine
if fine_range is not None:
ax[1].plot(values, label=key)
ax[1].set_xlabel('epochs')
ax[1].axis([0, len(values), fine_range[0], fine_range[1]])
ax[1].set_title("fine range")
ax[1].legend()
return fig
