SimpleITK  1.3.0.dev466
Python/FFTConvolution.py
1 #!/usr/bin/env python
2 #=========================================================================
3 #
4 # Copyright Insight Software Consortium
5 #
6 # Licensed under the Apache License, Version 2.0 (the "License");
7 # you may not use this file except in compliance with the License.
8 # You may obtain a copy of the License at
9 #
10 # http://www.apache.org/licenses/LICENSE-2.0.txt
11 #
12 # Unless required by applicable law or agreed to in writing, software
13 # distributed under the License is distributed on an "AS IS" BASIS,
14 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15 # See the License for the specific language governing permissions and
16 # limitations under the License.
17 #
18 #=========================================================================
19 
20 from __future__ import print_function
21 from __future__ import division
22 
23 import SimpleITK as sitk
24 import sys
25 import os
26 import math
27 
28 if len ( sys.argv ) < 4:
29  print( "Usage: FFTConvolution <input> <kernel> <output>" )
30  sys.exit ( 1 )
31 
32 
33 inputFileName = sys.argv[1]
34 kernelFileName = sys.argv[2]
35 outputFileName = sys.argv[3]
36 
37 ### Input Image ###
38 # read the input image
39 img = sitk.ReadImage( inputFileName )
40 
41 # save the type of pixel the input is, so that we can cast the result
42 # back out to the same type
43 pixelID = img.GetPixelID()
44 
45 # pad the image
46 img = sitk.MirrorPad( img, [128] *2, [128]*2 )
47 size = img.GetSize();
48 
49 # perform the FFT
50 fftimg = sitk.ForwardFFT( sitk.Cast( img, sitk.sitkFloat32 ) )
51 
52 
53 ### Kernel Image ###
54 # Read the kernel image file
55 kernel = sitk.ReadImage( kernelFileName )
56 
57 # flip kernel about all axis
58 kernel = sitk.Flip( kernel, [True]*2 )
59 
60 
61 # normalize the kernel to sum to ~1
63 stats.Execute( kernel )
64 kernel = sitk.Cast( kernel / stats.GetSum(), sitk.sitkFloat32 )
65 
66 upadding = [0]*2
67 upadding[0] = int( math.floor( (size[0] - kernel.GetSize()[0])/2.0 ) )
68 upadding[1] = int( math.floor( (size[1] - kernel.GetSize()[1])/2.0 ) )
69 
70 lpadding = [0]*2
71 lpadding[0] = int( math.ceil( (size[0] - kernel.GetSize()[0])/2.0 ) )
72 lpadding[1] = int( math.ceil( (size[1] - kernel.GetSize()[1])/2.0 ) )
73 
74 # pad the kernel to prevent edge artifacts
75 kernel = sitk.ConstantPad( kernel, upadding, lpadding, 0.0 )
76 
77 # perform FFT on kernel
78 fftkernel = sitk.ForwardFFT( sitk.FFTShift( kernel ) )
79 
80 # meta-data must match for multiplication
81 fftkernel.SetSpacing( fftimg.GetSpacing() )
82 fftkernel.SetOrigin( fftimg.GetOrigin() )
83 fftkernel.SetDirection( fftimg.GetDirection() )
84 
85 
86 ### Convolution ###
87 # Finally perform the convolution in Fourier space by multiplication
88 img = sitk.InverseFFT( fftimg*fftkernel )
89 
90 # remove the padding
91 img = sitk.Crop( img, [128]*2, [128]*2 )
92 
93 ### Writing ###
94 # write the output image the same type as the input
95 sitk.WriteImage( sitk.Cast( img, pixelID ), outputFileName )
96 
97 
98 
99 if ( not "SITK_NOSHOW" in os.environ ):
100  sitk.Show( sitk.ReadImage( inputFileName ), "original" )
101  sitk.Show( sitk.ReadImage( kernelFileName ), "kernel" )
102  sitk.Show( sitk.Cast( img, pixelID ), "FFT_Convolution" )