| 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() |
|
|
