def __init__(self, experiment_path, tractography_name): self.description = tractography_name self.experiment_path = experiment_path self.name = os.path.basename(self.experiment_path) # self.logger_file = pjoin(self.experiment_path, "logger.pkl") self.results_file = pjoin(self.experiment_path, "results.json") self.tractometer_scores_file = pjoin(self.experiment_path, "tractometer", "scores", "{}.json".format(tractography_name)) self.hyperparams_file = pjoin(self.experiment_path, "hyperparams.json") # self.model_hyperparams_file = pjoin(self.experiment_path, "GRU_Regression", "hyperparams.json") self.status_file = pjoin(self.experiment_path, "training", "status.json") self.early_stopping_file = pjoin(self.experiment_path, "training", "tasks", "early_stopping.json") self.results = {} if os.path.isfile(self.results_file): self.results = load_dict_from_json_file(self.results_file) self.hyperparams = load_dict_from_json_file(self.hyperparams_file) # self.model_hyperparams = load_dict_from_json_file(self.model_hyperparams_file) self.status = load_dict_from_json_file(self.status_file) self.tractometer_scores = {} if os.path.isfile(self.tractometer_scores_file): self.tractometer_scores = load_dict_from_json_file(self.tractometer_scores_file) elif os.path.isfile(self.tractometer_scores_file[:-5] + '.pkl'): self.tractometer_scores = pickle.load(open(self.tractometer_scores_file[:-5] + '.pkl', 'rb')) else: print("No tractometer results yet for: {}".format(self.tractometer_scores_file)) self.early_stopping = {} if os.path.isfile(self.early_stopping_file): self.early_stopping = load_dict_from_json_file(self.early_stopping_file)
def main(): parser = build_args_parser() args = parser.parse_args() print(args) # Get experiment folder experiment_path = args.name if not os.path.isdir(experiment_path): # If not a directory, it must be the name of the experiment. experiment_path = pjoin(".", "experiments", args.name) if not os.path.isdir(experiment_path): parser.error('Cannot find experiment: {0}!'.format(args.name)) # Load experiments hyperparameters try: hyperparams = smartutils.load_dict_from_json_file(pjoin(experiment_path, "hyperparams.json")) except FileNotFoundError: hyperparams = smartutils.load_dict_from_json_file(pjoin(experiment_path, "..", "hyperparams.json")) with Timer("Loading dataset", newline=True): volume_manager = VolumeManager() dataset = datasets.load_tractography_dataset([args.streamlines], volume_manager, name="dataset", use_sh_coeffs=hyperparams['use_sh_coeffs']) print("Dataset size:", len(dataset)) with Timer("Loading model"): model = None if hyperparams['model'] == 'gru_regression': from learn2track.models import GRU_Regression model = GRU_Regression.create(experiment_path, volume_manager=volume_manager) elif hyperparams['model'] == 'gru_mixture': from learn2track.models import GRU_Mixture model = GRU_Mixture.create(experiment_path, volume_manager=volume_manager) elif hyperparams['model'] == 'gru_multistep': from learn2track.models import GRU_Multistep_Gaussian model = GRU_Multistep_Gaussian.create(experiment_path, volume_manager=volume_manager) model.k = 1 model.m = 1 elif hyperparams['model'] == 'ffnn_regression': from learn2track.models import FFNN_Regression model = FFNN_Regression.create(experiment_path, volume_manager=volume_manager) else: raise NameError("Unknown model: {}".format(hyperparams['model'])) print(str(model)) tractogram_file = pjoin(experiment_path, args.out) if not os.path.isfile(tractogram_file) or args.force: if args.method == 'prediction': tractogram = prediction_tractogram(hyperparams, model, dataset, args.batch_size, args.prediction) elif args.method == 'evaluation': tractogram = evaluation_tractogram(hyperparams, model, dataset, args.batch_size, args.metric) else: raise ValueError("Unrecognized method: {}".format(args.method)) tractogram.affine_to_rasmm = dataset.subjects[0].signal.affine nib.streamlines.save(tractogram, tractogram_file) else: print("Tractogram already exists. (use --force to generate it again)")
def create(cls, path): loaddir = pjoin(path, "model") meta = utils.load_dict_from_json_file(pjoin(loaddir, "meta.json")) assert meta['name'] == cls.__name__ hyperparams = utils.load_dict_from_json_file(pjoin(loaddir, "hyperparams.json")) model = cls(**hyperparams) model.load(path) return model
def main(): parser = build_parser() args = parser.parse_args() print(args) # Get experiment folder experiment_path = args.name if not os.path.isdir(experiment_path): # If not a directory, it must be the name of the experiment. experiment_path = pjoin(".", "experiments", args.name) if not os.path.isdir(experiment_path): parser.error("Cannot find experiment: {0}!".format(args.name)) # Load experiments hyperparameters try: hyperparams = smartutils.load_dict_from_json_file(pjoin(experiment_path, "hyperparams.json")) except FileNotFoundError: hyperparams = smartutils.load_dict_from_json_file(pjoin(experiment_path, "..", "hyperparams.json")) with Timer("Loading dataset", newline=True): volume_manager = VolumeManager() dataset = datasets.load_tractography_dataset( [args.subject], volume_manager, name="dataset", use_sh_coeffs=hyperparams["use_sh_coeffs"] ) print("Dataset size:", len(dataset)) with Timer("Loading model"): if hyperparams["model"] == "gru_regression": from learn2track.models import GRU_Regression model = GRU_Regression.create(experiment_path, volume_manager=volume_manager) else: raise NameError("Unknown model: {}".format(hyperparams["model"])) with Timer("Building evaluation function"): loss = loss_factory(hyperparams, model, dataset) batch_scheduler = batch_schedulers.TractographyBatchScheduler( dataset, batch_size=1000, noisy_streamlines_sigma=None, use_data_augment=False, # Otherwise it doubles the number of losses :-/ seed=1234, shuffle_streamlines=False, normalize_target=hyperparams["normalize"], ) loss_view = views.LossView(loss=loss, batch_scheduler=batch_scheduler) losses = loss_view.losses.view() with Timer("Saving streamlines"): tractogram = Tractogram(dataset.streamlines, affine_to_rasmm=dataset.subjects[0].signal.affine) tractogram.data_per_streamline["loss"] = losses nib.streamlines.save(tractogram, args.out)
def create(cls, path, **kwargs): loaddir = pjoin(path, cls.__name__) hyperparams = smartutils.load_dict_from_json_file(pjoin(loaddir, "hyperparams.json")) hyperparams.update(kwargs) model = cls(**hyperparams) model.load(path) return model
def maybe_create_experiment_folder(args, exclude=[], retrocompatibility_defaults={}): # Extract experiments hyperparameters hyperparams = OrderedDict(sorted(vars(args).items())) # Remove hyperparams that should not be part of the hash for name in exclude: if name in hyperparams: del hyperparams[name] # Get/generate experiment name experiment_name = args.name if experiment_name is None: experiment_name = generate_uid_from_string(repr(hyperparams)) # Create experiment folder experiment_path = pjoin(".", "experiments", experiment_name) resuming = False if os.path.isdir(experiment_path) and not args.force: resuming = True print("### Resuming experiment ({0}). ###\n".format(experiment_name)) # Check if provided hyperparams match those in the experiment folder hyperparams_loaded = smartutils.load_dict_from_json_file(pjoin(experiment_path, "hyperparams.json")) for name in exclude: if name in hyperparams_loaded: del hyperparams_loaded[name] for new_hyperparams, default_value in retrocompatibility_defaults.items(): if new_hyperparams not in hyperparams_loaded: hyperparams_loaded[new_hyperparams] = default_value if hyperparams != hyperparams_loaded: print("{\n" + "\n".join(["{}: {}".format(k, hyperparams[k]) for k in sorted(hyperparams.keys())]) + "\n}") print("{\n" + "\n".join(["{}: {}".format(k, hyperparams_loaded[k]) for k in sorted(hyperparams_loaded.keys())]) + "\n}") print("The arguments provided are different than the one saved. Use --force if you are certain.\nQuitting.") sys.exit(1) else: if os.path.isdir(experiment_path): shutil.rmtree(experiment_path) os.makedirs(experiment_path) smartutils.save_dict_to_json_file(pjoin(experiment_path, "hyperparams.json"), hyperparams) return experiment_path, hyperparams, resuming
def main(): parser = build_argparser() args = parser.parse_args() # Get experiment folder experiment_path = args.name if not os.path.isdir(experiment_path): # If not a directory, it must be the name of the experiment. experiment_path = pjoin(".", "experiments", args.name) if not os.path.isdir(experiment_path): parser.error('Cannot find experiment: {0}!'.format(args.name)) # Load experiments hyperparameters try: hyperparams = smartutils.load_dict_from_json_file(pjoin(experiment_path, "hyperparams.json")) except FileNotFoundError: hyperparams = smartutils.load_dict_from_json_file(pjoin(experiment_path, "..", "hyperparams.json")) with Timer("Loading DWIs"): # Load gradients table dwi_name = args.dwi if dwi_name.endswith(".gz"): dwi_name = dwi_name[:-3] if dwi_name.endswith(".nii"): dwi_name = dwi_name[:-4] bvals_filename = dwi_name + ".bvals" bvecs_filename = dwi_name + ".bvecs" bvals, bvecs = dipy.io.gradients.read_bvals_bvecs(bvals_filename, bvecs_filename) dwi = nib.load(args.dwi) if hyperparams["use_sh_coeffs"]: # Use 45 spherical harmonic coefficients to represent the diffusion signal. weights = neurotools.get_spherical_harmonics_coefficients(dwi, bvals, bvecs).astype(np.float32) else: # Resample the diffusion signal to have 100 directions. weights = neurotools.resample_dwi(dwi, bvals, bvecs).astype(np.float32) affine_rasmm2dwivox = np.linalg.inv(dwi.affine) with Timer("Loading model"): if hyperparams["model"] == "gru_regression": from learn2track.models import GRU_Regression model_class = GRU_Regression elif hyperparams['model'] == 'gru_mixture': from learn2track.models import GRU_Mixture model_class = GRU_Mixture elif hyperparams['model'] == 'gru_multistep': from learn2track.models import GRU_Multistep_Gaussian model_class = GRU_Multistep_Gaussian elif hyperparams['model'] == 'ffnn_regression': from learn2track.models import FFNN_Regression model_class = FFNN_Regression else: raise ValueError("Unknown model!") kwargs = {} volume_manager = neurotools.VolumeManager() volume_manager.register(weights) kwargs['volume_manager'] = volume_manager # Load the actual model. model = model_class.create(pjoin(experiment_path), **kwargs) # Create new instance and restore model. print(str(model)) mask = None if args.mask is not None: with Timer("Loading mask"): mask_nii = nib.load(args.mask) mask = mask_nii.get_data() # Compute the affine allowing to evaluate the mask at some coordinates correctly. # affine_maskvox2dwivox = mask_vox => rasmm space => dwi_vox affine_maskvox2dwivox = np.dot(affine_rasmm2dwivox, mask_nii.affine) if args.dilate_mask: import scipy mask = scipy.ndimage.morphology.binary_dilation(mask).astype(mask.dtype) with Timer("Generating seeds"): seeds = [] for filename in args.seeds: if filename.endswith('.trk') or filename.endswith('.tck'): tfile = nib.streamlines.load(filename) # Send the streamlines to voxel since that's where we'll track. tfile.tractogram.apply_affine(affine_rasmm2dwivox) # Use extremities of the streamlines as seeding points. seeds += [s[0] for s in tfile.streamlines] seeds += [s[-1] for s in tfile.streamlines] else: # Assume it is a binary mask. rng = np.random.RandomState(args.seeding_rng_seed) nii_seeds = nib.load(filename) # affine_seedsvox2dwivox = mask_vox => rasmm space => dwi_vox affine_seedsvox2dwivox = np.dot(affine_rasmm2dwivox, nii_seeds.affine) indices = np.array(np.where(nii_seeds.get_data())).T for idx in indices: seeds_in_voxel = idx + rng.uniform(-0.5, 0.5, size=(args.nb_seeds_per_voxel, 3)) seeds_in_voxel = nib.affines.apply_affine(affine_seedsvox2dwivox, seeds_in_voxel) seeds.extend(seeds_in_voxel) seeds = np.array(seeds, dtype=theano.config.floatX) with Timer("Tracking in the diffusion voxel space"): voxel_sizes = np.asarray(dwi.header.get_zooms()[:3]) if not np.all(voxel_sizes == dwi.header.get_zooms()[0]): print("* Careful voxel are anisotropic {}!".format(tuple(voxel_sizes))) # Since we are tracking in diffusion voxel space, convert step_size (in mm) to voxel. if args.step_size is not None: step_size = np.float32(args.step_size / voxel_sizes.max()) # Also convert max length (in mm) to voxel. max_nb_points = int(args.max_length / step_size) else: step_size = None max_nb_points = args.max_length if args.theta is not None: theta = np.deg2rad(args.theta) elif args.curvature is not None and args.curvature > 0: theta = get_max_angle_from_curvature(args.curvature, step_size) else: theta = np.deg2rad(45) print("Angle: {}".format(np.rad2deg(theta))) print("Step size (vox): {}".format(step_size)) print("Max nb. points: {}".format(max_nb_points)) is_outside_mask = make_is_outside_mask(mask, affine_maskvox2dwivox, threshold=args.mask_threshold) is_too_long = make_is_too_long(max_nb_points) is_too_curvy = make_is_too_curvy(np.rad2deg(theta)) is_stopping = make_is_stopping({STOPPING_MASK: is_outside_mask, STOPPING_LENGTH: is_too_long, STOPPING_CURVATURE: is_too_curvy}) is_stopping.max_nb_points = max_nb_points # Small hack tractogram = batch_track(model, weights, seeds, step_size=step_size, is_stopping=is_stopping, batch_size=args.batch_size, args=args) # Streamlines have been generated in voxel space. # Transform them them back to RAS+mm space using the dwi's affine. tractogram.affine_to_rasmm = dwi.affine tractogram.to_world() # Performed in-place. nb_streamlines = len(tractogram) print("Generated {:,} (compressed) streamlines".format(nb_streamlines)) with Timer("Cleaning streamlines", newline=True): # Flush streamlines that has no points. tractogram = tractogram[np.array(list(map(len, tractogram))) > 0] print("Removed {:,} empty streamlines".format(nb_streamlines - len(tractogram))) # Remove small streamlines nb_streamlines = len(tractogram) lengths = dipy.tracking.streamline.length(tractogram.streamlines) tractogram = tractogram[lengths >= args.min_length] lengths = lengths[lengths >= args.min_length] if len(lengths) > 0: print("Average length: {:.2f} mm.".format(lengths.mean())) print("Minimum length: {:.2f} mm. Maximum length: {:.2f}".format(lengths.min(), lengths.max())) print("Removed {:,} streamlines smaller than {:.2f} mm".format(nb_streamlines - len(tractogram), args.min_length)) if args.discard_stopped_by_curvature: nb_streamlines = len(tractogram) stopping_curvature_flag_is_set = is_flag_set(tractogram.data_per_streamline['stopping_flags'][:, 0], STOPPING_CURVATURE) tractogram = tractogram[stopping_curvature_flag_is_set] print("Removed {:,} streamlines stopped for having a curvature higher than {:.2f} degree".format(nb_streamlines - len(tractogram), np.rad2deg(theta))) if args.filter_threshold is not None: # Remove streamlines that produces a reconstruction error higher than a certain threshold. nb_streamlines = len(tractogram) losses = compute_loss_errors(tractogram.streamlines, model, hyperparams) print("Mean loss: {:.4f} ± {:.4f}".format(np.mean(losses), np.std(losses, ddof=1) / np.sqrt(len(losses)))) tractogram = tractogram[losses <= args.filter_threshold] print("Removed {:,} streamlines producing a loss lower than {:.2f} mm".format(nb_streamlines - len(tractogram), args.filter_threshold)) with Timer("Saving {:,} (compressed) streamlines".format(len(tractogram))): filename = args.out if args.out is None: prefix = args.prefix if prefix is None: dwi_name = os.path.basename(args.dwi) if dwi_name.endswith(".nii.gz"): dwi_name = dwi_name[:-7] else: # .nii dwi_name = dwi_name[:-4] prefix = os.path.basename(os.path.dirname(args.dwi)) + dwi_name prefix = prefix.replace(".", "_") mask_type = args.seeds[0].replace(".", "_").replace("_", "") if "int" in args.seeds[0]: mask_type = "int" elif "wm" in args.seeds[0]: mask_type = "wm" elif "rois" in args.seeds[0]: mask_type = "rois" filename = "{}-{}_seeding-{}_step-{:.2f}mm_nbSeeds-{}_maxAngle-{:.1f}deg_keepCurv-{}_filtered-{}_minLen-{}_pftRetry-{}_pftHist-{}_useMaxComponent-{}.tck".format( os.path.basename(args.name.rstrip('/'))[:6], prefix, mask_type, args.step_size, args.nb_seeds_per_voxel, np.rad2deg(theta), not args.discard_stopped_by_curvature, args.filter_threshold, args.min_length, args.pft_nb_retry, args.pft_nb_backtrack_steps, args.use_max_component ) save_path = pjoin(experiment_path, filename) try: # Create dirs, if needed. os.makedirs(os.path.dirname(save_path)) except: pass print("Saving to {}".format(save_path)) nib.streamlines.save(tractogram, save_path)
def main(): parser = build_parser() args = parser.parse_args() print(args) # Get experiment folder experiment_path = args.name if not os.path.isdir(experiment_path): # If not a directory, it must be the name of the experiment. experiment_path = pjoin(".", "experiments", args.name) if not os.path.isdir(experiment_path): parser.error('Cannot find experiment: {0}!'.format(args.name)) # Load experiments hyperparameters try: hyperparams = smartutils.load_dict_from_json_file(pjoin(experiment_path, "hyperparams.json")) except FileNotFoundError: hyperparams = smartutils.load_dict_from_json_file(pjoin(experiment_path, "..", "hyperparams.json")) with Timer("Loading dataset", newline=True): volume_manager = VolumeManager() dataset = datasets.load_tractography_dataset(args.subjects, volume_manager, name="dataset", use_sh_coeffs=hyperparams['use_sh_coeffs']) print("Dataset size:", len(dataset)) with Timer("Loading model"): model = None if hyperparams['model'] == 'gru_regression': from learn2track.models import GRU_Regression model = GRU_Regression.create(experiment_path, volume_manager=volume_manager) elif hyperparams['model'] == 'gru_mixture': from learn2track.models import GRU_Mixture model = GRU_Mixture.create(experiment_path, volume_manager=volume_manager) elif hyperparams['model'] == 'gru_multistep': from learn2track.models import GRU_Multistep_Gaussian model = GRU_Multistep_Gaussian.create(experiment_path, volume_manager=volume_manager) model.k = 1 model.m = 1 elif hyperparams['model'] == 'ffnn_regression': from learn2track.models import FFNN_Regression model = FFNN_Regression.create(experiment_path, volume_manager=volume_manager) else: raise NameError("Unknown model: {}".format(hyperparams['model'])) with Timer("Building evaluation function"): # Override K for gru_multistep if 'k' in hyperparams: hyperparams['k'] = 1 batch_scheduler = batch_scheduler_factory(hyperparams, dataset, train_mode=False, batch_size_override=args.batch_size) loss = loss_factory(hyperparams, model, dataset, loss_type=args.loss_type) l2_error = views.LossView(loss=loss, batch_scheduler=batch_scheduler) with Timer("Evaluating...", newline=True): results_file = pjoin(experiment_path, "results.json") results = {} if os.path.isfile(results_file) and not args.force: print("Loading saved results... (use --force to re-run evaluation)") results = smartutils.load_dict_from_json_file(results_file) tag = "" if args.loss_type == 'expected_value' or hyperparams['model'] == 'gru_regression': tag = "_EV_L2_error" elif args.loss_type == 'maximum_component': tag = "_MC_L2_error" elif hyperparams['model'] == 'gru_mixture' or hyperparams['model'] == 'gru_multistep': tag = "_NLL" entry = args.dataset_name + tag if entry not in results or args.force: with Timer("Evaluating {}".format(entry)): dummy_status = Status() # Forces recomputing results results[entry] = {'mean': float(l2_error.mean.view(dummy_status)), 'stderror': float(l2_error.stderror.view(dummy_status))} smartutils.save_dict_to_json_file(results_file, results) # Update results file. print("{}: {:.4f} ± {:.4f}".format(entry, results[entry]['mean'], results[entry]['stderror']))
def main(): parser = build_argparser() args = parser.parse_args() # Get experiment folder experiment_path = args.name if not os.path.isdir(experiment_path): # If not a directory, it must be the name of the experiment. experiment_path = pjoin(".", "experiments", args.name) if not os.path.isdir(experiment_path): parser.error('Cannot find experiment: {0}!'.format(args.name)) # Load experiments hyperparameters try: hyperparams = smartutils.load_dict_from_json_file( pjoin(experiment_path, "hyperparams.json")) except FileNotFoundError: hyperparams = smartutils.load_dict_from_json_file( pjoin(experiment_path, "..", "hyperparams.json")) with Timer("Loading DWIs"): # Load gradients table dwi_name = args.dwi if dwi_name.endswith(".gz"): dwi_name = dwi_name[:-3] if dwi_name.endswith(".nii"): dwi_name = dwi_name[:-4] bvals_filename = dwi_name + ".bvals" bvecs_filename = dwi_name + ".bvecs" bvals, bvecs = dipy.io.gradients.read_bvals_bvecs( bvals_filename, bvecs_filename) dwi = nib.load(args.dwi) if hyperparams["use_sh_coeffs"]: # Use 45 spherical harmonic coefficients to represent the diffusion signal. weights = neurotools.get_spherical_harmonics_coefficients( dwi, bvals, bvecs).astype(np.float32) else: # Resample the diffusion signal to have 100 directions. weights = neurotools.resample_dwi(dwi, bvals, bvecs).astype(np.float32) with Timer("Loading model"): if hyperparams["model"] == "ffnn_classification": from learn2track.models import FFNN_Classification model_class = FFNN_Classification else: raise ValueError("Unknown model!") kwargs = {} volume_manager = neurotools.VolumeManager() volume_manager.register(weights) kwargs['volume_manager'] = volume_manager # Load the actual model. model = model_class.create( pjoin(experiment_path), **kwargs) # Create new instance and restore model. print(str(model)) with Timer("Generating mask"): symb_input = T.matrix(name="input") model_symb_pred = model.get_output(symb_input) f = theano.function(inputs=[symb_input], outputs=[model_symb_pred]) generated_mask = np.zeros(dwi.shape[:3]).astype(np.float32) # all_coords.shape = (n_coords, 3) all_coords = np.argwhere(generated_mask == 0) volume_ids = np.zeros((all_coords.shape[0], 1)) all_coords_and_volume_ids = np.concatenate((all_coords, volume_ids), axis=1).astype(np.float32) batch_size = args.batch_size if args.batch_size else len( all_coords_and_volume_ids) probs = [] while batch_size > 1: print("Trying to to process batches of size {} out of {}".format( batch_size, len(all_coords_and_volume_ids))) nb_batches = int( np.ceil(len(all_coords_and_volume_ids) / batch_size)) try: for batch_count in range(nb_batches): start = batch_count * batch_size end = (batch_count + 1) * batch_size probs.extend(f(all_coords_and_volume_ids[start:end])[-1]) print("Generated batch {} out of {}".format( batch_count + 1, nb_batches)) break except MemoryError: print("{} coordinates at the same time is too much!".format( batch_size)) batch_size //= 2 except RuntimeError: print("{} coordinates at the same time is too much!".format( batch_size)) batch_size //= 2 if not probs: raise RuntimeError("Could not generate predictions...") generated_mask[np.where(generated_mask == 0)] = np.array(probs) > 0.5 with Timer("Saving generated mask"): filename = args.out if args.out is None: prefix = args.prefix if prefix is None: dwi_name = os.path.basename(args.dwi) if dwi_name.endswith(".nii.gz"): dwi_name = dwi_name[:-7] else: # .nii dwi_name = dwi_name[:-4] prefix = os.path.basename(os.path.dirname(args.dwi)) + dwi_name prefix = prefix.replace(".", "_") filename = "{}.nii.gz".format(prefix) save_path = pjoin(experiment_path, filename) try: # Create dirs, if needed. os.makedirs(os.path.dirname(save_path)) except: pass print("Saving to {}".format(save_path)) mask = nib.Nifti1Image(generated_mask, dwi.affine) nib.save(mask, save_path)