ENH: Output TRA files on PV gen

AutoscoperM/AutoscoperM.py

@@ -594,6 +594,14 @@ def onGenerateVRG(self):
 
         slicer.util.messageBox("Success!")
 
+        if not self.logic.IsSequenceVolume(volumeNode):
+            firstNode = volumeNode
+        else:
+            firstNode, _ = self.logic.getItemInSequence(volumeNode, 0)
+        firstNode.SetAndObserveTransformNodeID(tfmNode.GetID())
+        firstNode.HardenTransform()
+        slicer.mrmlScene.RemoveNode(tfmNode)
+
     def onGenerateConfig(self):
         """
         Generates a complete config file (including all partial volumes, radiographs,
@@ -995,6 +1003,7 @@ def saveSubVolumesFromSegmentation(
         outputDir: str,
         volumeSubDir: str = "Volumes",
         transformSubDir: str = "Transforms",
+        trackingSubDir: str = "Tracking",
         progressCallback: Optional[callable] = None,
     ) -> bool:
         """
@@ -1021,6 +1030,14 @@ def progressCallback(x):
         segmentIDs = vtk.vtkStringArray()
         segmentationNode.GetSegmentation().GetSegmentIDs(segmentIDs)
         numSegments = segmentIDs.GetNumberOfValues()
+
+        tfmFiles = glob.glob(os.path.join(outputDir, transformSubDir, "*.tfm"))
+        tfms = [tfm if os.path.basename(tfm).split(".")[0] == "Origin2Dicom" else None for tfm in tfmFiles]
+        try:
+            origin2DicomTransformFile = next(item for item in tfms if item is not None)
+        except StopIteration:
+            origin2DicomTransformFile = None
+
         for idx in range(numSegments):
             segmentID = segmentIDs.GetValue(idx)
             segmentName = segmentationNode.GetSegmentation().GetSegment(segmentID).GetName()
@@ -1043,6 +1060,27 @@ def progressCallback(x):
             filename = os.path.join(outputDir, transformSubDir, segmentName + ".tfm")
             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,
+            )
+
             # update progress bar
             progressCallback((idx + 1) / numSegments * 100)
         # Set the  volumeNode to be the active volume

AutoscoperM/AutoscoperMLib/IO.py

@@ -255,6 +255,58 @@ 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 writeTemporyFile(filename: str, data: vtk.vtkImageData) -> str:
     """
     Writes a temporary file to the slicer temp directory

pyproject.toml

@@ -32,6 +32,8 @@ extend-ignore = [
     "PIE790",  # unnecessary-pass
     "PLR0915", # too-many-statements
     "EXE001",  # Allow use of shebang at top of file. Needed for Scripted CLI
+    "PLR0913", # Too many arguments to function call
+    "PLR0912"  # Too many branches
 ]
 target-version = "py39"
 line-length = 120
@@ -50,4 +52,4 @@ isort.known-first-party = [
   "TrackingEvaluation",
   "TrackingEvaluationLib",
 ]
-pylint.max-args = 7
+pylint.max-args = 7