from SimpleITK import ImageFileReader, GetArrayFromImage
import math
import numpy as np
import xml.etree.ElementTree as ET
import os
from mayavi import mlab
from tvtk.api import tvtk, write_data
import stl
from tvtk.common import configure_input
import tkFileDialog
def get_transformation_matrix(p):
''' Transform from optotrak global coordinates to optotrak position sensor coordinates'''
q1 = p[0]
q2 = p[1]
q3 = p[2]
q4 = p[3]
q5 = p[4]
q6 = p[5]
T = np.zeros((4, 4))
T[0, 0] = math.cos(q6) * math.cos(q5)
T[1, 0] = math.sin(q6) * math.cos(q5)
T[2, 0] = -math.sin(q5)
T[0, 1] = math.cos(q6) * math.sin(q5) * math.sin(q4) - math.sin(q6) * math.cos(q4)
T[1, 1] = math.sin(q6) * math.sin(q5) * math.sin(q4) + math.cos(q6) * math.cos(q4)
T[2, 1] = math.cos(q5) * math.sin(q4)
T[0, 2] = math.cos(q6) * math.sin(q5) * math.cos(q4) + math.sin(q6) * math.sin(q4)
T[1, 2] = math.sin(q6) * math.sin(q5) * math.cos(q4) - math.cos(q6) * math.sin(q4)
T[2, 2] = math.cos(q5) * math.cos(q4)
T[0, 3] = q1
T[1, 3] = q2
T[2, 3] = q3
T[3, 3] = 1
return T
def get_positions(fname):
tree = ET.parse(fname)
root = tree.getroot()
locations = []
positions = []
for loc in list(list(root)[0]):
pos = loc.find('Anatomical')
locations.append(pos.get('file'))
pos2 = pos.find('USPosition')
positions.append([float(pos2.find('x').get('value')), float(pos2.find('y').get('value')), float(pos2.find('z').get('value')), float(pos2.find('roll').get('value')), float(pos2.find('pitch').get('value')), float(pos2.find('yaw').get('value'))])
return locations, positions
def main(subj_path, modality, segment):
filename_9L4 = '../Registration/Probe STLS/9L4.transformed.stl'
reader_9L4 = tvtk.STLReader()
reader_9L4.file_name = filename_9L4
reader_9L4.update()
# filename_14L5 = '../Registration/Probe STLS/14L5.transformed.stl'
# reader_14L5 = tvtk.STLReader()
# reader_14L5.file_name = filename_14L5
# reader_14L5.update()
reader = ImageFileReader()
reader.SetFileName(tkFileDialog.askopenfilename(title="Choose Nifti File of %s segment" % (segment), initialdir=os.path.join(subj_path, modality)))
img_all = reader.Execute()
origin = img_all.GetOrigin()
spacing = img_all.GetSpacing()
data = GetArrayFromImage(img_all)
data = data.T
fname = os.path.join(subj_path, os.path.split(subj_path)[1]+'_'+''.join([c for c in segment if c.isupper()])+'_US_'+modality[0:2]+'.xml')
print fname
locations, positions = get_positions(fname)
# reader = tvtk.STLReader()
# reader.file_name = tkFileDialog.askopenfilename(title="Choose STL File of %s segment" % (segment),
# initialdir=os.path.join(subj_path, modality))
# reader.update()
for loc_idx in range(len(positions)):
pos = positions[loc_idx]
print locations[loc_idx]
T = get_transformation_matrix(pos)
# Project the z-axis onto the x-y plane
y_vec = np.dot(T, np.array([0,1,0, 0]))[0:3]
z_vec = np.dot(T, np.array([0,0,1,0]))[0:3]
print z_vec
fig = mlab.figure(bgcolor=(0, 0, 0), size=(500, 500))
# to speed things up
fig.scene.disable_render = True
# mapper2 = tvtk.PolyDataMapper()
# configure_input(mapper2, reader.output)
#
# prop = tvtk.Property(opacity=0.3)
# actor2 = tvtk.Actor(mapper=mapper2, property=prop)
#
# fig.scene.add_actor(actor2)
filename = filename_9L4
stl_file = stl.Mesh.from_file(filename, calculate_normals=True)
A = T
A[0][3] = A[0][3] * 1000
A[1][3] = A[1][3] * 1000
A[2][3] = A[2][3] * 1000
stl_file.transform(A)
stl_file.save(filename[0:-4] + '_1.stl')
filename3 = filename[0:-4] + '_1.stl'
reader2 = tvtk.STLReader()
reader2.file_name = filename3
reader2.update()
mapper2 = tvtk.PolyDataMapper()
configure_input(mapper2, reader2.output)
prop = tvtk.Property(opacity=0.3, color=(1,0,0))
actor2 = tvtk.Actor(mapper=mapper2, property=prop)
fig.scene.add_actor(actor2)
src = mlab.pipeline.scalar_field(data)
# Our data is not equally spaced in all directions:
src.spacing = spacing*np.array([-1,-1,1])
src.origin = origin*np.array([-1,-1,1])
src.update_image_data = True
#----------------------------------------------------------------------
# Display a cut plane of the raw data
ipy = mlab.pipeline.image_plane_widget(src, colormap='bone', plane_orientation='y_axes')
ipw = mlab.pipeline.image_plane_widget(src, colormap='bone',
plane_orientation='z_axes')
ipw.ipw.origin = np.array(pos[0:3])*1000 + y_vec * 18.5 - 50 * z_vec
ipw.ipw.point1 = np.array(pos[0:3])* 1000 - y_vec*18.5 - 50 * z_vec
ipw.ipw.point2 = np.array(pos[0:3])*1000 + y_vec * 18.5 + 100 * z_vec
ipw.ipw.update_placement()
sphere = pos[0:3]*1000
# mlab.points3d(sphere[0], sphere[1], sphere[2], 1, scale_factor=20, figure=fig, opacity=1, color=(1,0,0), mode = 'sphere')
# mlab.points3d(origin[0], origin[1], origin[2], 1, scale_factor=20, figure=fig, opacity=1, color=(0, 1, 0),
# mode='sphere')
mlab.view(-165, 32, 350, [143, 133, 73])
mlab.roll(180)
fig.scene.disable_render = False
fig.scene.reset_zoom()
#----------------------------------------------------------------------
# # To make the link between the Mayavi pipeline and the much more
# # complex VTK pipeline, we display both:
# mlab.show_pipeline(rich_view=False)
mlab.show()
if not os.path.isdir(os.path.join(os.path.split(fname)[0], modality[0:2]+'_Slices')):
os.makedirs(os.path.join(os.path.split(fname)[0], modality[0:2]+'_Slices'))
#
# LOC = ''.join([c for c in locations[loc_idx] if c.isupper()])
# LOC_f = LOC[-1] + LOC[0]
side_src = ipw.ipw._get_reslice_output()
write_data(side_src, os.path.join(os.path.join(os.path.split(fname)[0], modality[0:2]+'_Slices'), locations[loc_idx][0:-5]+'_'+modality[0:2]+'.vtk'))
if __name__ == "__main__":
dir = '/home/morrile2/Documents/MULTIS Data/MULTIS_invitro'
subj_path = 'CMULTIS002-2'
modality = 'MRI'
segment = 'UpperLeg'
main(os.path.join(dir, subj_path), modality, segment)