visualize.py

#! /usr/bin/python
# -*- coding: utf8 -*-

import argparse
import os
import re
import scipy.io as sio
import numpy as np
from sklearn.metrics import cohen_kappa_score
from sklearn.metrics import confusion_matrix, f1_score
import matplotlib.pyplot as plt
W=0
N1=1
N2=2
N3=3
REM=4
classes= ['W','N1', 'N2','N3','REM']
n_classes = len(classes)

def plot_attention(attention_map, input_tags = None, output_tags = None):
    attn_len = len(attention_map)

    # Plot the attention_map
    # plt.clf()
    f = plt.figure(figsize=(15, 10))
    ax = f.add_subplot(1, 1, 1)
    ax.tick_params(axis='both', which='major', labelsize=15)
    ax.tick_params(axis='both', which='minor', labelsize=15)

    # Add image
    i = ax.imshow(attention_map, interpolation='nearest', cmap='gray')

    # Add colorbar
    cbaxes = f.add_axes([0.2, -0.02, 0.6, 0.03])
    cbar = f.colorbar(i, cax=cbaxes, orientation='horizontal')
    cbar.ax.set_xlabel('Alpha value (Probability output of the "softmax")', labelpad=2)
    cbar.ax.tick_params(labelsize=15)
    cbar.ax.xaxis.label.set_size(16)

    # Add labels
    ax.set_yticks(range(attn_len))
    if output_tags != None:
      ax.set_yticklabels(output_tags[:attn_len])

    ax.set_xticks(range(attn_len))
    if input_tags != None:
      ax.set_xticklabels(input_tags[:attn_len], rotation=45)


    ax.set_xlabel('Input: EEG Epochs (a sequence)')
    ax.xaxis.label.set_size(20)
    ax.set_ylabel('Output: Satge Scoring')
    ax.yaxis.label.set_size(20)

    # add grid and legend
    ax.grid()
    HERE = os.path.realpath(os.path.join(os.path.realpath(__file__), '..'))
    dir_save = os.path.join(HERE, 'attention_maps')
    if (os.path.exists(dir_save) == False):
        os.mkdir(dir_save)
    f.savefig(os.path.join(dir_save, 'a_map_6_5.pdf'), bbox_inches='tight')
    # f.show()
    plt.show()

def visualize(data_dir):

    # Remove non-output files, and perform ascending sort
    allfiles = os.listdir(data_dir)
    outputfiles = []
    for idx, f in enumerate(allfiles):
        if re.match("^output_.+\d+\.npz", f):
            outputfiles.append(os.path.join(data_dir, f))
    outputfiles.sort()


    with np.load(outputfiles[0]) as f:
        y_true = f["y_true"]
        y_pred = f["y_pred"]
        alignments_alphas_all = f["alignments_alphas_all"] # (batch_num,B,max_time_step,max_time_step)

    batch_len = len(alignments_alphas_all)
    char2numY = dict(zip(classes, range(len(classes))))
    num2charY = dict(zip(char2numY.values(), char2numY.keys()))
    shape = alignments_alphas_all.shape
    max_time_step = shape[2]
    batch_num = 6
    alignments_alphas = alignments_alphas_all[batch_num] # get results for the batch of batch_num
    y_true_ = np.reshape(y_true,[-1,shape[1],shape[2]])
    y_pred_ = np.reshape(y_pred, [-1, shape[1], shape[2]])

    input_tags = [[num2charY[i] for i in seq] for seq in y_true_[batch_num,:]]
    output_tags = [[num2charY[i] for i in seq] for seq in y_pred_[batch_num,:]]

    plot_attention(alignments_alphas[5, :, :], input_tags[5], output_tags[5])



def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_dir", type=str, default="outputs_2013/outputs_eeg_fpz_cz",
                        help="Directory where to load prediction outputs")
    args = parser.parse_args()

    if args.data_dir is not None:
        visualize(data_dir=args.data_dir)


if __name__ == "__main__":
    main()