|
import torch |
|
import scipy.io as io |
|
import argparse |
|
import logging |
|
from utils import load_dataset_test, save_image_mat |
|
from fMRIVAE_Model import BetaVAE |
|
import os |
|
|
|
def main(): |
|
|
|
parser = argparse.ArgumentParser(description='VAE for fMRI generation') |
|
parser.add_argument('--batch-size', type=int, metavar='N', help='how many samples per saved file?') |
|
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('--data-path', type=str, metavar='DIR', help='path to dataset') |
|
parser.add_argument('--z-path', type=str, default='./result/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('--img-path', type=str, default='./result/recon', help='path to save reconstructed images') |
|
parser.add_argument('--mode', type=str, default='both', help='choose from \'encode\',\'decode\' or \'both\'') |
|
parser.add_argument('--debug', action='store_true', help='Enable debug mode for detailed logging') |
|
|
|
args = parser.parse_args() |
|
|
|
if not os.path.isdir(args.z_path): |
|
os.system('mkdir '+ args.z_path + ' -p') |
|
if (args.mode != 'encode') and not os.path.isdir(args.img_path): |
|
os.system('mkdir '+ args.img_path + ' -p') |
|
|
|
|
|
logging_level = logging.DEBUG if args.debug else logging.INFO |
|
logging.basicConfig(level=logging_level, format='%(asctime)s - %(levelname)s - %(message)s') |
|
|
|
logging.debug("Starting the VAE inference script.") |
|
args = parser.parse_args() |
|
logging.debug(f"Parsed arguments: {args}") |
|
|
|
try: |
|
torch.manual_seed(args.seed) |
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
logging.debug(f"Using device: {device}") |
|
|
|
logging.debug(f"Loading VAE model from {args.resume}.") |
|
model = BetaVAE(z_dim=args.zdim, nc=1).to(device) |
|
if os.path.isfile(args.resume): |
|
checkpoint = torch.load(args.resume, map_location=device) |
|
model.load_state_dict(checkpoint['state_dict']) |
|
logging.debug("Checkpoint loaded.") |
|
else: |
|
logging.error(f"Checkpoint not found at {args.resume}") |
|
raise RuntimeError("Checkpoint not found.") |
|
|
|
if (args.mode == 'encode') or (args.mode == 'both'): |
|
logging.debug("Starting encoding process...") |
|
test_loader = load_dataset_test(args.data_path, args.batch_size) |
|
logging.debug(f"Loaded test dataset from {args.data_path}") |
|
for batch_idx, (xL, xR) in enumerate(test_loader): |
|
xL = xL.to(device) |
|
xR = xR.to(device) |
|
z_distribution = model._encode(xL, xR) |
|
save_data = {'z_distribution': z_distribution.detach().cpu().numpy()} |
|
io.savemat(os.path.join(args.z_path, f'save_z{batch_idx}.mat'), save_data) |
|
logging.debug(f"Encoded batch {batch_idx}") |
|
|
|
if (args.mode == 'decode') or (args.mode == 'both'): |
|
logging.debug("Starting decoding process...") |
|
filelist = [f for f in os.listdir(args.z_path) if f.split('_')[0] == 'save'] |
|
logging.debug(f"Filelist: {filelist}") |
|
for batch_idx, filename in enumerate(filelist): |
|
logging.debug(f"Decoding file {filename}") |
|
z_dist = io.loadmat(os.path.join(args.z_path, f'save_z{batch_idx}.mat')) |
|
z_dist = z_dist['z_distribution'] |
|
mu = z_dist[:, :args.zdim] |
|
z = torch.tensor(mu).to(device) |
|
x_recon_L, x_recon_R = model._decode(z) |
|
save_image_mat(x_recon_R, x_recon_L, args.img_path, batch_idx) |
|
logging.debug(f"Decoded and saved batch {batch_idx}") |
|
|
|
except Exception as e: |
|
logging.error(f"An error occurred: {e}") |
|
raise |
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
