Skip to content

Commit 0ffb1ac

Browse files
axel-grcaxel-grc
committed
ENH: Add rtkspectralrooster python application
1 parent d684918 commit 0ffb1ac

5 files changed

Lines changed: 326 additions & 0 deletions

File tree

‎applications/rtkspectralrooster/rtkspectralrooster.py‎

Lines changed: 322 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,322 @@
1+
#!/usr/bin/env python
2+
import argparse
3+
import itk
4+
from itk import RTK as rtk
5+
6+
7+
def build_parser():
8+
parser = rtk.RTKArgumentParser(
9+
description=(
10+
"Reconstructs a 3D + material vector volume from a vector projection stack,"
11+
" alternating between conjugate gradient optimization and regularization."
12+
)
13+
)
14+
15+
parser.add_argument(
16+
"--geometry", "-g", help="XML geometry file name", required=True
17+
)
18+
parser.add_argument("--output", "-o", help="Output file name", required=True)
19+
parser.add_argument(
20+
"--niter", "-n", help="Number of main loop iterations", type=int, default=5
21+
)
22+
parser.add_argument(
23+
"--cgiter",
24+
help="Number of conjugate gradient nested iterations",
25+
type=int,
26+
default=4,
27+
)
28+
parser.add_argument(
29+
"--cudacg",
30+
help="Perform conjugate gradient calculations on GPU",
31+
action="store_true",
32+
)
33+
parser.add_argument("--input", "-i", help="Input volume (materials) file")
34+
parser.add_argument(
35+
"--projection", "-p", help="Vector projections file", required=True
36+
)
37+
parser.add_argument(
38+
"--nodisplaced",
39+
help="Disable the displaced detector filter",
40+
action="store_true",
41+
)
42+
43+
# Regularization
44+
parser.add_argument(
45+
"--nopositivity", help="Do not enforce positivity", action="store_true"
46+
)
47+
parser.add_argument("--tviter", help="TV iterations", type=int, default=10)
48+
parser.add_argument(
49+
"--gamma_space", help="Spatial TV regularization parameter", type=float
50+
)
51+
parser.add_argument("--threshold", help="Wavelets soft threshold", type=float)
52+
parser.add_argument("--order", help="Wavelets order", type=int, default=5)
53+
parser.add_argument("--levels", help="Wavelets levels", type=int, default=3)
54+
parser.add_argument(
55+
"--gamma_time", help="Temporal TV regularization parameter", type=float
56+
)
57+
parser.add_argument(
58+
"--lambda_time", help="Temporal L0 regularization parameter", type=float
59+
)
60+
parser.add_argument("--l0iter", help="L0 iterations", type=int, default=5)
61+
parser.add_argument("--gamma_tnv", help="TNV regularization parameter", type=float)
62+
63+
# Projector choices
64+
rtk.add_rtkprojectors_group(parser)
65+
rtk.add_rtk3Doutputimage_group(parser)
66+
return parser
67+
68+
69+
def process(args_info: argparse.Namespace):
70+
PixelValueType = itk.F
71+
Dimension = 3
72+
73+
DecomposedProjectionType = itk.VectorImage[PixelValueType, Dimension]
74+
MaterialsVolumeType = itk.VectorImage[PixelValueType, Dimension]
75+
VolumeSeriesType = itk.Image[PixelValueType, Dimension + 1]
76+
ProjectionStackType = itk.Image[PixelValueType, Dimension]
77+
78+
# Projections reader
79+
if args_info.verbose:
80+
print(f"Reading decomposed projections from {args_info.projection}...")
81+
proj_reader = itk.ImageFileReader[DecomposedProjectionType].New()
82+
proj_reader.SetFileName(args_info.projection)
83+
proj_reader.Update()
84+
decomposedProjection = proj_reader.GetOutput()
85+
86+
NumberOfMaterials = decomposedProjection.GetVectorLength()
87+
88+
# Geometry
89+
if args_info.verbose:
90+
print(f"Reading geometry from {args_info.geometry}...")
91+
geometry = rtk.read_geometry(args_info.geometry)
92+
93+
# Create 4D input. Fill it either with an existing materials volume
94+
# read from a file or a blank image
95+
vecVol2VolSeries = rtk.VectorImageToImageFilter[
96+
MaterialsVolumeType, VolumeSeriesType
97+
].New()
98+
99+
if args_info.input is not None:
100+
if args_info.like is not None:
101+
print("WARNING: --like ignored since --input was given")
102+
if args_info.verbose:
103+
print(f"Reading input volume {args_info.input}...")
104+
reference = itk.ImageFileReader[MaterialsVolumeType].New()
105+
reference.SetFileName(args_info.input)
106+
reference.Update()
107+
vecVol2VolSeries.SetInput(reference.GetOutput())
108+
vecVol2VolSeries.Update()
109+
input = vecVol2VolSeries.GetOutput()
110+
elif args_info.like is not None:
111+
if args_info.verbose:
112+
print(f"Reading reference volume {args_info.like} to infer geometry...")
113+
reference = itk.ImageFileReader[MaterialsVolumeType].New()
114+
reference.SetFileName(args_info.like)
115+
reference.Update()
116+
vecVol2VolSeries.SetInput(reference.GetOutput())
117+
vecVol2VolSeries.UpdateOutputInformation()
118+
constantImageSource = rtk.ConstantImageSource[VolumeSeriesType].New()
119+
constantImageSource.SetInformationFromImage(vecVol2VolSeries.GetOutput())
120+
constantImageSource.Update()
121+
input = constantImageSource.GetOutput()
122+
else:
123+
# Create new empty volume
124+
constantImageSource = rtk.ConstantImageSource[VolumeSeriesType].New()
125+
126+
imageSize = itk.Size[4]()
127+
imageSize.Fill(args_info.size[0])
128+
for i in range(min(len(args_info.size), Dimension)):
129+
imageSize[i] = args_info.size[i]
130+
imageSize[Dimension] = NumberOfMaterials
131+
132+
imageSpacing = itk.Vector[itk.D, 4]()
133+
imageSpacing.Fill(args_info.spacing[0])
134+
for i in range(min(len(args_info.spacing), Dimension)):
135+
imageSpacing[i] = args_info.spacing[i]
136+
imageSpacing[Dimension] = 1.0
137+
138+
imageOrigin = itk.Point[itk.D, 4]()
139+
for i in range(Dimension):
140+
imageOrigin[i] = imageSpacing[i] * (imageSize[i] - 1) * -0.5
141+
if args_info.origin is not None:
142+
for i in range(min(len(args_info.origin), Dimension)):
143+
imageOrigin[i] = args_info.origin[i]
144+
imageOrigin[Dimension] = 0.0
145+
146+
imageDirection = itk.Matrix[itk.D, 4, 4]()
147+
imageDirection.SetIdentity()
148+
if args_info.direction is not None:
149+
for i in range(Dimension):
150+
for j in range(Dimension):
151+
imageDirection[i][j] = args_info.direction[i * Dimension + j]
152+
153+
constantImageSource.SetOrigin(imageOrigin)
154+
constantImageSource.SetSpacing(imageSpacing)
155+
constantImageSource.SetDirection(imageDirection)
156+
constantImageSource.SetSize(imageSize)
157+
constantImageSource.SetConstant(0.0)
158+
constantImageSource.Update()
159+
input = constantImageSource.GetOutput()
160+
161+
# Duplicate geometry and transform the N M-vector projections into N*M scalar projections
162+
# Each material will occupy one frame of the 4D reconstruction, therefore all projections
163+
# of one material need to have the same phase.
164+
# Note : the 4D CG filter is optimized when projections with identical phases are packed together
165+
166+
# Geometry
167+
initialNumberOfProjections = (
168+
decomposedProjection.GetLargestPossibleRegion().GetSize()[Dimension - 1]
169+
)
170+
for material in range(1, NumberOfMaterials):
171+
for proj in range(initialNumberOfProjections):
172+
geometry.AddProjectionInRadians(
173+
geometry.GetSourceToIsocenterDistances()[proj],
174+
geometry.GetSourceToDetectorDistances()[proj],
175+
geometry.GetGantryAngles()[proj],
176+
geometry.GetProjectionOffsetsX()[proj],
177+
geometry.GetProjectionOffsetsY()[proj],
178+
geometry.GetOutOfPlaneAngles()[proj],
179+
geometry.GetInPlaneAngles()[proj],
180+
geometry.GetSourceOffsetsX()[proj],
181+
geometry.GetSourceOffsetsY()[proj],
182+
)
183+
geometry.SetCollimationOfLastProjection(
184+
geometry.GetCollimationUInf()[proj],
185+
geometry.GetCollimationUSup()[proj],
186+
geometry.GetCollimationVInf()[proj],
187+
geometry.GetCollimationVSup()[proj],
188+
)
189+
190+
# Signal
191+
fakeSignal = []
192+
for material in range(NumberOfMaterials):
193+
v = round(material / NumberOfMaterials * 1000.0) / 1000.0
194+
for proj in range(initialNumberOfProjections):
195+
fakeSignal.append(v)
196+
197+
# Projections
198+
vproj2proj = rtk.VectorImageToImageFilter[
199+
DecomposedProjectionType, ProjectionStackType
200+
].New()
201+
vproj2proj.SetInput(decomposedProjection)
202+
vproj2proj.Update()
203+
204+
# Release the memory holding the stack of original projections
205+
decomposedProjection.ReleaseData()
206+
207+
# Compute the interpolation weights
208+
signalToInterpolationWeights = rtk.SignalToInterpolationWeights.New()
209+
signalToInterpolationWeights.SetSignal(fakeSignal)
210+
signalToInterpolationWeights.SetNumberOfReconstructedFrames(NumberOfMaterials)
211+
signalToInterpolationWeights.Update()
212+
213+
# Set the forward and back projection filters to be used
214+
# Instantiate ROOSTER with CUDA image types if available, otherwise CPU types
215+
if hasattr(itk, "CudaImage"):
216+
cudaVolumeSeriesType = itk.CudaImage[PixelValueType, Dimension + 1]
217+
cudaProjectionStackType = itk.CudaImage[PixelValueType, Dimension]
218+
rooster = rtk.FourDROOSTERConeBeamReconstructionFilter[
219+
cudaVolumeSeriesType, cudaProjectionStackType
220+
].New()
221+
rooster.SetInputVolumeSeries(itk.cuda_image_from_image(input))
222+
rooster.SetInputProjectionStack(
223+
itk.cuda_image_from_image(vproj2proj.GetOutput())
224+
)
225+
else:
226+
rooster = rtk.FourDROOSTERConeBeamReconstructionFilter[
227+
VolumeSeriesType, ProjectionStackType
228+
].New()
229+
rooster.SetInputVolumeSeries(input)
230+
rooster.SetInputProjectionStack(vproj2proj.GetOutput())
231+
232+
# Configure projectors from args
233+
rtk.SetForwardProjectionFromArgParse(args_info, rooster)
234+
rtk.SetBackProjectionFromArgParse(args_info, rooster)
235+
236+
rooster.SetCG_iterations(args_info.cgiter)
237+
rooster.SetMainLoop_iterations(args_info.niter)
238+
rooster.SetCudaConjugateGradient(args_info.cudacg)
239+
rooster.SetDisableDisplacedDetectorFilter(args_info.nodisplaced)
240+
241+
rooster.SetGeometry(geometry)
242+
rooster.SetWeights(signalToInterpolationWeights.GetOutput())
243+
rooster.SetSignal(fakeSignal)
244+
245+
# For each optional regularization step, set whether or not
246+
# it should be performed, and provide the necessary inputs
247+
248+
# Positivity
249+
rooster.SetPerformPositivity(not args_info.nopositivity)
250+
251+
# No motion mask is used, since there is no motion
252+
rooster.SetPerformMotionMask(False)
253+
254+
# Spatial TV
255+
if args_info.gamma_space is not None:
256+
rooster.SetGammaTVSpace(args_info.gamma_space)
257+
rooster.SetTV_iterations(args_info.tviter)
258+
rooster.SetPerformTVSpatialDenoising(True)
259+
else:
260+
rooster.SetPerformTVSpatialDenoising(False)
261+
262+
# Spatial wavelets
263+
if args_info.threshold is not None:
264+
rooster.SetSoftThresholdWavelets(args_info.threshold)
265+
rooster.SetOrder(args_info.order)
266+
rooster.SetNumberOfLevels(args_info.levels)
267+
rooster.SetPerformWaveletsSpatialDenoising(True)
268+
else:
269+
rooster.SetPerformWaveletsSpatialDenoising(False)
270+
271+
# Temporal TV
272+
if args_info.gamma_time is not None:
273+
rooster.SetGammaTVTime(args_info.gamma_time)
274+
rooster.SetTV_iterations(args_info.tviter)
275+
rooster.SetPerformTVTemporalDenoising(True)
276+
else:
277+
rooster.SetPerformTVTemporalDenoising(False)
278+
279+
# Temporal L0
280+
if args_info.lambda_time is not None:
281+
rooster.SetLambdaL0Time(args_info.lambda_time)
282+
rooster.SetL0_iterations(args_info.l0iter)
283+
rooster.SetPerformL0TemporalDenoising(True)
284+
else:
285+
rooster.SetPerformL0TemporalDenoising(False)
286+
287+
# Total nuclear variation
288+
if args_info.gamma_tnv is not None:
289+
rooster.SetGammaTNV(args_info.gamma_tnv)
290+
rooster.SetTV_iterations(args_info.tviter)
291+
rooster.SetPerformTNVDenoising(True)
292+
else:
293+
rooster.SetPerformTNVDenoising(False)
294+
295+
if args_info.verbose:
296+
print("Running ROOSTER reconstruction...")
297+
rooster.Update()
298+
299+
# Convert to result to a vector image
300+
volSeries2VecVol = rtk.ImageToVectorImageFilter[
301+
VolumeSeriesType, MaterialsVolumeType
302+
].New()
303+
volSeries2VecVol.SetInput(rooster.GetOutput())
304+
volSeries2VecVol.Update()
305+
306+
# Write
307+
if args_info.verbose:
308+
print(f"Writing output to {args_info.output}...")
309+
writer = itk.ImageFileWriter[MaterialsVolumeType].New()
310+
writer.SetFileName(args_info.output)
311+
writer.SetInput(volSeries2VecVol.GetOutput())
312+
writer.Update()
313+
314+
315+
def main(argv=None):
316+
parser = build_parser()
317+
args_info = parser.parse_args(argv)
318+
process(args_info)
319+
320+
321+
if __name__ == "__main__":
322+
main()

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@ rtkshowgeometry = "itk.rtkshowgeometry:main"
8282
rtksart = "itk.rtksart:main"
8383
rtksimulatedgeometry = "itk.rtksimulatedgeometry:main"
8484
rtkspectraldenoiseprojections = "itk.rtkspectraldenoiseprojections:main"
85+
rtkspectralrooster = "itk.rtkspectralrooster:main"
8586
rtksubselect = "itk.rtksubselect:main"
8687
rtktotalvariationdenoising = "itk.rtktotalvariationdenoising:main"
8788
rtktotalnuclearvariationdenoising = "itk.rtktotalnuclearvariationdenoising:main"

‎wrapping/__init_rtk__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@
6161
"rtksart",
6262
"rtksimulatedgeometry",
6363
"rtkspectraldenoiseprojections",
64+
"rtkspectralrooster",
6465
"rtksubselect",
6566
"rtktotalnuclearvariationdenoising",
6667
"rtktotalvariationdenoising",

‎wrapping/itkImageToImageFilterRTK.wrap‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,7 @@ itk_wrap_class("itk::ImageToImageFilter" POINTER)
136136
itk_wrap_template("I${ITKM_${t}}2I${ITKM_${t}}1" "itk::Image<${ITKT_${t}}, 2>, itk::Image<${ITKT_${t}}, 1>")
137137
endif()
138138
itk_wrap_template("I${ITKM_${t}}3VI${ITKM_${t}}2" "itk::Image<${ITKT_${t}}, 3>, itk::VectorImage<${ITKT_${t}}, 2>")
139+
itk_wrap_template("I${ITKM_${t}}4VI${ITKM_${t}}3" "itk::Image<${ITKT_${t}}, 4>, itk::VectorImage<${ITKT_${t}}, 3>")
139140
itk_wrap_template("VI${ITKM_${t}}2I${ITKM_${t}}3" "itk::VectorImage<${ITKT_${t}}, 2>, itk::Image<${ITKT_${t}}, 3>")
140141
itk_wrap_template("VI${ITKM_${t}}3I${ITKM_${t}}4" "itk::VectorImage<${ITKT_${t}}, 3>, itk::Image<${ITKT_${t}}, 4>")
141142
endforeach()

‎wrapping/rtkImageToVectorImageFilter.wrap‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,5 +2,6 @@ itk_wrap_class("rtk::ImageToVectorImageFilter" POINTER)
22
foreach(t ${WRAP_ITK_REAL})
33
itk_wrap_template("I${ITKM_${t}}2VI${ITKM_${t}}2" "itk::Image<${ITKT_${t}}, 2>, itk::VectorImage<${ITKT_${t}}, 2>")
44
itk_wrap_template("I${ITKM_${t}}3VI${ITKM_${t}}2" "itk::Image<${ITKT_${t}}, 3>, itk::VectorImage<${ITKT_${t}}, 2>")
5+
itk_wrap_template("I${ITKM_${t}}4VI${ITKM_${t}}3" "itk::Image<${ITKT_${t}}, 4>, itk::VectorImage<${ITKT_${t}}, 3>")
56
endforeach()
67
itk_end_wrap_class()

0 commit comments

Comments
 (0)