asd-eeg/2_visualize.py at master · nirdslab/asd-eeg · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
import os

import numpy as np
from matplotlib import pyplot as plt

data_dir = 'out/csv/'


def run():
    file_names = [data_dir + x for x in os.listdir(data_dir)]

    electrodes = sorted(["F7", "F8", "T7", "T8", "TP9", "TP10", "P7", "P8", "C3", "C4"])
    file_names = sorted([x for x in sorted(file_names) if x.split('_')[2] in electrodes and "BASELINE" not in x])

    max_data = [[[0, np.zeros((150, 150))] for _ in range(len(electrodes))] for _ in range(2)]
    sum_data = [[[0, np.zeros((150, 150))] for _ in range(len(electrodes))] for _ in range(2)]

    for file_name in file_names:
        participant, epoch, electrode, label, chunk = file_name.split('/')[-1].split('_')
        chunk = int(chunk[:-4])
        label = int(label)
        diagnosis = 'ASD' if label else 'TD'

        data = np.loadtxt(file_name, delimiter=',').transpose()

        # Create individual PNG image
        name = '(%s) %s - %s [%s, %s/6] ' % (diagnosis, participant, epoch, electrode, chunk + 1)
        plt.figure(name)
        plt.title(name)
        plt.imshow(data, cmap='seismic')
        f_name = 'out/png/%s_%s_%s_%s_%s.png' % (diagnosis, electrode, epoch, chunk, participant)
        plt.savefig(f_name)
        plt.close(name)
        print(f_name + '\tDONE!')

        i = electrodes.index(electrode)

        # Max Aggregate for each Electrode
        max_data[label][i][0] += 1
        max_data[label][i][1] = np.maximum(max_data[label][i][1], data)

        # Avg Aggregate for each Electrode
        sum_data[label][i][0] += 1
        sum_data[label][i][1] = np.add(sum_data[label][i][1], data)

    for label, electrodes_data in enumerate(max_data):
        for i, electrode_data in enumerate(electrodes_data):
            electrode = electrodes[i]
            (cnt, mat) = electrode_data
            diagnosis = 'ASD' if label else 'TD'

            name = '(%s) Max Aggregate [%s] ' % (diagnosis, electrode)
            plt.figure(name)
            plt.title(name)
            plt.imshow(mat, cmap='seismic')

            f_name = 'out/png-max/%s_%s' % (diagnosis, electrode)
            plt.savefig(f_name)
            plt.close(name)

    for label, electrodes_data in enumerate(sum_data):
        for i, electrode_data in enumerate(electrodes_data):
            electrode = electrodes[i]
            (cnt, mat) = electrode_data
            mat /= cnt
            diagnosis = 'ASD' if label else 'TD'

            name = '(%s) Mean Aggregate [%s] ' % (diagnosis, electrode)
            plt.figure(name)
            plt.title(name)
            plt.imshow(mat, cmap='seismic')

            f_name = 'out/png-avg/%s_%s' % (diagnosis, electrode)
            plt.savefig(f_name)
            plt.close(name)


if __name__ == '__main__':
    os.makedirs('out/png', exist_ok=True)
    os.makedirs('out/png-avg', exist_ok=True)
    os.makedirs('out/png-max', exist_ok=True)
    run()