Add stylus/probe visualization

invesalius/data/bases.py

@@ -126,9 +126,7 @@ def calculate_fre(fiducials_raw, fiducials, ref_mode_id, m_change, m_icp=None):
 
     dist = np.zeros([3, 1])
     for i in range(0, 6, 2):
-        p_m, _ = dcr.corregistrate_dynamic(
-            (m_change, 0), fiducials_raw[i : i + 2], ref_mode_id, icp
-        )
+        p_m, _ = dcr.corregistrate_probe(m_change, None, fiducials_raw[i : i + 2], ref_mode_id, icp)
         dist[int(i / 2)] = np.sqrt(np.sum(np.power((p_m[:3] - fiducials[int(i / 2), :]), 2)))
 
     return float(np.sqrt(np.sum(dist**2) / 3))
@@ -199,19 +197,16 @@ def object_registration(fiducials, orients, coord_raw, m_change):
         fids_raw[ic, :] = dco.dynamic_reference_m2(coords[ic, :], coords[3, :])[:3]
 
     # compute initial alignment of probe fixed in the object in source frame
-
-    # XXX: Some duplicate processing is done here: the Euler angles are calculated once by
-    #      the lines below, and then again in dco.coordinates_to_transformation_matrix.
-    #
-    a, b, g = np.radians(coords[3, 3:])
-    r_s0_raw = tr.euler_matrix(a, b, g, axes="rzyx")
-
     s0_raw = dco.coordinates_to_transformation_matrix(
         position=coords[3, :3],
         orientation=coords[3, 3:],
         axes="rzyx",
     )
 
+    # copy rotation submatrix from s0_raw
+    r_s0_raw = np.eye(4)
+    r_s0_raw[:3, :3] = s0_raw[:3, :3]
+
     # compute change of basis for object fiducials in source frame
     base_obj_raw, q_obj_raw = base_creation(fids_raw[:3, :3])
     r_obj_raw = np.identity(4)

invesalius/data/coregistration.py

@@ -48,7 +48,6 @@ def object_marker_to_center(coord_raw, obj_ref_mode, t_obj_raw, s0_raw, r_s0_raw
     :return: 4 x 4 numpy double array
     :rtype: numpy.ndarray
     """
-
     as1, bs1, gs1 = np.radians(coord_raw[obj_ref_mode, 3:])
     r_probe = tr.euler_matrix(as1, bs1, gs1, "rzyx")
     t_probe_raw = tr.translation_matrix(coord_raw[obj_ref_mode, :3])
@@ -171,13 +170,14 @@ def image_to_tracker(m_change, coord_raw, target, icp, obj_data):
     return m_target_in_tracker
 
 
-def corregistrate_object_dynamic(inp, coord_raw, ref_mode_id, icp):
-    m_change, obj_ref_mode, t_obj_raw, s0_raw, r_s0_raw, s0_dyn, m_obj_raw, r_obj_img = inp
+def corregistrate_probe(m_change, r_stylus, coord_raw, ref_mode_id, icp=[None, None]):
+    if r_stylus is None:
+        # utils.debug("STYLUS ORIENTATION NOT DEFINED!")
+        r_stylus = np.eye(3)
 
-    # transform raw marker coordinate to object center
-    m_probe = object_marker_to_center(coord_raw, obj_ref_mode, t_obj_raw, s0_raw, r_s0_raw)
+    m_probe = compute_marker_transformation(coord_raw, 0)
 
-    # transform object center to reference marker if specified as dynamic reference
+    # transform probe to reference system if dynamic ref_mode
     if ref_mode_id:
         m_probe_ref = object_to_reference(coord_raw, m_probe)
     else:
@@ -186,71 +186,109 @@ def corregistrate_object_dynamic(inp, coord_raw, ref_mode_id, icp):
     # invert y coordinate
     m_probe_ref[2, -1] = -m_probe_ref[2, -1]
 
-    # corregistrate from tracker to image space
-    m_img = tracker_to_image(m_change, m_probe_ref, r_obj_img, m_obj_raw, s0_dyn)
+    # translate m_probe_ref from tracker to image space
+    m_img = m_change @ m_probe_ref
     m_img = apply_icp(m_img, icp)
 
+    # Rotate from trk system where stylus points in x-axis to vtk-system where stylus points in y-axis
+    R = tr.euler_matrix(*np.radians([0, 0, -90]), axes="rxyz")[:3, :3]
+
+    # rotate m_probe_ref from tracker to image space
+    r_img = r_stylus @ R @ m_probe_ref[:3, :3] @ np.linalg.inv(R)
+    m_img[:3, :3] = r_img[:3, :3]
+
     # compute rotation angles
-    angles = tr.euler_from_matrix(m_img, axes="sxyz")
+    angles = np.degrees(tr.euler_from_matrix(m_img, axes="sxyz"))
 
     # create output coordinate list
     coord = (
         m_img[0, -1],
         m_img[1, -1],
         m_img[2, -1],
-        np.degrees(angles[0]),
-        np.degrees(angles[1]),
-        np.degrees(angles[2]),
+        angles[0],
+        angles[1],
+        angles[2],
     )
 
     return coord, m_img
 
 
-def compute_marker_transformation(coord_raw, obj_ref_mode):
-    m_probe = dco.coordinates_to_transformation_matrix(
-        position=coord_raw[obj_ref_mode, :3],
-        orientation=coord_raw[obj_ref_mode, 3:],
-        axes="rzyx",
+def corregistrate_object_dynamic(inp, coord_raw, i_obj, icp):
+    """
+    Corregistrate the object at coord_raw[i_obj] in dynamic ref_mode
+    """
+    m_change, obj_ref_mode, t_obj_raw, s0_raw, r_s0_raw, s0_dyn, m_obj_raw, r_obj_img = inp
+
+    # transform raw marker coordinate to object center
+    m_probe = object_marker_to_center(coord_raw, i_obj, t_obj_raw, s0_raw, r_s0_raw)
+
+    # transform object center to reference marker
+    m_probe_ref = object_to_reference(coord_raw, m_probe)
+
+    # invert y coordinate
+    m_probe_ref[2, -1] = -m_probe_ref[2, -1]
+
+    # corregistrate from tracker to image space
+    m_img = tracker_to_image(m_change, m_probe_ref, r_obj_img, m_obj_raw, s0_dyn)
+    m_img = apply_icp(m_img, icp)
+
+    # compute rotation angles
+    angles = np.degrees(tr.euler_from_matrix(m_img, axes="sxyz"))
+
+    # create output coordinate list
+    coord = (
+        m_img[0, -1],
+        m_img[1, -1],
+        m_img[2, -1],
+        angles[0],
+        angles[1],
+        angles[2],
     )
-    return m_probe
 
+    return coord, m_img
 
-def corregistrate_dynamic(inp, coord_raw, ref_mode_id, icp):
-    m_change, obj_ref_mode = inp
 
-    # transform raw marker coordinate to object center
-    m_probe = compute_marker_transformation(coord_raw, obj_ref_mode)
+def corregistrate_object_static(inp, coord_raw, i_obj, icp):
+    """
+    Corregistrate the object at coord_raw[i_obj] in static ref_mode
+    """
+    m_change, obj_ref_mode, t_obj_raw, s0_raw, r_s0_raw, s0_dyn, m_obj_raw, r_obj_img = inp
 
-    # transform object center to reference marker if specified as dynamic reference
-    if ref_mode_id:
-        m_ref = compute_marker_transformation(coord_raw, 1)
-        m_probe_ref = np.linalg.inv(m_ref) @ m_probe
-    else:
-        m_probe_ref = m_probe
+    # transform raw marker coordinate to object center
+    m_probe = object_marker_to_center(coord_raw, i_obj, t_obj_raw, s0_raw, r_s0_raw)
 
     # invert y coordinate
-    m_probe_ref[2, -1] = -m_probe_ref[2, -1]
+    m_probe[2, -1] = -m_probe[2, -1]
 
     # corregistrate from tracker to image space
-    m_img = m_change @ m_probe_ref
+    m_img = tracker_to_image(m_change, m_probe, r_obj_img, m_obj_raw, s0_dyn)
     m_img = apply_icp(m_img, icp)
 
     # compute rotation angles
-    angles = tr.euler_from_matrix(m_img, axes="sxyz")
+    angles = np.degrees(tr.euler_from_matrix(m_img, axes="sxyz"))
 
     # create output coordinate list
     coord = (
         m_img[0, -1],
         m_img[1, -1],
         m_img[2, -1],
-        np.degrees(angles[0]),
-        np.degrees(angles[1]),
-        np.degrees(angles[2]),
+        angles[0],
+        angles[1],
+        angles[2],
     )
 
     return coord, m_img
 
 
+def compute_marker_transformation(coord_raw, obj_ref_mode):
+    m_probe = dco.coordinates_to_transformation_matrix(
+        position=coord_raw[obj_ref_mode, :3],
+        orientation=coord_raw[obj_ref_mode, 3:],
+        axes="rzyx",
+    )
+    return m_probe
+
+
 def apply_icp(m_img, icp):
     use_icp, m_icp = icp
     if use_icp:
@@ -295,6 +333,7 @@ def __init__(
         self,
         ref_mode_id,
         tracker,
+        n_coils,
         coreg_data,
         view_tracts,
         queues,
@@ -308,6 +347,7 @@ def __init__(
         threading.Thread.__init__(self, name="CoordCoregObject")
         self.ref_mode_id = ref_mode_id
         self.tracker = tracker
+        self.n_coils = n_coils
         self.coreg_data = coreg_data
         self.coord_queue = queues[0]
         self.view_tracts = view_tracts
@@ -336,9 +376,20 @@ def __init__(
 
     def run(self):
         coreg_data = self.coreg_data
-        view_obj = 1
+        corregistrate_object = (
+            corregistrate_object_dynamic if self.ref_mode_id else corregistrate_object_static
+        )
+
+        # compute n_coils_effective, the no. of coils to actually process:
+        # check how many coords we get from tracker (-2 for probe & head)
+        n_coils_trk = self.tracker.TrackerCoordinates.GetCoordinates()[0].shape[0] - 2
+
+        obj_ref_mode = coreg_data[2]
+        # if obj_ref_mode=0: only coregister one coil
+        # else: process the other (n_coils - 1) coils too
+        # min(obj_ref_mode, 1) = 0 if obj_ref_mode==0 else 1
+        n_coils_effective = min(n_coils_trk, 1 + min(obj_ref_mode, 1) * (self.n_coils - 1))
 
-        # print('CoordCoreg: event {}'.format(self.event.is_set()))
         while not self.event.is_set():
             try:
                 if not self.icp_queue.empty():
@@ -347,12 +398,27 @@ def run(self):
                 if not self.object_at_target_queue.empty():
                     self.target_flag = self.object_at_target_queue.get_nowait()
 
-                # print(f"Set the coordinate")
                 coord_raw, marker_visibilities = self.tracker.TrackerCoordinates.GetCoordinates()
-                coord, m_img = corregistrate_object_dynamic(
-                    coreg_data, coord_raw, self.ref_mode_id, [self.use_icp, self.m_icp]
+
+                # m_change = coreg_data[1], r_stylus = coreg_data[0]
+                coord_probe, m_img_probe = corregistrate_probe(
+                    coreg_data[1], coreg_data[0], coord_raw, self.ref_mode_id
+                )
+                coord_coil, m_img_coil = corregistrate_object(
+                    coreg_data[1:], coord_raw, obj_ref_mode, [self.use_icp, self.m_icp]
                 )
 
+                coords = [coord_probe, coord_coil]
+                m_imgs = [m_img_probe, m_img_coil]
+
+                # the possible other coils are i_obj=3 onwards at coord_raw[i_obj]
+                for i_obj in range(3, n_coils_effective + 2):
+                    coord_coil, m_img_coil = corregistrate_object(
+                        coreg_data[1:], coord_raw, i_obj, [self.use_icp, self.m_icp]
+                    )
+                    coords.append(coord_coil)
+                    m_imgs.append(m_img_coil)
+
                 # XXX: This is not the best place to do the logic related to approaching the target when the
                 #      debug tracker is in use. However, the trackers (including the debug trackers) operate in
                 #      the tracker space where it is hard to make the tracker approach the target in the image space.
@@ -363,77 +429,26 @@ def run(self):
                 #
                 if self.tracker_id == const.DEBUGTRACKAPPROACH and self.target is not None:
                     if self.last_coord is None:
-                        self.last_coord = np.array(coord)
+                        self.last_coord = np.array(coord_coil)
                     else:
                         coord = self.last_coord + (self.target - self.last_coord) * 0.05
+                        coords[1] = coord
                         self.last_coord = coord
 
                     angles = [np.radians(coord[3]), np.radians(coord[4]), np.radians(coord[5])]
                     translate = coord[0:3]
-                    m_img = tr.compose_matrix(angles=angles, translate=translate)
+                    m_imgs[1] = tr.compose_matrix(angles=angles, translate=translate)
 
-                m_img_flip = m_img.copy()
-                m_img_flip[1, -1] = -m_img_flip[1, -1]
-                # self.pipeline.set_message(m_img_flip)
+                self.coord_queue.put_nowait([coords, marker_visibilities, m_imgs])
 
-                self.coord_queue.put_nowait([coord, marker_visibilities, m_img, view_obj])
-                # print('CoordCoreg: put {}'.format(count))
-                # count += 1
+                coord = coords[1]  # main coil
+                m_img = m_imgs[1]
+                # LUKATODO: should coord = coords[track_coil]
+                # should the stuff below ever be done for stylus, but not coil?
 
-                if self.view_tracts:
-                    self.coord_tracts_queue.put_nowait(m_img_flip)
-                if self.e_field_loaded:
-                    self.efield_queue.put_nowait([m_img, coord])
-                if not self.icp_queue.empty():
-                    self.icp_queue.task_done()
-            except queue.Full:
-                pass
-            # The sleep has to be in both threads
-            sleep(self.sle)
-
-
-class CoordinateCorregistrateNoObject(threading.Thread):
-    def __init__(
-        self, ref_mode_id, tracker, coreg_data, view_tracts, queues, event, sle, icp, e_field_loaded
-    ):
-        threading.Thread.__init__(self, name="CoordCoregNoObject")
-        self.ref_mode_id = ref_mode_id
-        self.tracker = tracker
-        self.coreg_data = coreg_data
-        self.coord_queue = queues[0]
-        self.view_tracts = view_tracts
-        self.coord_tracts_queue = queues[1]
-        self.event = event
-        self.sle = sle
-        self.icp_queue = queues[2]
-        self.use_icp = icp.use_icp
-        self.m_icp = icp.m_icp
-        self.efield_queue = queues[3]
-        self.e_field_loaded = e_field_loaded
-
-    def run(self):
-        coreg_data = self.coreg_data
-        view_obj = 0
-
-        # print('CoordCoreg: event {}'.format(self.event.is_set()))
-        while not self.event.is_set():
-            try:
-                if self.icp_queue.empty():
-                    None
-                else:
-                    self.use_icp, self.m_icp = self.icp_queue.get_nowait()
-                # print(f"Set the coordinate")
-                # print(self.icp, self.m_icp)
-                coord_raw, marker_visibilities = self.tracker.TrackerCoordinates.GetCoordinates()
-                coord, m_img = corregistrate_dynamic(
-                    coreg_data, coord_raw, self.ref_mode_id, [self.use_icp, self.m_icp]
-                )
-                # print("Coord: ", coord)
                 m_img_flip = m_img.copy()
                 m_img_flip[1, -1] = -m_img_flip[1, -1]
 
-                self.coord_queue.put_nowait([coord, marker_visibilities, m_img, view_obj])
-
                 if self.view_tracts:
                     self.coord_tracts_queue.put_nowait(m_img_flip)
                 if self.e_field_loaded: