import os
import glob
import pydicom
import aux_dicom
import numpy as np
import pandas as pd
import nibabel as nib
import vtk
from vtk.util.numpy_support import vtk_to_numpy
import warnings
from skimage.measure import find_contours
class Tag:
def __init__(self, subject_id=1, data_dir='../datasets/CMAC', work_dir='results/data', results_group='MEVIS'):
self.subject_id = subject_id
self.data_dir = os.path.abspath(data_dir)
self.work_dir = os.path.abspath(work_dir)
self.raw_dir = os.path.join(self.data_dir, 'raw', f'v{subject_id}', '3DTAG')
self.gt_dir = os.path.join(self.data_dir, 'GT', '3DTAG', f'v{subject_id}')
self.results_dir = os.path.join(self.data_dir, 'RESULTS', results_group, '3DTAG', f'v{subject_id}')
self.results_group = results_group
self.affine_axial = nib.load(os.path.join(self.raw_dir, 'NIFTI', 'NIFTI00.nii')).affine
def load_pts_gt(self, lmks_obs=1, lmks_frame_id=0):
"""Load GT mesh and landmarks."""
coords = {
'MEVIS': 'VTK_COORDINATES',
'UPF': 'VTK_COORDINATES',
'IUCL': 'DICOM_COORDINATES',
'INRIA': 'INRIA_COORDINATES'
}[self.results_group]
mesh_gt_mm, region_ids, subpart_ids = load_pts_mm(
os.path.join(self.gt_dir, 'MESH', coords, f'v{self.subject_id}.vtk'),
regionID=True, subpartID=True
)
self.mesh_gt_mm = mesh_gt_mm
self.regionIDs = region_ids
self.subpartIDs = subpart_ids
self.lmks_gt_mm = load_pts_mm(
os.path.join(self.gt_dir, 'LMKS', coords, f'obs{lmks_obs}_groundTruth{lmks_frame_id:02d}.vtk')
)
self.mesh_gt_mm[:, :2] *= -1
self.lmks_gt_mm[:, :2] *= -1
def load_pts(self, frame_id=0, obs=1):
"""Load predicted mesh and landmarks."""
self.mesh_pred_mm = load_pts_mm(
os.path.join(self.results_dir, 'MESH', f'finalMesh{frame_id:03d}.vtk')
)
self.lmks_pred_mm = load_pts_mm(
os.path.join(self.results_dir, 'LMKS', f'obs{obs}_results{frame_id:03d}.vtk')
)
self.mesh_pred_mm[:, :2] *= -1
self.lmks_pred_mm[:, :2] *= -1
def load_nifti_3d(self, frame_id, set_affine=None):
nifti = nib.load(os.path.join(self.raw_dir, 'NIFTI', f'NIFTI{frame_id:02d}.nii'))
return nib.Nifti1Image(nifti.get_fdata(), set_affine if set_affine is not None else nifti.affine)
def load_nifti_4d(self, n_frames=20, set_affine=None):
shape = self.load_nifti_3d(frame_id=0).shape
arr_4d = np.zeros(shape + (n_frames,))
for frame_id in range(n_frames):
arr_3d = self.load_nifti_3d(frame_id=frame_id).get_fdata()
arr_4d[:, :, :, frame_id] = arr_3d
affine = set_affine if set_affine is not None else self.affine_axial
return nib.Nifti1Image(arr_4d, affine)
class Cine:
def __init__(self, subject_id=1, data_dir='cMAC', work_dir='results/data', results_group='UPF'):
self.subject_id = subject_id
self.data_dir = os.path.abspath(data_dir)
self.work_dir = os.path.abspath(work_dir)
self.raw_dir = os.path.join(self.data_dir, 'raw', f'v{subject_id}', 'cSAX')
self.gt_dir = os.path.join(self.data_dir, 'GT', 'SSFP', f'v{subject_id}')
self.results_dir = os.path.join(self.data_dir, 'RESULTS', results_group, 'SSFP', f'v{subject_id}')
self.results_group = results_group
self.ijk2ras, self.affine, self.affine_axial = aux_dicom.read_affine_info(os.path.join(self.raw_dir, 'time_1', '*'))
def load_pts_gt(self, lmks_frame_id=0, obs=1):
"""Load GT mesh and landmarks."""
coords = {
'MEVIS': 'VTK_COORDINATES',
'UPF': 'VTK_COORDINATES',
'IUCL': 'DICOM_COORDINATES',
'INRIA': 'INRIA_COORDINATES'
}[self.results_group]
mesh_gt_mm, region_ids, subpart_ids = load_pts_mm(
os.path.join(self.gt_dir, 'MESH', 'VTK_COORDINATES', f'v{self.subject_id}.vtk'),
regionID=True, subpartID=True
)
self.mesh_gt_mm = mesh_gt_mm
self.regionIDs = region_ids
self.subpartIDs = subpart_ids
self.lmks_gt_mm = load_pts_mm(
os.path.join(self.gt_dir, 'LMKS', coords, f'obs{obs}_groundTruth{lmks_frame_id:02d}.vtk')
)
def load_pts(self, frame_id=0, obs=1):
"""Load predicted mesh and landmarks."""
self.mesh_pred_mm = load_pts_mm(
os.path.join(self.results_dir, 'MESH', f'finalMesh{frame_id:03d}.vtk')
)
self.lmks_pred_mm = load_pts_mm(
os.path.join(self.results_dir, 'LMKS', f'obs{obs}_results{frame_id:03d}.vtk')
)
def load_nifti_3d(self, frame_id, set_affine=None):
dicom_files = glob.glob(os.path.join(self.raw_dir, f'time_{frame_id + 1}', '*'))
cine_arrs_2d = [pydicom.dcmread(file).pixel_array.T for file in dicom_files]
cine_arr_3d = np.stack(cine_arrs_2d, -1)
return nib.Nifti1Image(cine_arr_3d, set_affine if set_affine is not None else self.affine)
def load_nifti_4d(self, n_frames=20, set_affine=None):
shape = self.load_nifti_3d(frame_id=0).shape
arr_4d = np.zeros(shape + (n_frames,))
for frame_id in range(n_frames):
arr_3d = self.load_nifti_3d(frame_id=frame_id).get_fdata()
arr_4d[:, :, :, frame_id] = arr_3d
affine = set_affine if set_affine is not None else self.affine
return nib.Nifti1Image(arr_4d, affine)
def load_pts_mm(pts_filename, regionID=False, subpartID=False):
"""Load points (mm) in 3D space stored in vtk file."""
reader = vtk.vtkDataSetReader()
reader.SetFileName(pts_filename)
reader.Update()
pts_mm = vtk_to_numpy(reader.GetOutput().GetPoints().GetData())
output = [pts_mm]
if regionID:
output.append(vtk_to_numpy(reader.GetOutput().GetPointData().GetArray('regionID')))
if subpartID:
output.append(vtk_to_numpy(reader.GetOutput().GetPointData().GetArray('subpartID')))
return output if len(output) > 1 else pts_mm
def pts_pix_to_mm(pts_pix, affine):
"""Apply affine transformation to pixel coordinates."""
pts_mm = np.dot(np.c_[pts_pix, np.ones(len(pts_pix))], affine.T)[:, :3]
return pts_mm
def pts_mm_to_pix(pts_mm, affine):
"""Apply inverse affine transformation to mm coordinates."""
pts_pix = np.dot(np.c_[pts_mm, np.ones(len(pts_mm))], np.linalg.inv(affine).T)[:, :3]
return pts_pix
def cine_pix_to_tag_pix(cine_pts_pix, cine_affine, tag_affine):
"""Convert cine pixel space points to tag pixel space points."""
return pts_mm_to_pix(pts_pix_to_mm(cine_pts_pix, cine_affine), tag_affine)
def tag_pix_to_cine_pix(tag_pts_pix, tag_affine, cine_affine):
"""Convert tag pixel space points to cine pixel space points."""
return pts_mm_to_pix(pts_pix_to_mm(tag_pts_pix, tag_affine), cine_affine)
def cine_cnts_pix_to_tag_cnts_pix(cine_cts_pix, cine_affine, tag_affine):
"""Convert contours in cine pixel space to points in tag pixel space."""
tag_cts_pix = []
for slice_contours in cine_cts_pix:
slice_tag_cts = [cine_pix_to_tag_pix(contour, cine_affine, tag_affine) for contour in slice_contours]
tag_cts_pix.append(slice_tag_cts)
return tag_cts_pix
def estimate_affine(ins, outs):
"""Estimate affine transformation matrix."""
l = len(ins)
B = np.vstack([np.transpose(ins), np.ones(l)])
D = 1.0 / np.linalg.det(B)
entry = lambda r, d: np.linalg.det(np.delete(np.vstack([r, B]), (d + 1), axis=0))
M = [[(-1) ** i * D * entry(R, i) for i in range(l)] for R in np.transpose(outs)]
return np.array(M + [[0, 0, 0, 1]])
###############################
####### EXTRA FUNCTIONS #######
###############################
def export_pts_mm_to_slicer_fiducials(pts_mm, fiducials_filename_fscv):
"""Save points (mm) in 3D space as a slicer fiducial file (.fscv)."""
with open(fiducials_filename_fscv, 'w') as file:
file.write('# Markups fiducial file version = 4.6\n')
file.write('# CoordinateSystem = 0\n')
file.write('# columns = id,x,y,z,ow,ox,oy,oz,vis,sel,lock,label,desc,associatedNodeID\n')
for idx, point in enumerate(pts_mm):
file.write(f'vtkMRMLMarkupsFiducialNode_0,{point[0]},{point[1]},{point[2]},0,0,0,1,1,1,0,F-{idx},,vtkMRMLScalarVolumeNode1\n')
def import_pts_mm_from_slicer_fiducials(fiducials_filename_fscv):
"""Load points (mm) in 3D space from a slicer fiducial file (.fscv)."""
return pd.read_csv(fiducials_filename_fscv, skiprows=3, header=None).iloc[:, 1:4].to_numpy()
def import_cnts_pix_from_label(label_nifti_filename, label_id=1):
"""Load and convert 3D label of myocardial contours to contour points in pixel space."""
label = nib.load(label_nifti_filename).get_fdata()
epicardium, endocardium = [], []
for z in range(label.shape[-1]):
contours = find_contours(label[:, :, z] == label_id, 0.8)
if len(contours) == 2:
epicardium.append(np.c_[contours[0], z * np.ones((len(contours[0]), 1))])
endocardium.append(np.c_[contours[1], z * np.ones((len(contours[1]), 1))])
return epicardium, endocardium
Datasets
Models
