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()