ENH: Improve and export Slicer2AUT transform

AutoscoperM/AutoscoperM.py

@@ -1061,25 +1061,21 @@ def progressCallback(x):
             IO.writeTFMFile(filename, spacing, origin)
             self.showVolumeIn3D(segmentVolume)
 
-            bounds = [0] * 6
-            segmentVolume.GetRASBounds(bounds)
-            segmentVolumeSize = [abs(bounds[i + 1] - bounds[i]) for i in range(0, len(bounds), 2)]
-
             # Write TRA
             tfm = vtk.vtkMatrix4x4()
             tfm.SetElement(0, 3, origin[0])
             tfm.SetElement(1, 3, origin[1])
             tfm.SetElement(2, 3, origin[2])
 
-            IO.createTRAFile(
-                volName=segmentName,
-                trialName=None,
-                outputDir=outputDir,
-                trackingsubDir=trackingSubDir,
-                volSize=segmentVolumeSize,
-                Origin2DicomTransformFile=origin2DicomTransformFile,
-                transform=tfm,
+            if origin2DicomTransformFile is not None:
+                tfm = self.applyOrigin2DicomTransform(tfm, origin2DicomTransformFile, ApplyDicom2Origin=True)
+            slicer2autoscoperFilename = self.createSlicer2AutoscoperTransform(
+                segmentVolume, os.path.join(outputDir, transformSubDir)
             )
+            tfm = self.applySlicer2AutoscoperTransform(tfm, slicer2autoscoperFilename, ApplyAutoscoper2Slicer=False)
+            filename = f"{segmentName}.tra"
+            filename = os.path.join(outputDir, trackingSubDir, filename)
+            IO.writeTRA(filename, [tfm])
 
             # update progress bar
             progressCallback((idx + 1) / numSegments * 100)
@@ -1427,3 +1423,79 @@ def IsVolumeCentered(node: Union[slicer.vtkMRMLVolumeNode, slicer.vtkMRMLSequenc
         if AutoscoperMLogic.IsSequenceVolume(node):
             return AutoscoperMLogic.getItemInSequence(node, 0)[0].IsCentered()
         return node.IsCentered()
+
+    @staticmethod
+    def applyOrigin2DicomTransform(
+        transform: vtk.vtkMatrix4x4, Origin2DicomTransformFile: str, ApplyDicom2Origin: bool = False
+    ) -> vtk.vtkMatrix4x4:
+        origin2DicomTransformNode = slicer.util.loadNodeFromFile(Origin2DicomTransformFile)
+
+        if ApplyDicom2Origin:
+            origin2DicomTransformNode.Inverse()
+
+        origin2DicomTransformMatrix = vtk.vtkMatrix4x4()
+        origin2DicomTransformNode.GetMatrixTransformToParent(origin2DicomTransformMatrix)
+
+        vtk.vtkMatrix4x4.Multiply4x4(origin2DicomTransformMatrix, transform, transform)
+
+        slicer.mrmlScene.RemoveNode(origin2DicomTransformNode)
+        return transform
+
+    @staticmethod
+    def applySlicer2AutoscoperTransform(
+        transform: vtk.vtkMatrix4x4, Slicer2AutoscoperFile: str, ApplyAutoscoper2Slicer: bool = False
+    ) -> vtk.vtkMatrix4x4:
+        """Utility function for converting a transform between the Slicer and Autoscoper coordinate systems."""
+        from itertools import product
+
+        slicer2AutoscoperNode = slicer.util.loadNodeFromFile(Slicer2AutoscoperFile)
+
+        if ApplyAutoscoper2Slicer:
+            slicer2AutoscoperNode.Inverse()
+
+        slicer2Autoscoper = vtk.vtkMatrix4x4()
+        slicer2AutoscoperNode.GetMatrixTransformToParent(slicer2Autoscoper)
+
+        # Extract the rotation matrices so we are not affecting the translation vector
+        transformR = vtk.vtkMatrix3x3()
+        slicer2autoscoperR = vtk.vtkMatrix3x3()
+        for i, j in product(range(3), range(3)):
+            transformR.SetElement(i, j, transform.GetElement(i, j))
+            slicer2autoscoperR.SetElement(i, j, slicer2Autoscoper.GetElement(i, j))
+
+        vtk.vtkMatrix3x3.Multiply3x3(slicer2autoscoperR, transformR, transformR)
+
+        for i, j in product(range(3), range(3)):
+            transform.SetElement(i, j, transformR.GetElement(i, j))
+
+        # Apply the translation vector
+        for i in range(3):
+            transform.SetElement(i, 3, transform.GetElement(i, 3) + slicer2Autoscoper.GetElement(i, 3))
+
+        slicer.mrmlScene.RemoveNode(slicer2AutoscoperNode)
+
+        return transform
+
+    @staticmethod
+    def createSlicer2AutoscoperTransform(volumeNode: slicer.vtkMRMLVolumeNode, exportDir: str):
+        # Slicer 2 Autoscoper Transform
+        # https://github.com/BrownBiomechanics/Autoscoper/issues/280
+        bounds = [0] * 6
+        volumeNode.GetRASBounds(bounds)
+        volSize = [abs(bounds[i + 1] - bounds[i]) for i in range(0, len(bounds), 2)]
+
+        slicer2autoscoper = vtk.vtkMatrix4x4()
+        slicer2autoscoper.Identity()
+        # Rotation matrix for a 180 x-axis rotation
+        slicer2autoscoper.SetElement(1, 1, -slicer2autoscoper.GetElement(1, 1))
+        slicer2autoscoper.SetElement(1, 2, -slicer2autoscoper.GetElement(1, 2))
+        slicer2autoscoper.SetElement(2, 2, -slicer2autoscoper.GetElement(2, 2))
+        slicer2autoscoper.SetElement(0, 3, -volSize[0])  # Offset -X
+
+        # Write out Slicer2AUT
+        slicer2autoscoperNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLLinearTransformNode")
+        slicer2autoscoperNode.SetMatrixTransformToParent(slicer2autoscoper)
+        slicer2autoscoperFilename = os.path.join(exportDir, f"{volumeNode.GetName()}-Slicer2AUT.tfm")
+        slicer.util.saveNode(slicer2autoscoperNode, slicer2autoscoperFilename)
+        slicer.mrmlScene.RemoveNode(slicer2autoscoperNode)
+        return slicer2autoscoperFilename

AutoscoperM/AutoscoperMLib/IO.py

@@ -1,6 +1,7 @@
 import glob
 import logging
 import os
+from itertools import product
 
 import numpy as np
 import slicer
@@ -255,56 +256,13 @@ def writeTFMFile(filename: str, spacing: list[float], origin: list[float]):
     slicer.mrmlScene.RemoveNode(transformNode)
 
 
-def createTRAFile(
-    volName: str,
-    trialName: str,
-    outputDir: str,
-    trackingsubDir: str,
-    volSize: list[float],
-    Origin2DicomTransformFile: str,
-    transform: vtk.vtkMatrix4x4,
-):
-    transformNode = slicer.vtkMRMLLinearTransformNode()
-    transformNode.SetMatrixTransformToParent(transform)
-    slicer.mrmlScene.AddNode(transformNode)
-
-    if Origin2DicomTransformFile is not None:
-        origin2DicomTransformNode = slicer.util.loadNodeFromFile(Origin2DicomTransformFile)
-        origin2DicomTransformNode.Inverse()
-        transformNode.SetAndObserveTransformNodeID(origin2DicomTransformNode.GetID())
-        transformNode.HardenTransform()
-        slicer.mrmlScene.RemoveNode(origin2DicomTransformNode)
-
-    filename = f"{trialName}_{volName}.tra" if trialName is not None else f"{volName}.tra"
-    filename = os.path.join(outputDir, trackingsubDir, filename)
-
-    if not os.path.exists(os.path.join(outputDir, trackingsubDir)):
-        os.mkdir(os.path.join(outputDir, trackingsubDir))
-
-    tfmMat = vtk.vtkMatrix4x4()
-    transformNode.GetMatrixTransformToParent(tfmMat)
-
-    writeTRA(filename, volSize, tfmMat)
-
-    slicer.mrmlScene.RemoveNode(transformNode)
-
-
-def writeTRA(filename: str, volSize: list[float], transform: vtk.vtkMatrix4x4):
-    # Slicer 2 Autoscoper Transform
-    # https://github.com/BrownBiomechanics/Autoscoper/issues/280
-    transform.SetElement(1, 1, -transform.GetElement(1, 1))  # Flip Y
-    transform.SetElement(2, 2, -transform.GetElement(2, 2))  # Flip Z
-
-    transform.SetElement(0, 3, transform.GetElement(0, 3) - volSize[0])  # Offset X
-
-    # Write TRA
-    rowwise = []
-    for i in range(4):  # Row
-        for j in range(4):  # Col
-            rowwise.append(str(transform.GetElement(i, j)))
-
-    with open(filename, "w+") as f:
-        f.write(",".join(rowwise))
+def writeTRA(fileName: str, transforms: list[vtk.vtkMatrix4x4]) -> None:
+    rowWiseStrings = []
+    for transform in transforms:
+        rowWiseStrings.append([str(transform.GetElement(i, j)) for i, j in product(range(4), range(4))])
+    with open(fileName, "w") as traFile:
+        for row in rowWiseStrings:
+            traFile.write(",".join(row) + "\n")
 
 
 def writeTemporyFile(filename: str, data: vtk.vtkImageData) -> str: