Skip to content

Commit

Permalink
ENH: Add Python command line tool for conjugate gradient
Browse files Browse the repository at this point in the history
  • Loading branch information
Simon Rit authored and SimonRit committed Jul 15, 2023
1 parent 2f0e76a commit 07bcb36
Show file tree
Hide file tree
Showing 8 changed files with 495 additions and 3 deletions.
9 changes: 8 additions & 1 deletion CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -255,12 +255,19 @@ if(RTK_BUILD_APPLICATIONS)
endif()
if(ITK_WRAP_PYTHON)
# Install Python application
install(FILES applications/rtkelektasynergygeometry/rtkelektasynergygeometry.py
install(FILES applications/rtkconjugategradient/rtkconjugategradient.py
applications/rtkelektasynergygeometry/rtkelektasynergygeometry.py
applications/rtkorageometry/rtkorageometry.py
applications/rtksimulatedgeometry/rtksimulatedgeometry.py
applications/rtkvarianobigeometry/rtkvarianobigeometry.py
DESTINATION ${RTK_INSTALL_LIB_DIR}
COMPONENT PythonWheelRuntimeLibraries)
WRAP_ITK_PYTHON_BINDINGS_INSTALL(/itk "RTK"
applications/rtk3Doutputimage_group.py
applications/rtkinputprojections_group.py
applications/rtkiterations_group.py
applications/rtkprojectors_group.py
wrapping/__init_rtk__.py)
endif()

# --------------------------------------------------------
Expand Down
61 changes: 61 additions & 0 deletions applications/rtk3Doutputimage_group.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
import itk

__all__ = [
'add_rtk3Doutputimage_group',
'SetConstantImageSourceFromArgParse',
]

# Mimicks rtk3Doutputimage_section.ggo
def add_rtk3Doutputimage_group(parser):
rtk3Doutputimage_group = parser.add_argument_group('Output 3D image properties')
rtk3Doutputimage_group.add_argument('--origin', help='Origin (default=centered)', type=float, nargs='+')
rtk3Doutputimage_group.add_argument('--dimension', help='Dimension', type=int, nargs='+', default=[256])
rtk3Doutputimage_group.add_argument('--spacing', help='Spacing', type=float, nargs='+', default=[1])
rtk3Doutputimage_group.add_argument('--direction', help='Direction', type=float, nargs='+')
rtk3Doutputimage_group.add_argument('--like', help='Copy information from this image (origin, dimension, spacing, direction)')

# Mimicks SetConstantImageSourceFromGgo
def SetConstantImageSourceFromArgParse(source, args_info):
ImageType = type(source.GetOutput())

Dimension = ImageType.GetImageDimension()

imageDimension = itk.Size[Dimension]()
imageDimension.Fill(args_info.dimension[0])
for i in range(min(len(args_info.dimension), Dimension)):
imageDimension[i] = args_info.dimension[i]

imageSpacing = itk.Vector[itk.D, Dimension]()
imageSpacing.Fill(args_info.spacing[0])
for i in range(min(len(args_info.spacing), Dimension)):
imageSpacing[i] = args_info.spacing[i]

imageOrigin = itk.Point[itk.D, Dimension]()
for i in range(Dimension):
imageOrigin[i] = imageSpacing[i] * (imageDimension[i] - 1) * -0.5
if args_info.origin is not None:
for i in range(min(len(args_info.origin), Dimension)):
imageOrigin[i] = args_info.origin[i]

imageDirection = source.GetOutput().GetDirection()
if args_info.direction is not None:
for i in range(Dimension):
for j in range(Dimension):
imageDirection[i][j] = args_info.direction[i * Dimension + j]

source.SetOrigin(imageOrigin)
source.SetSpacing(imageSpacing)
source.SetDirection(imageDirection)
source.SetSize(imageDimension)
source.SetConstant(0.)

# Copy output image information from an existing file, if requested
# Overwrites parameters given in command line, if any
if args_info.like is not None:
LikeReaderType = itk.ImageFileReader[ImageType]
likeReader = LikeReaderType.New()
likeReader.SetFileName(args_info.like)
likeReader.UpdateOutputInformation()
source.SetInformationFromImage(likeReader.GetOutput())

source.UpdateOutputInformation()
133 changes: 133 additions & 0 deletions applications/rtkconjugategradient/rtkconjugategradient.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,133 @@
#!/usr/bin/env python
import argparse
import sys
import itk
from itk import RTK as rtk

def GetCudaImageFromImage(img):
if hasattr(itk, 'CudaImage'):
# TODO: avoid this Update by implementing an itk::ImageToCudaImageFilter
img.Update()
cuda_img = itk.CudaImage[itk.itkCType.GetCTypeForDType(img.dtype),img.ndim].New()
cuda_img.SetPixelContainer(img.GetPixelContainer())
cuda_img.CopyInformation(img)
cuda_img.SetBufferedRegion(img.GetBufferedRegion())
cuda_img.SetRequestedRegion(img.GetRequestedRegion())
return cuda_img
else:
return img

if __name__ == '__main__':
# argument parsing
parser = argparse.ArgumentParser(description=
'Reconstructs a 3D volume from a sequence of projections with a conjugate gradient technique',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)

parser.add_argument('--verbose', '-v', help='Verbose execution', action='store_true')
parser.add_argument('--geometry', '-g', help='Geometry file name', required=True)
parser.add_argument('--output', '-o', help='Output file name', required=True)
parser.add_argument('--niterations', '-n', type=int, default=5, help='Number of iterations')
parser.add_argument('--input', '-i', help='Input volume')
parser.add_argument('--weights', '-w', help='Weights file for Weighted Least Squares (WLS)')
parser.add_argument('--gamma', help='Laplacian regularization weight')
parser.add_argument('--tikhonov', help='Tikhonov regularization weight')
parser.add_argument('--nocudacg', help='Do not perform conjugate gradient calculations on GPU', action='store_true')
parser.add_argument('--mask', '-m', help='Apply a support binary mask: reconstruction kept null outside the mask)'
)
parser.add_argument('--nodisplaced', help='Disable the displaced detector filter', action='store_true')

rtk.add_rtkinputprojections_group(parser)
rtk.add_rtk3Doutputimage_group(parser)
rtk.add_rtkprojectors_group(parser)
rtk.add_rtkiterations_group(parser)

args_info = parser.parse_args()

OutputPixelType = itk.F
Dimension = 3

OutputImageType = itk.Image[OutputPixelType, Dimension]

# Projections reader
reader = rtk.ProjectionsReader[OutputImageType].New()
rtk.SetProjectionsReaderFromArgParse(reader, args_info)

# Geometry
if args_info.verbose:
print(f'Reading geometry information from {args_info.geometry}')
geometryReader = rtk.ThreeDCircularProjectionGeometryXMLFileReader.New()
geometryReader.SetFilename(args_info.geometry)
geometryReader.GenerateOutputInformation()
geometry = geometryReader.GetOutputObject()

# Create input: either an existing volume read from a file or a blank image
if args_info.input is not None:
# Read an existing image to initialize the volume
InputReaderType = itk.ImageFileReader[OutputImageType]
inputReader = InputReaderType.New()
inputReader.SetFileName(args_info.input)
inputFilter = inputReader
else:
# Create new empty volume
ConstantImageSourceType = rtk.ConstantImageSource[OutputImageType]
constantImageSource = ConstantImageSourceType.New()
rtk.SetConstantImageSourceFromArgParse(constantImageSource, args_info)
inputFilter = constantImageSource

# Read weights if given, otherwise default to weights all equal to one
if args_info.weights is not None:
WeightsReaderType = itk.ImageFileReader[OutputImageType]
weightsReader = WeightsReaderType.New()
weightsReader.SetFileName(args_info.weights)
weightsSource = weightsReader
else:
ConstantWeightsSourceType = rtk.ConstantImageSource[OutputImageType]
constantWeightsSource = ConstantWeightsSourceType.New()

# Set the weights to be like the projections
reader.UpdateOutputInformation()

constantWeightsSource.SetInformationFromImage(reader.GetOutput())
constantWeightsSource.SetConstant(1.)
weightsSource = constantWeightsSource

# Read Support Mask if given
if args_info.mask is not None:
supportmask = itk.imread(args_info.mask)

# Set the forward and back projection filters to be used
if hasattr(itk, 'CudaImage'):
OutputCudaImageType = itk.CudaImage[OutputPixelType, Dimension]
ConjugateGradientFilterType = rtk.ConjugateGradientConeBeamReconstructionFilter[OutputCudaImageType]
else:
ConjugateGradientFilterType = rtk.ConjugateGradientConeBeamReconstructionFilter[OutputImageType]

conjugategradient = ConjugateGradientFilterType.New()
rtk.SetForwardProjectionFromArgParse(args_info, conjugategradient)
rtk.SetBackProjectionFromArgParse(args_info, conjugategradient)
rtk.SetIterationsReportFromArgParse(args_info, conjugategradient)
conjugategradient.SetInput(GetCudaImageFromImage(inputFilter.GetOutput()))
conjugategradient.SetInput(1, GetCudaImageFromImage(reader.GetOutput()))
conjugategradient.SetInput(2, GetCudaImageFromImage(weightsSource.GetOutput()))
conjugategradient.SetCudaConjugateGradient(not args_info.nocudacg)
if args_info.mask is not None:
conjugategradient.SetSupportMask(GetCudaImageFromImage(supportmask))

if args_info.gamma is not None:
conjugategradient.SetGamma(args_info.gamma)
if args_info.tikhonov is not None:
conjugategradient.SetTikhonov(args_info.tikhonov)

conjugategradient.SetGeometry(geometry)
conjugategradient.SetNumberOfIterations(args_info.niterations)
conjugategradient.SetDisableDisplacedDetectorFilter(args_info.nodisplaced)

# TODO REPORT_ITERATIONS(conjugategradient, ConjugateGradientFilterType, OutputImageType)

conjugategradient.Update()

# Write
writer = itk.ImageFileWriter[OutputImageType].New()
writer.SetFileName(args_info.output)
writer.SetInput(conjugategradient.GetOutput())
writer.Update()
144 changes: 144 additions & 0 deletions applications/rtkinputprojections_group.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,144 @@
import itk

__all__ = [
'add_rtkinputprojections_group',
'GetProjectionsFileNamesFromArgParse',
]

# Mimicks rtkinputprojections_section.ggo
def add_rtkinputprojections_group(parser):
rtkinputprojections_group = parser.add_argument_group('Input projections and their pre-processing')
rtkinputprojections_group.add_argument('--path', '-p', help='Path containing projections', required=True)
rtkinputprojections_group.add_argument('--regexp', '-r', help='Regular expression to select projection files in path', required=True)
rtkinputprojections_group.add_argument('--nsort', help='Numeric sort for regular expression matches', action='store_true')
rtkinputprojections_group.add_argument('--submatch', help='Index of the submatch that will be used to sort matches', type=int, default=0)
rtkinputprojections_group.add_argument('--nolineint', help='Disable raw to line integral conversion, just casts to float', type=bool, default=False)
rtkinputprojections_group.add_argument('--newdirection', help='New value of input projections (before pre-processing)', type=float, nargs='+')
rtkinputprojections_group.add_argument('--neworigin', help='New origin of input projections (before pre-processing)', type=float, nargs='+')
rtkinputprojections_group.add_argument('--newspacing', help='New spacing of input projections (before pre-processing)', type=float, nargs='+')
rtkinputprojections_group.add_argument('--lowercrop', help='Lower boundary crop size', type=int, nargs='+', default=[0])
rtkinputprojections_group.add_argument('--uppercrop', help='Upper boundary crop size', type=int, nargs='+', default=[0])
rtkinputprojections_group.add_argument('--binning', help='Shrink / Binning factos in each direction', type=int, nargs='+', default=[1])
rtkinputprojections_group.add_argument('--wpc', help='Water precorrection coefficients (default is no correction)', type=float, nargs='+')
rtkinputprojections_group.add_argument('--spr', help='Boellaard scatter correction: scatter-to-primary ratio', type=float, default=0)
rtkinputprojections_group.add_argument('--nonneg', help='Boellaard scatter correction: non-negativity threshold', type=float)
rtkinputprojections_group.add_argument('--airthres', help='Boellaard scatter correction: air threshold', type=float)
rtkinputprojections_group.add_argument('--i0', help='I0 value (when assumed constant per projection), 0 means auto', type=float)
rtkinputprojections_group.add_argument('--idark', help='IDark value, i.e., value when beam is off', type=float, default=0)
rtkinputprojections_group.add_argument('--component', help='Vector component to extract, for multi-material projections', type=int, default=0)
rtkinputprojections_group.add_argument('--radius', help='Radius of neighborhood for conditional median filtering', type=int, nargs='+', default=[0])
rtkinputprojections_group.add_argument('--multiplier', help='Threshold multiplier for conditional median filtering', type=float, default=0)

# Mimicks GetProjectionsFileNamesFromGgo
def GetProjectionsFileNamesFromArgParse(args_info):
# Generate file names
names = itk.RegularExpressionSeriesFileNames.New()
names.SetDirectory(args_info.path)
names.SetNumericSort(args_info.nsort)
names.SetRegularExpression(args_info.regexp)
names.SetSubMatch(args_info.submatch)

if args_info.verbose:
print(f'Regular expression matches {len(names.GetFileNames())} file(s)...')

# Check submatch in file names TODO

fileNames = names.GetFileNames()
# rtk.RegisterIOFactories() TODO
idxtopop = []
i = 0
for fn in fileNames:
imageio = itk.ImageIOFactory.CreateImageIO(fn, itk.CommonEnums.IOFileMode_ReadMode)
if imageio is None:
print(f'Ignoring file: {fn}')
idxtopop.append(i)
++i

for id in idxtopop:
fileNames.pop(id)

return fileNames

# Mimicks SetProjectionsReaderFromGgo
def SetProjectionsReaderFromArgParse(reader, args_info):
fileNames = GetProjectionsFileNamesFromArgParse(args_info)

# Vector component extraction
if args_info.component is not None:
reader.SetVectorComponent(args_info.component)

# Change image information
Dimension = reader.GetOutput().GetImageDimension()
if args_info.newdirection is not None:
direction = itk.Matrix[itk.D,Dimension,Dimension]()
direction.Fill(args_info.newdirection[0])
for i in range(len(args_info.newdirection)):
direction[i / Dimension][i % Dimension] = args_info.newdirection_arg[i]
reader.SetDirection(direction)

if args_info.newspacing is not None:
spacing = itk.Vector[itk.D,Dimension]()
spacing.Fill(args_info.newspacing[0])
for i in range(len(args_info.newspacing)):
spacing[i] = args_info.newspacing[i]
reader.SetSpacing(spacing)

if args_info.neworigin is not None:
origin = itk.Point[itk.D,Dimension]()
origin.Fill(args_info.neworigin[0])
for i in range(len(args_info.neworigin)):
origin[i] = args_info.neworigin[i]
reader.SetOrigin(origin)

# Crop boundaries
upperCrop = [0]*Dimension
lowerCrop = [0]*Dimension
if args_info.lowercrop is not None:
for i in range(len(args_info.lowercrop)):
lowerCrop[i] = args_info.lowercrop[i]
reader.SetLowerBoundaryCropSize(lowerCrop)
if args_info.uppercrop is not None:
for i in range(len(args_info.uppercrop)):
upperCrop[i] = args_info.uppercrop[i]
reader.SetUpperBoundaryCropSize(upperCrop)

# Conditional median
medianRadius = reader.GetMedianRadius()
if args_info.radius is not None:
for i in range(len(args_info.radius)):
medianRadius[i] = args_info.radius[i]
reader.SetMedianRadius(medianRadius)
if args_info.multiplier is not None:
reader.SetConditionalMedianThresholdMultiplier(args_info.multiplier)

# Shrink / Binning
binFactors = reader.GetShrinkFactors()
if args_info.binning is not None:
for i in range(len(args_info.binning)):
binFactors[i] = args_info.binning[i]
reader.SetShrinkFactors(binFactors)

# Boellaard scatter correction
if args_info.spr is not None:
reader.SetScatterToPrimaryRatio(args_info.spr)
if args_info.nonneg is not None:
reader.SetNonNegativityConstraintThreshold(args_info.nonneg)
if args_info.airthres is not None:
reader.SetAirThreshold(args_info.airthres)

# I0 and IDark
if args_info.i0 is not None:
reader.SetI0(args_info.i0)
reader.SetIDark(args_info.idark)

# Line integral flag
if args_info.nolineint:
reader.ComputeLineIntegralOff()

# Water precorrection
if args_info.wpc is not None:
reader.SetWaterPrecorrectionCoefficients(coeffs)

# Pass list to projections reader
reader.SetFileNames(fileNames)
reader.UpdateOutputInformation()
Loading

0 comments on commit 07bcb36

Please sign in to comment.