def predict_each_datapart(args, net, network_config, input_batch, datapart_idx, batch_size, patch_size, predict_transform_space):
moving_image = torch.load(args.moving_image_dataset[datapart_idx])
target_image = torch.load(args.target_image_dataset[datapart_idx])
optimization_momentum = torch.load(args.deformation_parameter[datapart_idx])
for slice_idx in range(0, moving_image.size()[0]):
print(slice_idx)
moving_slice = moving_image[slice_idx].numpy()
target_slice = target_image[slice_idx].numpy()
if predict_transform_space:
moving_slice = util.convert_to_registration_space(moving_slice)
target_slice = util.convert_to_registration_space(target_slice)
predicted_momentum = util.predict_momentum(moving_slice, target_slice, input_batch, batch_size, patch_size, net, predict_transform_space);
m0_reg = common.FieldFromNPArr(predicted_momentum['image_space'], ca.MEM_DEVICE);
moving_image_ca = common.ImFromNPArr(moving_slice, ca.MEM_DEVICE)
target_image_ca = common.ImFromNPArr(target_slice, ca.MEM_DEVICE)
registration_result = registration_methods.geodesic_shooting(moving_image_ca, target_image_ca, m0_reg, args.shoot_steps, ca.MEM_DEVICE, network_config)
target_inv = common.AsNPCopy(registration_result['I1_inv'])
print(target_inv.shape)
if predict_transform_space:
target_inv = util.convert_to_predict_space(target_inv)
print(target_inv.shape)
target_inv = torch.from_numpy(target_inv)
target_image[slice_idx] = target_inv
optimization_momentum[slice_idx] = optimization_momentum[slice_idx] - torch.from_numpy(predicted_momentum['prediction_space'])
torch.save(target_image, args.warped_back_target_output[datapart_idx])
torch.save(optimization_momentum, args.momentum_residual[datapart_idx])
def predict_image(args, moving_images, target_images, output_prefixes):
if (args.use_CPU_for_shooting):
mType = ca.MEM_HOST
else:
mType = ca.MEM_DEVICE
# load the prediction network
predict_network_config = torch.load(args.prediction_parameter)
prediction_net = create_net(args, predict_network_config);
batch_size = args.batch_size
patch_size = predict_network_config['patch_size']
input_batch = torch.zeros(batch_size, 2, patch_size, patch_size, patch_size).cuda()
# use correction network if required
if args.use_correction:
correction_network_config = torch.load(args.correction_parameter);
correction_net = create_net(args, correction_network_config);
else:
correction_net = None;
# start prediction
for i in range(0, len(moving_images)):
common.Mkdir_p(os.path.dirname(output_prefixes[i]))
if (args.affine_align):
# Perform affine registration to both moving and target image to the ICBM152 atlas space.
# Registration is done using Niftireg.
call(["reg_aladin",
"-noSym", "-speeeeed", "-ref", args.atlas ,
"-flo", moving_images[i],
"-res", output_prefixes[i]+"moving_affine.nii",
"-aff", output_prefixes[i]+'moving_affine_transform.txt'])
call(["reg_aladin",
"-noSym", "-speeeeed" ,"-ref", args.atlas ,
"-flo", target_images[i],
"-res", output_prefixes[i]+"target_affine.nii",
"-aff", output_prefixes[i]+'target_affine_transform.txt'])
moving_image = common.LoadITKImage(output_prefixes[i]+"moving_affine.nii", mType)
target_image = common.LoadITKImage(output_prefixes[i]+"target_affine.nii", mType)
else:
moving_image = common.LoadITKImage(moving_images[i], mType)
target_image = common.LoadITKImage(target_images[i], mType)
#preprocessing of the image
moving_image_np = preprocess_image(moving_image, args.histeq);
target_image_np = preprocess_image(target_image, args.histeq);
grid = moving_image.grid()
#moving_image = ca.Image3D(grid, mType)
#target_image = ca.Image3D(grid, mType)
moving_image_processed = common.ImFromNPArr(moving_image_np, mType)
target_image_processed = common.ImFromNPArr(target_image_np, mType)
moving_image.setGrid(grid)
target_image.setGrid(grid)
# Indicating whether we are using the old parameter files for the Neuroimage experiments (use .t7 files from matlab .h5 format)
predict_transform_space = False
if 'matlab_t7' in predict_network_config:
predict_transform_space = True
# run actual prediction
prediction_result = util.predict_momentum(moving_image_np, target_image_np, input_batch, batch_size, patch_size, prediction_net, predict_transform_space);
m0 = prediction_result['image_space']
#convert to registration space and perform registration
m0_reg = common.FieldFromNPArr(m0, mType);
#perform correction
if (args.use_correction):
registration_result = registration_methods.geodesic_shooting(moving_image_processed, target_image_processed, m0_reg, args.shoot_steps, mType, predict_network_config)
target_inv_np = common.AsNPCopy(registration_result['I1_inv'])
correct_transform_space = False
if 'matlab_t7' in correction_network_config:
correct_transform_space = True
correction_result = util.predict_momentum(moving_image_np, target_inv_np, input_batch, batch_size, patch_size, correction_net, correct_transform_space);
m0_correct = correction_result['image_space']
m0 += m0_correct;
m0_reg = common.FieldFromNPArr(m0, mType);
registration_result = registration_methods.geodesic_shooting(moving_image, target_image, m0_reg, args.shoot_steps, mType, predict_network_config)
#endif
write_result(registration_result, output_prefixes[i]);
def predict_image(args):
if (args.use_CPU_for_shooting):
mType = ca.MEM_HOST
else:
mType = ca.MEM_DEVICE
# load the prediction network
predict_network_config = torch.load(args.prediction_parameter)
prediction_net = create_net(args, predict_network_config)
batch_size = args.batch_size
patch_size = predict_network_config['patch_size']
input_batch = torch.zeros(batch_size, 2, patch_size, patch_size,
patch_size).cuda()
# start prediction
for i in range(0, len(args.moving_image)):
common.Mkdir_p(os.path.dirname(args.output_prefix[i]))
if (args.affine_align):
# Perform affine registration to both moving and target image to the ICBM152 atlas space.
# Registration is done using Niftireg.
call([
"reg_aladin", "-noSym", "-speeeeed", "-ref", args.atlas,
"-flo", args.moving_image[i], "-res",
args.output_prefix[i] + "moving_affine.nii", "-aff",
args.output_prefix[i] + 'moving_affine_transform.txt'
])
call([
"reg_aladin", "-noSym", "-speeeeed", "-ref", args.atlas,
"-flo", args.target_image[i], "-res",
args.output_prefix[i] + "target_affine.nii", "-aff",
args.output_prefix[i] + 'target_affine_transform.txt'
])
moving_image = common.LoadITKImage(
args.output_prefix[i] + "moving_affine.nii", mType)
target_image = common.LoadITKImage(
args.output_prefix[i] + "target_affine.nii", mType)
else:
moving_image = common.LoadITKImage(args.moving_image[i], mType)
target_image = common.LoadITKImage(args.target_image[i], mType)
#preprocessing of the image
moving_image_np = preprocess_image(moving_image, args.histeq)
target_image_np = preprocess_image(target_image, args.histeq)
grid = moving_image.grid()
moving_image_processed = common.ImFromNPArr(moving_image_np, mType)
target_image_processed = common.ImFromNPArr(target_image_np, mType)
moving_image.setGrid(grid)
target_image.setGrid(grid)
predict_transform_space = False
if 'matlab_t7' in predict_network_config:
predict_transform_space = True
# run actual prediction
prediction_result = util.predict_momentum(moving_image_np,
target_image_np, input_batch,
batch_size, patch_size,
prediction_net,
predict_transform_space)
m0 = prediction_result['image_space']
m0_reg = common.FieldFromNPArr(prediction_result['image_space'], mType)
registration_result = registration_methods.geodesic_shooting(
moving_image_processed, target_image_processed, m0_reg,
args.shoot_steps, mType, predict_network_config)
phi = common.AsNPCopy(registration_result['phiinv'])
phi_square = np.power(phi, 2)
for sample_iter in range(1, args.samples):
print(sample_iter)
prediction_result = util.predict_momentum(
moving_image_np, target_image_np, input_batch, batch_size,
patch_size, prediction_net, predict_transform_space)
m0 += prediction_result['image_space']
m0_reg = common.FieldFromNPArr(prediction_result['image_space'],
mType)
registration_result = registration_methods.geodesic_shooting(
moving_image_processed, target_image_processed, m0_reg,
args.shoot_steps, mType, predict_network_config)
phi += common.AsNPCopy(registration_result['phiinv'])
phi_square += np.power(
common.AsNPCopy(registration_result['phiinv']), 2)
m0_mean = np.divide(m0, args.samples)
m0_reg = common.FieldFromNPArr(m0_mean, mType)
registration_result = registration_methods.geodesic_shooting(
moving_image_processed, target_image_processed, m0_reg,
args.shoot_steps, mType, predict_network_config)
phi_mean = registration_result['phiinv']
phi_var = np.divide(phi_square, args.samples) - np.power(
np.divide(phi, args.samples), 2)
#save result
common.SaveITKImage(registration_result['I1'],
args.output_prefix[i] + "I1.mhd")
common.SaveITKField(phi_mean,
args.output_prefix[i] + "phiinv_mean.mhd")
common.SaveITKField(common.FieldFromNPArr(phi_var, mType),
args.output_prefix[i] + "phiinv_var.mhd")
def main():
# Extract the Monkey number and section number from the command line
global frgNum
global secOb
mkyNum = sys.argv[1]
secNum = sys.argv[2]
frgNum = int(sys.argv[3])
write = True
# if not os.path.exists(os.path.expanduser('~/korenbergNAS/3D_database/Working/configuration_files/SidescapeRelateBlockface/M{0}/section_{1}/include_configFile.yaml'.format(mkyNum,secNum))):
# cf = initial(secNum, mkyNum)
try:
secOb = Config.Load(
secSpec,
pth.expanduser(
'~/korenbergNAS/3D_database/Working/configuration_files/SidescapeRelateBlockface/M{0}/section_{1}/include_configFile.yaml'
.format(mkyNum, secNum)))
except IOError as e:
try:
temp = Config.LoadYAMLDict(pth.expanduser(
'~/korenbergNAS/3D_database/Working/configuration_files/SidescapeRelateBlockface/M{0}/section_{1}/include_configFile.yaml'
.format(mkyNum, secNum)),
include=False)
secOb = Config.MkConfig(temp, secSpec)
except IOError:
print 'It appears there is no configuration file for this section. Please initialize one and restart.'
sys.exit()
if frgNum == int(secOb.yamlList[frgNum][-6]):
Fragmenter()
try:
secOb = Config.Load(
secSpec,
pth.expanduser(
'~/korenbergNAS/3D_database/Working/configuration_files/SidescapeRelateBlockface/M{0}/section_{1}/include_configFile.yaml'
.format(mkyNum, secNum)))
except IOError:
print 'It appeas that the include yaml file list does not match your fragmentation number. Please check them and restart.'
sys.exit()
if not pth.exists(
pth.expanduser(secOb.ssiOutPath + 'frag{0}'.format(frgNum))):
common.Mkdir_p(
pth.expanduser(secOb.ssiOutPath + 'frag{0}'.format(frgNum)))
if not pth.exists(
pth.expanduser(secOb.bfiOutPath + 'frag{0}'.format(frgNum))):
common.Mkdir_p(
pth.expanduser(secOb.bfiOutPath + 'frag{0}'.format(frgNum)))
if not pth.exists(
pth.expanduser(secOb.ssiSrcPath + 'frag{0}'.format(frgNum))):
os.mkdir(pth.expanduser(secOb.ssiSrcPath + 'frag{0}'.format(frgNum)))
if not pth.exists(
pth.expanduser(secOb.bfiSrcPath + 'frag{0}'.format(frgNum))):
os.mkdir(pth.expanduser(secOb.bfiSrcPath + 'frag{0}'.format(frgNum)))
frgOb = Config.MkConfig(secOb.yamlList[frgNum], frgSpec)
ssiSrc, bfiSrc, ssiMsk, bfiMsk = Loader(frgOb, ca.MEM_HOST)
#Extract the saturation Image from the color iamge
bfiHsv = common.FieldFromNPArr(
matplotlib.colors.rgb_to_hsv(
np.rollaxis(np.array(np.squeeze(bfiSrc.asnp())), 0, 3)),
ca.MEM_HOST)
bfiHsv.setGrid(bfiSrc.grid())
bfiSat = ca.Image3D(bfiSrc.grid(), bfiHsv.memType())
ca.Copy(bfiSat, bfiHsv, 1)
#Histogram equalize, normalize and mask the blockface saturation image
bfiSat = cb.HistogramEqualize(bfiSat, 256)
bfiSat.setGrid(bfiSrc.grid())
bfiSat *= -1
bfiSat -= ca.Min(bfiSat)
bfiSat /= ca.Max(bfiSat)
bfiSat *= bfiMsk
bfiSat.setGrid(bfiSrc.grid())
#Write out the blockface region after adjusting the colors with a format that supports header information
if write:
common.SaveITKImage(
bfiSat,
pth.expanduser(secOb.bfiSrcPath +
'frag{0}/M{1}_01_bfi_section_{2}_frag{0}_sat.nrrd'.
format(frgNum, secOb.mkyNum, secOb.secNum)))
#Set the sidescape grid relative to that of the blockface
ssiSrc.setGrid(ConvertGrid(ssiSrc.grid(), bfiSat.grid()))
ssiMsk.setGrid(ConvertGrid(ssiMsk.grid(), bfiSat.grid()))
ssiSrc *= ssiMsk
#Write out the sidescape masked image in a format that stores the header information
if write:
common.SaveITKImage(
ssiSrc,
pth.expanduser(secOb.ssiSrcPath +
'frag{0}/M{1}_01_ssi_section_{2}_frag{0}.nrrd'.
format(frgNum, secOb.mkyNum, secOb.secNum)))
#Update the image parameters of the sidescape image for future use
frgOb.imSize = ssiSrc.size().tolist()
frgOb.imOrig = ssiSrc.origin().tolist()
frgOb.imSpac = ssiSrc.spacing().tolist()
updateFragOb(frgOb)
#Find the affine transform between the two fragments
bfiAff, ssiAff, aff = Affine(bfiSat, ssiSrc, frgOb)
updateFragOb(frgOb)
#Write out the affine transformed images in a format that stores header information
if write:
common.SaveITKImage(
bfiAff,
pth.expanduser(
secOb.bfiOutPath +
'frag{0}/M{1}_01_bfi_section_{2}_frag{0}_aff_ssi.nrrd'.format(
frgNum, secOb.mkyNum, secOb.secNum)))
common.SaveITKImage(
ssiAff,
pth.expanduser(
secOb.ssiOutPath +
'frag{0}/M{1}_01_ssi_section_{2}_frag{0}_aff_bfi.nrrd'.format(
frgNum, secOb.mkyNum, secOb.secNum)))
bfiVe = bfiAff.copy()
ssiVe = ssiSrc.copy()
cc.VarianceEqualize_I(bfiVe, sigma=frgOb.sigVarBfi, eps=frgOb.epsVar)
cc.VarianceEqualize_I(ssiVe, sigma=frgOb.sigVarSsi, eps=frgOb.epsVar)
#As of right now, the largest pre-computed FFT table is 2048, so resample onto that grid for registration
regGrd = ConvertGrid(
cc.MakeGrid(ca.Vec3Di(2048, 2048, 1), ca.Vec3Df(1, 1, 1),
ca.Vec3Df(0, 0, 0)), ssiSrc.grid())
ssiReg = ca.Image3D(regGrd, ca.MEM_HOST)
bfiReg = ca.Image3D(regGrd, ca.MEM_HOST)
cc.ResampleWorld(ssiReg, ssiVe)
cc.ResampleWorld(bfiReg, bfiVe)
#Create the default configuration object for IDiff Matching and then set some parameters
idCf = Config.SpecToConfig(IDiff.Matching.MatchingConfigSpec)
idCf.compute.useCUDA = True
idCf.io.outputPrefix = '/home/sci/blakez/IDtest/'
#Run the registration
ssiDef, phi = DefReg(ssiReg, bfiReg, frgOb, ca.MEM_DEVICE, idCf)
#Turn the deformation into a displacement field so it can be applied to the large tif with C++ code
affV = phi.copy()
cc.ApplyAffineReal(affV, phi, np.linalg.inv(frgOb.affine))
ca.HtoV_I(affV)
#Apply the found deformation to the input ssi
ssiSrc.toType(ca.MEM_DEVICE)
cc.HtoReal(phi)
affPhi = phi.copy()
ssiBfi = ssiSrc.copy()
upPhi = ca.Field3D(ssiSrc.grid(), phi.memType())
cc.ApplyAffineReal(affPhi, phi, np.linalg.inv(frgOb.affine))
cc.ResampleWorld(upPhi, affPhi, bg=2)
cc.ApplyHReal(ssiBfi, ssiSrc, upPhi)
# ssiPhi = ca.Image3D(ssiSrc.grid(), phi.memType())
# upPhi = ca.Field3D(ssiSrc.grid(), phi.memType())
# cc.ResampleWorld(upPhi, phi, bg=2)
# cc.ApplyHReal(ssiPhi, ssiSrc, upPhi)
# ssiBfi = ssiSrc.copy()
# cc.ApplyAffineReal(ssiBfi, ssiPhi, np.linalg.inv(frgOb.affine))
# #Apply affine to the deformation
# affPhi = phi.copy()
# cc.ApplyAffineReal(affPhi, phi, np.linalg.inv(frgOb.affine))
if write:
common.SaveITKImage(
ssiBfi,
pth.expanduser(
secOb.ssiOutPath +
'frag{0}/M{1}_01_ssi_section_{2}_frag{0}_def_bfi.nrrd'.format(
frgNum, secOb.mkyNum, secOb.secNum)))
cc.WriteMHA(
affPhi,
pth.expanduser(
secOb.ssiOutPath +
'frag{0}/M{1}_01_ssi_section_{2}_frag{0}_to_bfi_real.mha'.
format(frgNum, secOb.mkyNum, secOb.secNum)))
cc.WriteMHA(
affV,
pth.expanduser(
secOb.ssiOutPath +
'frag{0}/M{1}_01_ssi_section_{2}_frag{0}_to_bfi_disp.mha'.
format(frgNum, secOb.mkyNum, secOb.secNum)))
#Create the list of names that the deformation should be applied to
# nameList = ['M15_01_0956_SideLight_DimLED_10x_ORG.tif',
# 'M15_01_0956_TyrosineHydroxylase_Ben_10x_Stitching_c1_ORG.tif',
# 'M15_01_0956_TyrosineHydroxylase_Ben_10x_Stitching_c2_ORG.tif',
# 'M15_01_0956_TyrosineHydroxylase_Ben_10x_Stitching_c3_ORG.tif']
# appLarge(nameList, affPhi)
common.DebugHere()