diff --git a/AutoscoperM/AutoscoperM.py b/AutoscoperM/AutoscoperM.py index 7698e8f..3848ad9 100644 --- a/AutoscoperM/AutoscoperM.py +++ b/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 diff --git a/AutoscoperM/AutoscoperMLib/IO.py b/AutoscoperM/AutoscoperMLib/IO.py index ecfe158..0d94c8c 100644 --- a/AutoscoperM/AutoscoperMLib/IO.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 04cbd43..adc32bb 100644 --- a/pyproject.toml +++ b/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 \ No newline at end of file +pylint.max-args = 7