AlternativeModels/AAE/AAE_inference.py

import torch
import scipy.io as io
import numpy as np
import argparse
from utils import load_dataset_test, save_image_mat
from fMRIAAE_Model import *
import torch.utils.data
import os
import evals
import nibabel as nib

parser = argparse.ArgumentParser(description='AAE for fMRI generation')
parser.add_argument('--batch-size', type=int, default=120, metavar='N',
                            help='input batch size for training (default: 128)')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                            help='random seed (default: 1)')
parser.add_argument('--zdim', type=int, default=256, metavar='N',
                            help='dimension of latent variables')
parser.add_argument('--dataname', default='demo_data', type=str, help='dataset name')

parser.add_argument('--z-path', type=str, default='./result/demo_latent/', help='path to saved z files')
parser.add_argument('--resume', type=str, default='./checkpoint/checkpoint.pth.tar', help='the VAE checkpoint') 
parser.add_argument('--save-cifti', type=bool, default=False, help='whether to save the reconstructed images as CIFTI files')
parser.add_argument('--img-path',type=str,default='./result/recon',help = 'path to save CIFTI files')
parser.add_argument('--mode',type=str,default='encode',help = 'encode or decode')
args = parser.parse_args()
torch.manual_seed(args.seed)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
data_path="/mnt/leuthardte/Active/Cindy/Data/"+ args.dataname

model = AAE(z_dim=args.zdim, nc=1).to(device)
if os.path.isfile(args.resume):
    print("==> Loading checkpoint: ", args.resume)
    checkpoint = torch.load(args.resume, map_location=device)
    model.load_state_dict(checkpoint['state_dict'])
else:
    print('[ERROR] Checkpoint not found: ', args.resume)
    raise RuntimeError

if args.mode == 'encode' or args.mode == 'both':
    '''
    Encoder mode. 
    Image --> z
    '''
    test_loader = load_dataset_test(data_path, args.batch_size)

    print('Mode: Encode \n Distribution of Z will be saved at: ' + args.z_path)
    if not os.path.isdir(args.z_path):
        os.system('mkdir '+ args.z_path + ' -p')

    for batch_idx, (xL, xR) in enumerate(test_loader):
        xL = xL.to(device)
        xR = xR.to(device)
        _,mu,logvar = model.encoder(xL, xR)
        z_distribution = np.hstack((mu.detach().cpu().numpy(),logvar.detach().cpu().numpy()))
        save_data = {}
        save_data['z_distribution'] = z_distribution
        io.savemat(os.path.join(args.z_path, 'save_z{}.mat'.format(batch_idx)), save_data)

if args.mode == 'decode' or args.mode == 'both':
    '''
    Decoder mode.
    z --> reconstructed image
    '''
    print('Mode: Decode \n Reconstructed images will be saved at: ' + args.img_path)
    if not os.path.isdir(args.img_path):
        os.system('mkdir '+ args.img_path + ' -p')
        
    filelist = [f for f in os.listdir(args.z_path) if f.split('_')[0] == 'save']
    for batch_idx, filename in enumerate(filelist):

        # z_dist = io.loadmat(os.path.join(args.z_path, filename))
        z_dist = io.loadmat(os.path.join(args.z_path, 'save_z{}.mat'.format(batch_idx)))
        z_dist = z_dist['z_distribution']
        mu=z_dist[:, :args.zdim]
        logvar = z_dist[:, args.zdim:]

        z = torch.tensor(mu).to(device)
        x_recon_L, x_recon_R = model.decoder(z)
        # recon_dtseries = evals.backward_reformatting(x_recon_L,x_recon_R) # not sure why this was not working
        # print(recon_dtseries.shape)
        # if args.save_cifti:
        #     recon_cifti_name = args.dataname +'_AAE_Zdim' + str(args.zdim) + '_recon.dscalar.nii'
        #     empty_cifti = nib.load("empty_cifti.dscalar.nii")
        #     brain_model_axis = empty_cifti.header.get_axis(1)
        #     scalar_axis = nib.cifti2.cifti2_axes.ScalarAxis(['Column' + str(i+1) for i in range(recon_dtseries.shape[0])])
        #     img = nib.Cifti2Image(np.single(recon_dtseries), header=(scalar_axis,brain_model_axis))
        #     nib.save(img, args.img_path + recon_cifti_name + '_' + '{:02d}'.format(batch_idx))
        save_image_mat(x_recon_R, x_recon_L, args.img_path, batch_idx)