MRI_recon/code/Frequency-Diffusion/FSMNet/dataloaders/fastmri.py

import csv
import os
import random
import xml.etree.ElementTree as etree
from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
import pathlib

import h5py
import numpy as np
import torch
import yaml
from torch.utils.data import Dataset
from .transforms import build_transforms
from matplotlib import pyplot as plt

from .albu_transform import get_albu_transforms

def fetch_dir(key, data_config_file=pathlib.Path("fastmri_dirs.yaml")):
    """
    Data directory fetcher.

    This is a brute-force simple way to configure data directories for a
    project. Simply overwrite the variables for `knee_path` and `brain_path`
    and this function will retrieve the requested subsplit of the data for use.

    Args:
        key (str): key to retrieve path from data_config_file.
        data_config_file (pathlib.Path,
            default=pathlib.Path("fastmri_dirs.yaml")): Default path config
            file.

    Returns:
        pathlib.Path: The path to the specified directory.
    """
    if not data_config_file.is_file():
        default_config = dict(
            knee_path="/home/jc3/Data/",
            brain_path="/home/jc3/Data/",
        )
        with open(data_config_file, "w") as f:
            yaml.dump(default_config, f)

        raise ValueError(f"Please populate {data_config_file} with directory paths.")

    with open(data_config_file, "r") as f:
        data_dir = yaml.safe_load(f)[key]

    data_dir = pathlib.Path(data_dir)

    if not data_dir.exists():
        raise ValueError(f"Path {data_dir} from {data_config_file} does not exist.")

    return data_dir


def et_query(
        root: etree.Element,
        qlist: Sequence[str],
        namespace: str = "http://www.ismrm.org/ISMRMRD",
) -> str:
    """
    ElementTree query function.
    This can be used to query an xml document via ElementTree. It uses qlist
    for nested queries.
    Args:
        root: Root of the xml to search through.
        qlist: A list of strings for nested searches, e.g. ["Encoding",
            "matrixSize"]
        namespace: Optional; xml namespace to prepend query.
    Returns:
        The retrieved data as a string.
    """
    s = "."
    prefix = "ismrmrd_namespace"

    ns = {prefix: namespace}

    for el in qlist:
        s = s + f"//{prefix}:{el}"

    value = root.find(s, ns)
    if value is None:
        raise RuntimeError("Element not found")

    return str(value.text)


class SliceDataset(Dataset):
    def __init__(
            self,
            root,
            transform,
            challenge,
            sample_rate=1,
            mode='train'
    ):
        self.mode = mode
        self.albu_transforms = get_albu_transforms(self.mode, (320, 320))


        # challenge
        if challenge not in ("singlecoil", "multicoil"):
            raise ValueError('challenge should be either "singlecoil" or "multicoil"')
        self.recons_key = (
            "reconstruction_esc" if challenge == "singlecoil" else "reconstruction_rss"
        )
        # transform
        self.transform = transform

        self.examples = []

        self.cur_path = root
        if not os.path.exists(self.cur_path):
            self.cur_path = self.cur_path + "_selected"

        self.csv_file = "knee_data_split/singlecoil_" + self.mode + "_split_less.csv"

        with open(self.csv_file, 'r') as f:
            reader = csv.reader(f)

            id = 0

            for row in reader:
                pd_metadata, pd_num_slices = self._retrieve_metadata(os.path.join(self.cur_path, row[0] + '.h5'))

                pdfs_metadata, pdfs_num_slices = self._retrieve_metadata(os.path.join(self.cur_path, row[1] + '.h5'))

                for slice_id in range(min(pd_num_slices, pdfs_num_slices)):
                    self.examples.append(
                        (os.path.join(self.cur_path, row[0] + '.h5'), os.path.join(self.cur_path, row[1] + '.h5')
                         , slice_id, pd_metadata, pdfs_metadata, id))
                id += 1

        if sample_rate < 1:
            random.shuffle(self.examples)
            num_examples = round(len(self.examples) * sample_rate)

            self.examples = self.examples[0:num_examples]

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, i):

        # read pd
        pd_fname, pdfs_fname, slice, pd_metadata, pdfs_metadata, id = self.examples[i]

        with h5py.File(pd_fname, "r") as hf:
            pd_kspace = hf["kspace"][slice]

            pd_mask = np.asarray(hf["mask"]) if "mask" in hf else None

            pd_target = hf[self.recons_key][slice] if self.recons_key in hf else None

            attrs = dict(hf.attrs)

            attrs.update(pd_metadata)

        if self.transform is None:
            pd_sample = (pd_kspace, pd_mask, pd_target, attrs, pd_fname, slice)
        else:
            pd_sample = self.transform(pd_kspace, pd_mask, pd_target, attrs, pd_fname, slice)

        with h5py.File(pdfs_fname, "r") as hf:
            pdfs_kspace = hf["kspace"][slice]
            pdfs_mask = np.asarray(hf["mask"]) if "mask" in hf else None

            pdfs_target = hf[self.recons_key][slice] if self.recons_key in hf else None

            attrs = dict(hf.attrs)

            attrs.update(pdfs_metadata)

        if self.transform is None:
            pdfs_sample = (pdfs_kspace, pdfs_mask, pdfs_target, attrs, pdfs_fname, slice)
        else:
            pdfs_sample = self.transform(pdfs_kspace, pdfs_mask, pdfs_target, attrs, pdfs_fname, slice)

        # 0: input, 1: target, 2: mean, 3: std
        sample = self.albu_transforms(image=pdfs_sample[1].numpy(),
                                      image2=pd_sample[1].numpy(),
                                      image3=pdfs_sample[0].numpy(),
                                      image4=pd_sample[0].numpy())

        pdfs_sample = list(pdfs_sample)
        pd_sample = list(pd_sample)
        pdfs_sample[1] = sample['image']
        pd_sample[1]   = sample['image2']
        pdfs_sample[0] = sample['image3']
        pd_sample[0]   = sample['image4']

        # dataset pdf mean and std tensor(3.1980e-05) tensor(1.3093e-05)
        # print("dataset pdf mean and std", pdfs_sample[2], pdfs_sample[3])
        # print(pdfs_sample[1].shape, pdfs_sample[1].min(), pdfs_sample[1].max())

        return (pd_sample, pdfs_sample, id)

    def _retrieve_metadata(self, fname):
        with h5py.File(fname, "r") as hf:
            et_root = etree.fromstring(hf["ismrmrd_header"][()])

            enc = ["encoding", "encodedSpace", "matrixSize"]
            enc_size = (
                int(et_query(et_root, enc + ["x"])),
                int(et_query(et_root, enc + ["y"])),
                int(et_query(et_root, enc + ["z"])),
            )
            rec = ["encoding", "reconSpace", "matrixSize"]
            recon_size = (
                int(et_query(et_root, rec + ["x"])),
                int(et_query(et_root, rec + ["y"])),
                int(et_query(et_root, rec + ["z"])),
            )

            lims = ["encoding", "encodingLimits", "kspace_encoding_step_1"]
            enc_limits_center = int(et_query(et_root, lims + ["center"]))
            enc_limits_max = int(et_query(et_root, lims + ["maximum"])) + 1

            padding_left = enc_size[1] // 2 - enc_limits_center
            padding_right = padding_left + enc_limits_max

            num_slices = hf["kspace"].shape[0]

        metadata = {
            "padding_left": padding_left,
            "padding_right": padding_right,
            "encoding_size": enc_size,
            "recon_size": recon_size,
        }

        return metadata, num_slices


