BHS/BHS_reg.py at master · ManskeLab/BHS · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
#-----------------------------------------------------
# BHS_reg.py
#
# Created by:   Michael Kuczynski
# Created on:   05-02-2020
#
# Description: Perform 3 month XCT to 12 month XCT image registration for MCP joints.
#              First, an initial alignment of images is obtained by matching geometric
#              image centres. Final image alignment is obtained by optimizing the mean
#              squares error.
#              To obtain more accurate registration of the metacarpal bone, an image mask
#              of the metacarpal is provided as input.
#
#-----------------------------------------------------
# Usage: python BHS_reg.py arg1 arg2 arg3 arg4 arg5 arg6
#
# Where:  arg1 = The input 03MO XCT grayscale image
#         arg2 = The input 12MO XCT grayscale image
#         arg3 = The registered output 12MO to 03MO XCT grayscale image
#         arg4 = The input 03MO XCT segmented image
#         arg5 = The input 12MO XCT segmented image
#         arg6 = The registered output 12MO to 03MO XCT segmented image
#
# Notes: this script should work with NIfTI images as well, but hasn't been tested.
#-----------------------------------------------------
import os
import argparse

import SimpleITK as sitk

#----------------------------------------------#
# STEP 1: Read in images
#----------------------------------------------#
# Parse input arguments
parser = argparse.ArgumentParser()
parser.add_argument( 'input_03MO_gray', type=str, help='The input 03MO XCT grayscale image (path + filename)' )
parser.add_argument( 'input_12MO_gray', type=str, help='The input 12MO XCT grayscale image (path + filename)' )
parser.add_argument( 'output_12MO_to_03MO_reg', type=str, help='The registered output 12MO to 03MO XCT grayscale image (path + filename)' )
parser.add_argument( 'input_03MO_seg', type=str, help='The input 03MO XCT segmented image (path + filename)' )
parser.add_argument( 'input_12MO_seg', type=str, help='The input 12MO XCT segmented image (path + filename)' )
parser.add_argument( 'output_12MO_to_03MO_seg_reg', type=str, help='The registered output 12MO to 03MO XCT segmented image (path + filename)' )
args = parser.parse_args()

# Grayscale images
XCT_3MO_image_path = args.input_03MO_gray
XCT_12MO_image_path = args.input_12MO_gray
XCT_12MO_TO_03MO_REG_path = args.output_12MO_to_03MO_reg

# Segmented images
XCT_3MO_SEG_image_path = args.input_03MO_seg
XCT_12MO_SEG_image_path = args.input_12MO_seg
XCT_12MO_TO_03MO_SEG_REG_path = args.output_12MO_to_03MO_seg_reg

# Get the sample number from the input images
sampleNum = ( XCT_3MO_image_path.rsplit('BHS_', 1)[1] ).rsplit('_', 1)[0]

# Get the MCP joint number from the input images
mcp = ( XCT_3MO_image_path.rsplit('3MO_', 1)[1] ).rsplit('.', 1)[0]

# Get the output directory
outDirectory, outFilename = os.path.split(XCT_12MO_TO_03MO_REG_path)

# Registered 12MO image and transformation matrix (12MO -> 03MO XCT image space):
XCT_12MO_to_3MO_TMAT_path = os.path.join(outDirectory, 'BHS_' + sampleNum + '_' + mcp + '_REG.tfm')

# Read images
# 03MO
print('Reading in {}'.format(XCT_3MO_image_path))
XCT_3MO_image = sitk.ReadImage(XCT_3MO_image_path)

print('Reading in {}'.format(XCT_3MO_SEG_image_path))
XCT_3MO_SEG_image = sitk.ReadImage(XCT_3MO_SEG_image_path)

# 12MO
print('Reading in {}'.format(XCT_12MO_image_path))
XCT_12MO_image = sitk.ReadImage(XCT_12MO_image_path)

print('Reading in {}'.format(XCT_12MO_SEG_image_path))
XCT_12MO_SEG_image = sitk.ReadImage(XCT_12MO_SEG_image_path)

#----------------------------------------------#
# STEP 2: Perform landmark transformation
#----------------------------------------------#
# Set initial transform by matching geometric centres
initalTransform_12MO_to_03MO = sitk.CenteredTransformInitializer(XCT_3MO_image, XCT_12MO_image, sitk.Euler3DTransform(), sitk.CenteredTransformInitializerFilter.GEOMETRY)

#----------------------------------------------#
# STEP 3: Setup registration method
#----------------------------------------------#
# Connect all of the observers so that we can perform plotting during registration.
def command_iteration(method) :
  print( '{0:3} = {1:10.5f} : {2}'.format( method.GetOptimizerIteration(), method.GetMetricValue(), method.GetOptimizerPosition() ) )

# Set up registration (12MO -> 03MO)
reg = sitk.ImageRegistrationMethod()

# Similarity metric settings:
reg.SetMetricAsMeanSquares()
reg.SetMetricSamplingStrategy(reg.RANDOM)
reg.SetMetricSamplingPercentage(0.01)   # Make this value smaller for faster (less accurate) results

#Set Interpolator
reg.SetInterpolator(sitk.sitkLinear)

# Optimizer settings.
reg.SetOptimizerAsPowell(numberOfIterations=500)
reg.SetOptimizerScalesFromPhysicalShift()

# Setup for the multi-resolution framework.
reg.SetShrinkFactorsPerLevel(shrinkFactors = [1, 1])
reg.SetSmoothingSigmasPerLevel(smoothingSigmas=[1.0, 0])
reg.SmoothingSigmasAreSpecifiedInPhysicalUnitsOn()

reg.AddCommand( sitk.sitkIterationEvent, lambda: command_iteration(reg) )

#----------------------------------------------#
# STEP 4: Perform the registration
#----------------------------------------------#
# Register 12MO and 03MO grayscale images
# Don't optimize in-place, we would possibly like to run this cell multiple times.
reg.SetInitialTransform(initalTransform_12MO_to_03MO, inPlace=False)
print('Start registration')
finalTransform = reg.Execute( sitk.Cast(XCT_3MO_image, sitk.sitkFloat64), sitk.Cast(XCT_12MO_image, sitk.sitkFloat64) )
print('Writing to {}'.format(XCT_12MO_to_3MO_TMAT_path))
sitk.WriteTransform(finalTransform, XCT_12MO_to_3MO_TMAT_path)

#----------------------------------------------#
# STEP 5: Resample and write registered image
#----------------------------------------------#
# Resample 12MO grayscale image
print('Resampling')
XCT_12MO_resampled = sitk.Resample( XCT_12MO_image, XCT_3MO_image, finalTransform, sitk.sitkLinear, 0.0, XCT_12MO_image.GetPixelID() )

print('Writing to {}'.format(XCT_12MO_TO_03MO_REG_path))
sitk.WriteImage(XCT_12MO_resampled, XCT_12MO_TO_03MO_REG_path)

#----------------------------------------------#
# STEP 6: Transform the segmented 12MO image
#----------------------------------------------#
# Resample 12MO segmented image
print('Resampling')
XCT_12MO_SEG_resampled = sitk.Resample( XCT_12MO_SEG_image, XCT_3MO_SEG_image, finalTransform, sitk.sitkLinear, 0.0, XCT_12MO_SEG_image.GetPixelID() )

print('Writing to {}'.format(XCT_12MO_TO_03MO_SEG_REG_path))
sitk.WriteImage(XCT_12MO_SEG_resampled, XCT_12MO_TO_03MO_SEG_REG_path)