plot_ev_frames.py

import os
import numpy as np
import matplotlib.pyplot as plt


root_dir_in = f'{os.path.dirname(os.path.abspath(__file__))}/DATA/EventFrames2/TEST_FULL_SEQUENCE'
root_dir_out = f'{os.path.dirname(os.path.abspath(__file__))}/DATA/trials'

subjects = [2]
configs = [1, 2, 3, 4, 5]

for subject in subjects:
    for config in configs:
        f_name_frames_ev = f'{root_dir_in}/subject{subject:02}_seq{config:02}.npy'
        f_output_dir = f'{root_dir_out}/subject{subject:02}_seq{config:02}/frames'

        if not os.path.exists(f_output_dir):
            os.makedirs(f_output_dir)

        frames = np.load(f_name_frames_ev)
        for i, frame in enumerate(frames):
            fig, axs = plt.subplots(1, 1)
            mpos = frame[:,:,0,0]
            mpos[mpos == 0] = np.nan
            mpos = mpos.astype(np.float32)
            axs.imshow(mpos, alpha=0.5, cmap='Reds')

            mneg = frame[:,:,0,1]
            mneg[mneg == 0] = np.nan
            mneg = mneg.astype(np.float32)
            axs.imshow(mneg, alpha=0.5, cmap='Greens')

            fig.savefig(f'{f_output_dir}/frame_{i:05d}.png')
            plt.close(fig)