Skip to content

Commit

Permalink
WIP: 4d
Browse files Browse the repository at this point in the history
  • Loading branch information
@NicerNewerCar
NicerNewerCar committed Dec 20, 2023
1 parent 619546b commit 9a29e76
Show file tree
Hide file tree
Showing 5 changed files with 271 additions and 81 deletions.
214 changes: 189 additions & 25 deletions AutoscoperM/AutoscoperM.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import glob
import logging
import os
import re
import shutil
import time
import zipfile
Expand Down Expand Up @@ -221,6 +222,7 @@ def setup(self):
self.ui.segmentationButton.connect("clicked(bool)", self.onSegmentation)

self.ui.loadPVButton.connect("clicked(bool)", self.onLoadPV)
self.ui.volumeSelector.connect("currentNodeChanged(vtkMRMLNode*)", self.logic.check4D)

# Default output directory
self.ui.mainOutputSelector.setCurrentPath(
Expand Down Expand Up @@ -278,6 +280,10 @@ def initializeParameterNode(self):
# so that when the scene is saved and reloaded, these settings are restored.

self.setParameterNode(self.logic.getParameterNode())
if self.ui.volumeSelector.currentNode() is not None:
self.logic.check4D(self.ui.volumeSelector.currentNode())
if self.ui.mVRG_markupSelector.currentNode() is not None:
self.onMarkupNodeChanged(self.ui.mVRG_markupSelector.currentNode())

# Select default input nodes if nothing is selected yet to save a few clicks for the user
# NA
Expand Down Expand Up @@ -469,6 +475,7 @@ def onGenerateVRG(self):
nPossibleCameras = self.ui.posCamSpin.value
nOptimizedCameras = self.ui.optCamSpin.value
tmpDir = self.ui.vrgTempDir.text
tfmPath = self.ui.tfmSubDir.text
cameraSubDir = self.ui.cameraSubDir.text
vrgSubDir = self.ui.vrgSubDir.text
if not self.logic.validateInputs(
Expand All @@ -479,6 +486,7 @@ def onGenerateVRG(self):
nPossibleCameras=nPossibleCameras,
nOptimizedCameras=nOptimizedCameras,
tmpDir=tmpDir,
tfmPath=tfmPath,
cameraSubDir=cameraSubDir,
vrgSubDir=vrgSubDir,
):
Expand All @@ -491,8 +499,17 @@ def onGenerateVRG(self):
logging.error("Failed to generate VRG: more optimized cameras than possible cameras")
return

# center the volume
self.logic.createPathsIfNotExists(os.path.join(mainOutputDir, tfmPath))
self.logic.centerVolume(volumeNode, os.path.join(mainOutputDir, tfmPath))

numFrames = 1
currentNode = volumeNode
if self.logic.is_4d:
numFrames = volumeNode.GetNumberOfDataNodes()
currentNode = self.logic.getItemInSequence(volumeNode, 0)
bounds = [0] * 6
volumeNode.GetBounds(bounds)
currentNode.GetBounds(bounds)

# Generate all possible camera positions
camOffset = self.ui.camOffSetSpin.value
Expand All @@ -503,14 +520,18 @@ def onGenerateVRG(self):
self.updateProgressBar(10)

# Generate initial VRG for each camera
self.logic.generateVRGForCameras(
cameras,
volumeNode,
os.path.join(mainOutputDir, tmpDir),
width,
height,
progressCallback=self.updateProgressBar,
)
for i in range(numFrames):
filename = self.logic.cleanFilename(currentNode.GetName(), i)
self.logic.generateVRGForCameras(
cameras,
currentNode,
os.path.join(mainOutputDir, tmpDir),
[width, height],
filename=filename,
progressCallback=self.updateProgressBar,
)

currentNode = self.logic.getNextItemInSequence(volumeNode) if self.logic.is_4d else currentNode

# Optimize the camera positions
bestCameras = RadiographGeneration.optimizeCameras(
Expand Down Expand Up @@ -696,14 +717,36 @@ def onManualVRGGen(self):
if self.logic.vrgManualCameras is None:
self.onMarkupNodeChanged(markupsNode) # create the cameras

self.logic.generateVRGForCameras(
self.logic.vrgManualCameras,
volumeNode,
os.path.join(mainOutputDir, vrgDir),
width,
height,
progressCallback=self.updateProgressBar,
)
# Check if the volume is centered at the origin
bounds = [0] * 6
if self.logic.is_4d:
# get the bounds of the first frame
volumeNode.GetNthDataNode(0).GetRASBounds(bounds)
else:
volumeNode.GetRASBounds(bounds)

center = [(bounds[0] + bounds[1]) / 2, (bounds[2] + bounds[3]) / 2, (bounds[4] + bounds[5]) / 2]
center = [round(x) for x in center]
if center != [0, 0, 0]:
logging.warning("Volume is not centered at the origin. This may cause issues with Autoscoper.")

numFrames = 1
currentNode = volumeNode
if self.logic.is_4d:
numFrames = volumeNode.GetNumberOfDataNodes()
currentNode = self.logic.getItemInSequence(volumeNode, 0)

for i in range(numFrames):
filename = self.logic.cleanFilename(currentNode.GetName(), i)
self.logic.generateVRGForCameras(
self.logic.vrgManualCameras,
currentNode,
os.path.join(mainOutputDir, vrgDir),
[width, height],
filename=filename,
progressCallback=self.updateProgressBar,
)
currentNode = self.logic.getNextItemInSequence(volumeNode) if self.logic.is_4d else currentNode

self.updateProgressBar(100)

Expand All @@ -726,6 +769,7 @@ def onMarkupNodeChanged(self, node):
# get the volume nodes
volumeNode = self.ui.volumeSelector.currentNode()
self.logic.validateInputs(volumeNode=volumeNode)
volumeNode = self.logic.getItemInSequence(volumeNode, 0) if self.logic.is_4d else volumeNode
bounds = [0] * 6
volumeNode.GetBounds(bounds)
self.logic.vrgManualCameras = RadiographGeneration.generateCamerasFromMarkups(
Expand Down Expand Up @@ -773,6 +817,7 @@ def __init__(self):
self.AutoscoperProcess.setProcessChannelMode(qt.QProcess.ForwardedChannels)
self.AutoscoperSocket = None
self.vrgManualCameras = None
self.is_4d = False

def setDefaultParameters(self, parameterNode):
"""
Expand Down Expand Up @@ -1047,18 +1092,21 @@ def generateVRGForCameras(
cameras: list[RadiographGeneration.Camera],
volumeNode: slicer.vtkMRMLVolumeNode,
outputDir: str,
width: int,
height: int,
progressCallback: Optional[callable] = None,
size: list[int],
filename: str,
progressCallback=None,
) -> None:
"""
Generates VRG files for each camera in the cameras list
:param cameras: list of cameras
:param volumeNode: volume node
:param outputDir: output directory
:param width: width of the radiographs
:param height: height of the radiographs
:type outputDir: str
:param size: size of the VRG
:type size: list[int]
:param filename: filename of the VRG
:type filename: str
:param progressCallback: progress callback, defaults to None
"""
self.createPathsIfNotExists(outputDir)
Expand Down Expand Up @@ -1091,6 +1139,7 @@ def progressCallback(x):
cliModule = slicer.modules.virtualradiographgeneration
cliNodes = []
for cam in cameras:
logging.info(f"Generating {filename}.tif for {cam.id}")
cameraDir = os.path.join(outputDir, f"cam{cam.id}")
self.createPathsIfNotExists(cameraDir)
camera = cam.vtkCamera
Expand All @@ -1101,9 +1150,9 @@ def progressCallback(x):
"cameraViewUp": [camera.GetViewUp()[0], camera.GetViewUp()[1], camera.GetViewUp()[2]],
"cameraViewAngle": camera.GetViewAngle(),
"clippingRange": [camera.GetClippingRange()[0], camera.GetClippingRange()[1]],
"outputWidth": width,
"outputHeight": height,
"outputFileName": os.path.join(cameraDir, "1.tif"),
"width": size[0],
"height": size[1],
"outputFileName": os.path.join(cameraDir, f"{filename}.tif"),
}
cliNode = slicer.cli.run(cliModule, None, parameters) # run asynchronously
cliNodes.append(cliNode)
Expand Down Expand Up @@ -1194,3 +1243,118 @@ def convertNodeToData(self, volumeNode: slicer.vtkMRMLVolumeNode) -> vtk.vtkImag
imageData = imageReslice.GetOutput()

return imageData

def check4D(self, node: slicer.vtkMRMLNode) -> bool:
"""
Checks if the volume is 4D
"""
self.is_4d = type(node) == slicer.vtkMRMLSequenceNode

def centerVolume(self, volumeNode: slicer.vtkMRMLVolumeNode, transformPath: str) -> None:
"""
A requirement for Autoscoper is that the center of the volume is at the origin.
This method will center the volume and save the transform to the transformPath
:param volumeNode: volume node
:type volumeNode: slicer.vtkMRMLVolumeNode
:param transformPath: path to save the transform to
:type transformPath: str
:return: None
"""

# Get the bounds of the volume
bounds = [0] * 6
if self.is_4d:
volumeNode.GetNthDataNode(0).GetRASBounds(bounds)
else:
volumeNode.GetRASBounds(bounds)

# Get the center of the volume
center = [0] * 3
for i in range(3):
center[i] = (bounds[i * 2] + bounds[i * 2 + 1]) / 2

center_rounded = [round(x) for x in center] # don't want to move the volume if its off by a small amount
if center_rounded == [0, 0, 0]:
return # Already centered
# Create a transform node
transformNode = slicer.vtkMRMLTransformNode()
transformNode.SetName("CenteringTransform")
slicer.mrmlScene.AddNode(transformNode)

# Get the transform matrix
matrix = vtk.vtkMatrix4x4()

# Move the center of the volume to the origin
matrix.SetElement(0, 3, -center[0])
matrix.SetElement(1, 3, -center[1])
matrix.SetElement(2, 3, -center[2])

# Set the transform matrix
transformNode.SetMatrixTransformToParent(matrix)

# Apply the transform to the volume
num_frames = 1
curVol = volumeNode
if self.is_4d:
num_frames = volumeNode.GetNumberOfDataNodes()
for i in range(num_frames):
if self.is_4d:
curVol = self.getItemInSequence(volumeNode, i)
curVol.SetAndObserveTransformNodeID(transformNode.GetID())

# Harden the transform
slicer.modules.transforms.logic().hardenTransform(curVol)
curVol.SetAndObserveTransformNodeID(None)

# # Invert and save the transform
matrix.Invert()
transformNode.SetMatrixTransformToParent(matrix)
slicer.util.exportNode(transformNode, os.path.join(transformPath, "Origin2DICOMCenter.tfm"))

def getItemInSequence(self, sequenceNode: slicer.vtkMRMLSequenceNode, idx: int) -> slicer.vtkMRMLNode:
"""
Returns the item at the specified index in the sequence node
:param sequenceNode: sequence node
:type sequenceNode: slicer.vtkMRMLSequenceNode
:param idx: index
:type idx: int
:return: item at the specified index
:rtype: slicer.vtkMRMLNode
"""
if type(sequenceNode) != slicer.vtkMRMLSequenceNode:
logging.error("[AutoscoperM.logic.getItemInSequence] sequenceNode must be a sequence node")
return None

if idx >= sequenceNode.GetNumberOfDataNodes():
logging.error(f"[AutoscoperM.logic.getItemInSequence] index {idx} is out of range")
return None

browserNode = slicer.modules.sequences.logic().GetFirstBrowserNodeForSequenceNode(sequenceNode)
browserNode.SetSelectedItemNumber(idx)
return browserNode.GetProxyNode(sequenceNode)

def getNextItemInSequence(self, sequenceNode: slicer.vtkMRMLSequenceNode) -> slicer.vtkMRMLNode:
"""
Returns the next item in the sequence
:param sequenceNode: sequence node
:type sequenceNode: slicer.vtkMRMLSequenceNode
:return: next item in the sequence
:rtype: slicer.vtkMRMLNode
"""
if type(sequenceNode) != slicer.vtkMRMLSequenceNode:
logging.error("[AutoscoperM.logic.getNextItemInSequence] sequenceNode must be a sequence node")
return None

browserNode = slicer.modules.sequences.logic().GetFirstBrowserNodeForSequenceNode(sequenceNode)
browserNode.SelectNextItem()
return browserNode.GetProxyNode(sequenceNode)

def cleanFilename(self, volumeName: str, index: Optional[int] = None) -> str:
filename = (
re.sub(r"\s+", "_", f"{volumeName}_{index}") if index is not None else re.sub(r"\s+", "_", f"{volumeName}")
) # Remove spaces
filename = re.sub(r"[^\w]", "", filename) # Remove non alphanumeric characters
return re.sub(r"__+", "_", filename) # Remove double or more underscores
5 changes: 2 additions & 3 deletions AutoscoperM/AutoscoperMLib/RadiographGeneration.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -221,7 +221,6 @@ def optimizeCameras(
:return: Optimized cameras
:rtype: list[Camera]
"""
import glob
import os

if not progressCallback:
Expand All @@ -234,11 +233,11 @@ def progressCallback(_x):
cliNodes = []
for i in range(len(cameras)):
camera = cameras[i]
vrgFName = glob.glob(os.path.join(cameraDir, f"cam{camera.id}", "*.tif"))[0]
vrgDirName = os.path.join(cameraDir, f"cam{camera.id}")
cliNode = slicer.cli.run(
cliModule,
None,
{"whiteRadiographFileName": vrgFName},
{"whiteRadiographDirName": vrgDirName},
wait_for_completion=False,
)
cliNodes.append(cliNode)
Expand Down
3 changes: 2 additions & 1 deletion AutoscoperM/Resources/UI/AutoscoperM.ui
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
<item>
<widget class="QTabWidget" name="tabWidget">
<property name="currentIndex">
<number>0</number>
<number>1</number>
</property>
<widget class="QWidget" name="AutoscoperControlTab">
<attribute name="title">
Expand Down Expand Up @@ -151,6 +151,7 @@
<property name="nodeTypes">
<stringlist notr="true">
<string>vtkMRMLScalarVolumeNode</string>
<string>vtkMRMLSequenceNode</string>
</stringlist>
</property>
<property name="hideChildNodeTypes">
Expand Down
Loading

0 comments on commit 9a29e76

Please sign in to comment.