ENH: Output TRA files on PV gen

AutoscoperM/AutoscoperM.py

@@ -505,7 +505,6 @@ def onGenerateVRG(self):
             tfmPath = os.path.join(mainOutputDir, tfmPath, "Origin2Dicom.tfm")
             tfmNode.Inverse()
             slicer.util.saveNode(tfmNode, tfmPath)
-            slicer.mrmlScene.RemoveNode(tfmNode)
         else:
             # Just export the first frame
             currentNode, _ = self.logic.getItemInSequence(volumeNode, 0)
@@ -516,7 +515,6 @@ def onGenerateVRG(self):
             tfmPath = os.path.join(mainOutputDir, tfmPath, "Origin2Dicom.tfm")
             tfmNode.Inverse()
             slicer.util.saveNode(tfmNode, tfmPath)
-            slicer.mrmlScene.RemoveNode(tfmNode)
 
             # Harden and remove the transform from the sequence
             for i in range(1, volumeNode.GetNumberOfDataNodes()):
@@ -589,6 +587,14 @@ def onGenerateVRG(self):
         if self.ui.removeVrgTmp.isChecked():
             shutil.rmtree(os.path.join(mainOutputDir, tmpDir))
 
+        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,
@@ -683,6 +689,7 @@ def onSegmentation(self):
                     currentVolumeNode,
                     self.ui.segGen_ThresholdSpinBox.value,
                     self.ui.segGen_marginSizeSpin.value,
+                    progressCallback=self.updateProgressBar,
                 )
                 progress = (i + 1) / numFrames * 100
                 self.ui.progressBar.setValue(progress)
@@ -691,12 +698,6 @@ def onSegmentation(self):
                     slicer.mrmlScene.RemoveNode(segmentationNode)
                     currentVolumeNode = self.logic.getNextItemInSequence(volumeNode)
 
-            segmentationNode = SubVolumeExtraction.automaticSegmentation(
-                volumeNode,
-                self.ui.segGen_ThresholdSpinBox.value,
-                self.ui.segGen_marginSizeSpin.value,
-                progressCallback=self.updateProgressBar,
-            )
         elif self.ui.segGen_fileRadioButton.isChecked():
             segmentationFileDir = self.ui.segGen_lineEdit.currentPath
             if not self.logic.validatePaths(segmentationFileDir=segmentationFileDir):
@@ -979,6 +980,7 @@ def saveSubVolumesFromSegmentation(
         outputDir: str,
         volumeSubDir: str = "Volumes",
         transformSubDir: str = "Transforms",
+        trackingSubDir: str = "Tracking",
         progressCallback: Optional[callable] = None,
     ) -> bool:
         """
@@ -1005,6 +1007,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()
@@ -1027,6 +1037,26 @@ def progressCallback(x):
             filename = os.path.join(outputDir, transformSubDir, segmentName + ".tfm")
             IO.writeTFMFile(filename, spacing, origin)
             self.showVolumeIn3D(segmentVolume)
+
+            bounds = self.GetRASBounds(segmentVolume)
+            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,57 @@ 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
+    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