Prepare masks

Simulate 2D+t Gaussian Cartesian random subsampling masks at varying accelerations using the deepinv library.

A = 8 # acceleration

%cd ..
import os, pickle
from pathlib import Path
from tqdm import tqdm

import numpy as np
from scipy.io import savemat
import torch

from deepinv.physics.generator import GaussianMaskGenerator

from utils.utils import set_posix_windows

root = Path("M:/data/CMRxRecon")
set_name = "TrainingSet"
root = root / "SingleCoil/Cine" / set_name
acc_folder = "FullSample"
new_mask_folder = "TimeVaryingGaussianMask" + str(A).zfill(2)

seed = 0
rng = np.random.default_rng(seed)
torch.manual_seed(seed)
np.random.seed(seed)

# Open existing dataset cache
with open("dataset_cache_windows.pkl", "rb") as f:
    with set_posix_windows():
        dataset_cache = pickle.load(f)

# Iterate through existing dataset cache

done_fnames = []
for sample in tqdm(dataset_cache[root]):
    if sample.fname in done_fnames: continue
    done_fnames.append(sample.fname)

    mask_fname = sample.fname.replace(acc_folder, new_mask_folder)[:-4] + "_mask.mat"
    os.makedirs(os.path.dirname(mask_fname), exist_ok=True)

    generator = GaussianMaskGenerator(
        (2, sample.metadata["timeframes"], sample.metadata["width"], sample.metadata["height"]), 
        acceleration=A, 
        rng=rng,
    )

    mask = generator.step()["mask"].detach().numpy().astype("uint8")
    mask = mask[0, 0, :, :, :] #only save T,H,W
    mask = np.moveaxis(mask, 0, -1) #H, W, T
    
    savemat(mask_fname, {"mask": mask}, do_compression=True)

  0%|          | 0/1492 [00:00<?, ?it/s]100%|██████████| 1492/1492 [00:48<00:00, 31.03it/s]