def build_dataset(args, mode='train', sample_rate=1, use_kspace=False):
    assert mode in ['train', 'val', 'test'], 'unknown mode'
    transforms = build_transforms(args, mode, use_kspace)

    return SliceDataset(os.path.join(args.root_path, 'singlecoil_' + mode), transforms, 'singlecoil', sample_rate=sample_rate, mode=mode)


if __name__ == "__main__":
    ## make logger file
    from torch.utils.data import DataLoader
    from option import args
    import time
    from frequency_diffusion.degradation.k_degradation import get_ksu_kernel, apply_ksu_kernel, apply_tofre, \
        apply_to_spatial

    batch_size = 1
    db_train = build_dataset(args, mode='train')

    trainloader = DataLoader(db_train, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)

    for i_batch, sampled_batch in enumerate(trainloader):
        time2 = time.time()
        # print("time for data loading:", time2 - time1)

        pd, pdfs, _ = sampled_batch
        target = pdfs[1]

        mean = pdfs[2]
        std = pdfs[3]

        pd_img = pd[1].unsqueeze(1)
        pdfs_img = pdfs[0].unsqueeze(1)
        target = target.unsqueeze(1)

        b = pd_img.size(0)

        pd_img = pd_img  # [4, 1, 320, 320]
        pdfs_img = pdfs_img  # [4, 1, 320, 320]
        target = target  # [4, 1, 320, 320]

        # ----------- Degradation -------------
        num_timesteps = 1
        image_size = 320

        # Output a list of k-space kernels
        kspace_masks = get_ksu_kernel(num_timesteps, image_size,
                                      ksu_routine="LogSamplingRate",
                                      accelerated_factor=args.ACCELERATIONS[0],
                                      )  # args.ACCELERATIONS = [4] or [8]
        kspace_masks = torch.from_numpy(np.asarray(kspace_masks)).cuda()



        t = torch.randint(0, num_timesteps, (b,)).long()
        mask = kspace_masks[t]
        fft, mask = apply_tofre(target.clone(), mask)
        # fft = fft * mask + 0.0
        pdfs_img = apply_to_spatial(fft)
        pdfs_img_mask = apply_to_spatial(mask * fft)[0]




        print("mask = ", mask.shape, mask.min(), mask.max())
        print("pdfs_img_mask =", pdfs_img_mask.shape)

        import matplotlib.pyplot as plt

        # combine them together
        pd_img = pd_img.squeeze(1).cpu().numpy()
        pdfs_img = pdfs_img.squeeze(1).cpu().numpy()
        target = target.squeeze(1).cpu().numpy()

        plt.figure()

        plt.subplot(161)
        plt.imshow(pd_img[0], cmap='gray')
        plt.title('PD')
        plt.axis('off')
        plt.subplot(162)

        plt.imshow(pdfs_img_mask[0], cmap='gray')
        plt.title('PDFS_mask')
        plt.axis('off')

        plt.subplot(163)
        plt.imshow(pdfs_img[0], cmap='gray')
        plt.title('PDFS')
        plt.axis('off')

        plt.subplot(164)
        plt.imshow(pdfs_img_mask[0] - target[0], cmap='gray')
        plt.title('Diff')
        plt.axis('off')

        plt.subplot(165)
        plt.imshow(target[0], cmap='gray')
        plt.title('Target')
        plt.axis('off')

        plt.subplot(166)
        plt.imshow(pdfs_img[0] - target[0], cmap='gray')#mask[0][0], cmap='gray')
        plt.title('Target')
        plt.axis('off')

        plt.show()