demo.py

import os, sys, time
from math import ceil

import nibabel as nib
import SimpleITK as sitk
import pyqtgraph as pg
import numpy as np
import scipy.io as scio
from functions.apart_qsm import *
import torch

if __name__ == '__main__':
    # ####gpu
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') 
    input_path = r"F:\code\Lab_github\apartqsm_python\data"
    output_path = r"F:\code\Lab_github\apartqsm_python\results"
    #dcm_path = os.path.join(input_path, 'GRE')
    t2 = nib.load(os.path.join(input_path, 'T2.nii.gz')).get_fdata()
    TE = np.loadtxt(os.path.join(input_path, 'TE.txt'))
    voxel_size=(1.0,1.0,2.0)
    B0=3.0
    B0_dir=(0.0,0.0,1.0)
    tissue_phase=nib.load(os.path.join(input_path,'tissue_phase.nii.gz')).get_fdata()
    new_mask=nib.load(os.path.join(input_path,'newmask.nii.gz')).get_fdata()
    magnitude=nib.load(os.path.join(input_path,'mag.nii.gz')).get_fdata()
    qsm=nib.load(os.path.join(input_path,'qsm.nii.gz')).get_fdata()
    if tissue_phase.shape[3]>6:
        TE_phi = TE[2:]
        phase_shape=tissue_phase[...,2:].shape
    else:
        TE_phi = TE
        phase_shape = tissue_phase.shape
    
    #### APART-QSM
    params_input = build_params_input(
    TE_mag=TE,
    TE_phi=TE_phi,
    voxel_size= np.array(voxel_size, dtype=np.float32),
    B0=B0,
    B0_dir=np.array(B0_dir, dtype=np.float32),
    newmask=new_mask,
    phase_shape=phase_shape,
    mag_shape = magnitude.shape,
    qsm=qsm,
    magnitude=magnitude,
    voxel_size_used=voxel_size,
    t2_map=t2,
    matrix_size=magnitude.shape[:3])

    recon = ARART_Recon(
    mag=magnitude,
    phi_tissue=tissue_phase,
    QSM=qsm,
    params=params_input,
    t2_map=t2,
    output_path=output_path
    )
    Res_map = recon.run()