def shard(self, worker_index, num_workers): custom_print('SHARDS: Worker %s/%s' % (worker_index + 1, num_workers)) assert (worker_index < num_workers) self.files = [ f for i, f in enumerate(self.all_files) if (i % num_workers) == worker_index ]
def fit(self, X, y): if "hooks" in self.params and "LogTotalSteps" in self.params["hooks"]: self.params["hooks"]["LogTotalSteps"][ "batch_size"] = self.input_fn_config["batch_size"] self.params["hooks"]["LogTotalSteps"][ "epochs"] = self.input_fn_config["num_epochs"] self.params["hooks"]["LogTotalSteps"]["train_size"] = train_size(X) with BaseTF.lock: config = self.config if BaseTF.num_instances > 1: config["model_dir"] = os.path.join( config["model_dir"], "inst-" + str(self.instance_id)) self.est_config = config self.estimator = tf.estimator.Estimator( model_fn=self.model_fn, params=self.params, config=tf.estimator.RunConfig(**config)) if self.feature_spec is None: self.feature_spec = feature_spec_from(X) tf.logging.set_verbosity(tf.logging.ERROR) try: self.fit_main_training_loop(X, y) except KeyboardInterrupt: custom_print("\nEarly stop of training, saving model...") self.export_estimator() return self else: self.export_estimator() return self
def create_graph(self, resized_image_tensor): input_tensor_name = self.model_info['resized_input_tensor_name'] try: with gfile.FastGFile(self.model_path, 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) custom_print('[Pretrained] Model loaded: ', self.model_path) model_features = tf.import_graph_def( graph_def, name='inceptionv3', input_map={ input_tensor_name: resized_image_tensor, }, return_elements=[ self.model_info['bottleneck_tensor_name'], ], ) return tf.reshape( model_features[0], [-1, self.model_info['bottleneck_tensor_size']], ) except NotFoundError: custom_print(('Unable to open file %s; please download the model' + ' there (%s)') % (self.model_path, self.model_info['data_url'])) exit(0)
def fit_main_training_loop(self, X, y): self.estimator.train( input_fn=self.gen_input_fn(X, y, True, self.input_fn_config)) evaluation_fn = self.gen_input_fn(X, y, False, self.input_fn_config) if evaluation_fn is None: custom_print("No evaluation data available - skipping evaluation.") return evaluation = self.estimator.evaluate(input_fn=evaluation_fn) custom_print(evaluation)
def _generate_masked_tractography(self, reseed_endpoints=False, affine=None, predictor=None): """Generate the tractography using the white matter mask. Args: reseed_endpoints: Boolean. If True, use the end points of the fibers produced to generate another tractography. This is to symmetrize the process. """ i = 0 while self.ongoing_fibers: i += 1 predictions = predictor(self._build_next_X(affine))["predictions"] directions = self.get_directions_from_predictions( predictions, affine) # Update the positions of the fibers and check if they are still ongoing cur_ongoing = [] for j, fiber in enumerate(self.ongoing_fibers): new_position = fiber[-1] + directions[j] * self.args.step_size if i == 1 and self._is_border(new_position): # First step is ambiguous and leads into boarder -> flip it. new_position = fiber[ -1] - directions[j] * self.args.step_size # Only continue fibers inside the boundaries and short enough if self._is_border(new_position) or \ i * self.args.step_size > self.max_fiber_length: self.tractography.append(fiber) else: fiber.append(new_position) cur_ongoing.append(fiber) self.ongoing_fibers = cur_ongoing end = "\r" if i % 25 == 0: end = "\n" custom_print("Round num:", '%4d' % i, "; ongoing:", '%7d' % len(self.ongoing_fibers), "; completed:", '%7d' % len(self.tractography), end=end) if reseed_endpoints: ending_seeds = [[fiber[-1]] for fiber in self.tractography] self.ongoing_fibers = ending_seeds self._generate_masked_tractography(reseed_endpoints=False, affine=affine)
def predict(self, X, args): """Generate the tracktography with the current model on the given brain.""" # Check model check_is_fitted(self, ["n_incoming", "block_size"]) assert isinstance(X, dict) #predictor = tf.contrib.predictor.from_saved_model(self._restore_path) predictor = self.predictor(self.feature_spec) self.args = args try: self.wm_mask = X['mask'] self.nii = X['dwi'] if 'max_fiber_length' in args: self.max_fiber_length = args.max_fiber_length else: self.max_fiber_length = 400 except KeyError as err: custom_print("KeyError: {}".format(err)) # Get brain information self.brain_data = X['dwi'] # If no seeds are specified, build them from the wm mask if 'seeds' not in self.args: seeds = self._seeds_from_wm_mask() # The final result will be here self.tractography = [] # Fibers that are still under construction. At first seeds. self.ongoing_fibers = seeds if predictor is not None: # Start tractography generation if 'reseed_endpoints' in self.args: self._generate_masked_tractography( self.args.reseed_endpoints, affine=X["header"]["vox_to_ras"], predictor=predictor) else: self._generate_masked_tractography( affine=X["header"]["vox_to_ras"], predictor=predictor) # Save the Fibers fiber_path = os.path.join(self.save_path, "fibers.trk") save_fibers(self.tractography, X["header"], fiber_path)
def __init__(self, run_config, sumatra_outcome_config, *args, **kwargs): self.is_model_first_run = True self.run_config = run_config self.sumatra_outcome_config = sumatra_outcome_config self.sumatra_outcome = {} self.sumatra_outcome['run_tags'] = [] tf_run_config = copy.deepcopy(run_config['tf_estimator_run_config']) if 'session_config' in tf_run_config: session_config = tf_run_config['session_config'] gpu_config = session_config.get('gpu_options') if gpu_config is not None: del session_config['gpu_options'] session_config_obj = tf.ConfigProto(**session_config) if gpu_config is not None: for conf_key, conf_val in gpu_config.items(): setattr(session_config_obj.gpu_options, conf_key, conf_val) tf_run_config['session_config'] = session_config_obj super(Estimator, self).__init__( config=tf_run_config, *args, **kwargs ) # Data provider provider = self.input_fn_config['data_provider'] module = getattr(src.data.providers, provider) if module is None: custom_print( 'FATAL: Data provider ' + self.input_fn_config['data_provider'] + ' not found' ) assert(module is not None) self.data_provider = module.DataProvider(self.input_fn_config) ft_def.all_features.feature_info[ft_def.MRI]['shape'] = \ self.data_provider.get_mri_shape() self.feature_spec = { name: tf.placeholder( shape=[1] + ft_info['shape'], dtype=ft_info['type'] ) for name, ft_info in ft_def.all_features.feature_info.items() }
def _seeds_from_wm_mask(self): """Compute the seeds for the streamlining from the white matter mask. This is invoked only if no seeds are specified. The seeds are selected on the interface between white and gray matter, i.e. they are the white matter voxels that have at least one gray matter neighboring voxel. These points are furthermore perturbed with some gaussian noise to have a wider range of starting points. Returns: seeds: The list of voxel that are seeds. """ # Take te border voxels as seeds seeds = self._find_borders() custom_print("Number of seeds on the white matter mask:", len(seeds)) custom_print("Number of requested seeds:", self.args.n_fibers) new_idxs = np.random.choice(len(seeds), self.args.n_fibers, replace=True) new_seeds = [[ seeds[i] + np.clip(np.random.normal(0, 0.25, 3), -0.5, 0.5) ] for i in new_idxs] return new_seeds
def filter_xml(self, files, xml_image_id, filters, xml_class=None, xml_patient_id=None): self.image_id_enabled = set() self.image_id_to_class = {} discarded_count = 0 for f in glob.glob(files): tree = ET.parse(f) root = tree.getroot() image_id = int(xml_elem_unique(root, xml_image_id)) pass_all_filters = True for filter in filters: value = xml_elem_unique(root, filter['key'], allowNone=True) if not filters_match(value, filter['value']): pass_all_filters = False break if not pass_all_filters: discarded_count += 1 continue self.image_id_enabled.add(image_id) if xml_class is not None: self.image_id_to_class[image_id] = xml_elem_unique( root, xml_class, ) if xml_patient_id is not None: self.image_id_to_patient_id[image_id] = xml_elem_unique( root, xml_patient_id, ) custom_print('[filter_xml] %s images discarded' % (discarded_count))
def transform(self, X=None): steps_registered = { 'no_operation': self.no_operation, 'exec_command': self.exec_command, 'brain_extraction': self.brain_extraction, 'template_registration': self.template_registration, 'image_crop': self.image_crop, 'eddy_correct': self.eddy_correct, 'dtifit': self.dtifit, } for step in self.steps: self._mkdir(step['subfolder']) custom_print('Applying MRI pipeline to %s files' % (len(self.files))) all_images_ids = [] for i, mri_raw in enumerate(self.files): image_id, patient_id = self.extract_image_and_patient(mri_raw) custom_print('Image %s/%s [image_id = %s]' % (i, len(self.files), image_id)) if self.image_id_enabled is not None: if image_id not in self.image_id_enabled: custom_print('... Skipped') continue for step_id, step in enumerate(self.steps): if 'skip' in step: continue if step['from_subfolder'] == 'raw': path_from = mri_raw else: path_from = os.path.join( self.path, step['from_subfolder'], 'I{image_id}.nii.gz'.format(image_id=image_id), ) path_to = os.path.join( self.path, step['subfolder'], 'I{image_id}.nii.gz'.format(image_id=image_id), ) if not os.path.exists(path_from): continue if os.path.exists(path_to) and not step['overwrite']: continue steps_registered[step['type']]( path_from, path_to, image_id, step, ) all_images_ids.append(image_id) # Split train/test if self.split_train_test is not None: self.do_split_train_test(all_images_ids, **self.split_train_test)
def fit_main_training_loop(self, X, y): if self.streamer is not None: self.streamer.dump_split(self.save_path) self.streamer.dump_normalization(self.save_path) self.streamer.dump_train_val_test_split(self.save_path) n_epochs = self.input_fn_config["num_epochs"] self.n_epochs = n_epochs self.input_fn_config["num_epochs"] = 1 # Compute the number of steps per epoch #n_train_steps = self.streamer.get_number_train_batches() #self.config["keep_checkpoint_max"] = n_epochs #self.config["save_checkpoints_steps"] = n_train_steps + 1 output_dir = self.config["model_dir"] self.metric_logger = MetricLogger(output_dir, "Evaluation metrics") for i in range(n_epochs): self.current_epoch = i # train self.params["validation"] = False self.estimator = tf.estimator.Estimator( model_fn=self.model_fn, params=self.params, config=tf.estimator.RunConfig(**self.est_config)) self.estimator.train(input_fn=self.gen_input_fn( X, y, "train", self.input_fn_config)) # validation self.params["validation"] = True self.estimator = tf.estimator.Estimator( model_fn=self.model_fn, params=self.params, config=tf.estimator.RunConfig(**self.est_config)) if "do_validation" in self.sumatra_params and \ self.sumatra_params["do_validation"]: validation_fn = self.gen_input_fn(X, y, "validation", self.input_fn_config) if validation_fn is None: custom_print("No evaluation - skipping evaluation.") return validation = self.estimator.evaluate(input_fn=validation_fn, name="validation") print(validation) self.metric_logger.add_evaluations("validation", validation) # evaluate # evaluation on test set self.params["validation"] = False self.estimator = tf.estimator.Estimator( model_fn=self.model_fn, params=self.params, config=tf.estimator.RunConfig(**self.est_config)) if "no_test" not in self.sumatra_params: evaluation_fn = self.gen_input_fn(X, y, "test", self.input_fn_config) if evaluation_fn is None: custom_print("No evaluation - skipping evaluation.") return evaluation = self.estimator.evaluate(input_fn=evaluation_fn, name="test") print(evaluation) self.metric_logger.add_evaluations("test", evaluation) # persist evaluations to json file self.metric_logger.dump() sys.stdout.flush() if "keep_epoch_checkpoints" in self.data_params and \ self.data_params['keep_epoch_checkpoints']: # Copy checkpoint to new epoch folder dest_path = os.path.join(self.save_path, "epoch{}".format(i)) os.makedirs(dest_path) for fname in os.listdir(self.save_path): full_path = os.path.join(self.save_path, fname) if os.path.isfile( full_path) and "outcome" not in full_path: shutil.copy(full_path, dest_path) if "keep_embeddings" in self.data_params: self.compress_data(True) else: self.compress_data(False) if "keep_checkpoint" in self.data_params: if not self.data_params["keep_checkpoint"]: self.remove_checkpoints() else: self.remove_checkpoints() if "keep_tfevents" in self.data_params: if not self.data_params["keep_tfevents"]: self.remove_tfevents() else: self.remove_tfevents() self.streamer = None
def export_estimator(self): receiver_fn = input_receiver_fn(self.feature_spec) self._restore_path = self.estimator.export_savedmodel( self.save_path, receiver_fn) custom_print("Model saved to {}".format(self._restore_path))
def fit_main_training_loop(self, X, y): """ Trains and runs validation regularly at the same time """ self.evaluations = [] self.training_metrics = [] custom_print('[INFO] Main training loop. Model dir is %s' % ( self.config["model_dir"])) if 'reason' in self.sumatra_outcome_config: custom_print('[INFO] Provided run reason in config: %s' % ( self.sumatra_outcome_config['reason'], )) # Dump list of train/test images dataset = self.data_provider.export_dataset() if dataset is not None: with open('%s/dataset.json' % ( self.config["model_dir"] ), 'w') as outfile: json.dump(dataset, outfile) printed_count = [0] # Workaround to modify variable inside nested func def one_run_finished(t='.'): sys.stdout.write(t) printed_count[0] += 1 if printed_count[0] % 10 == 0: sys.stdout.write('\n') sys.stdout.flush() self.write_outcome() def do_evaluate(): evaluate_fn = self.gen_input_fn(X, y, False, self.input_fn_config) assert(evaluate_fn is not None) self.evaluations.append( self.estimator.evaluate(input_fn=evaluate_fn) ) one_run_finished('V') num_epochs = self.run_config['num_epochs'] validations_per_epoch = self.run_config['validations_per_epoch'] one_run_finished('\n') # 1st case, evaluation every few epochs if validations_per_epoch <= 1: validation_counter = 0 train_fn = self.gen_input_fn(X, y, True, self.input_fn_config) for i in range(num_epochs): self.estimator.train(input_fn=train_fn) one_run_finished() # Check if we need to run validation validation_counter += validations_per_epoch if validation_counter >= 1: validation_counter -= 1 do_evaluate() # 2nd case, several evaluations per epoch (should be an int then!) else: iters = num_epochs * validations_per_epoch for i in range(iters): train_fn = self.gen_input_fn( X, y, True, self.input_fn_config, shard=(i % validations_per_epoch, validations_per_epoch), ) self.estimator.train(input_fn=train_fn) one_run_finished() do_evaluate()
def do_split_train_test( self, image_ids, random_seed, pkl_prefix, test_images_def, ): if self.image_id_to_class is None: if self.set_all_single_class is None: custom_print('Train/Test split: No class loaded from XML.') return self.image_id_to_class = { img_id: self.set_all_single_class for img_id in image_ids } num_images = len(image_ids) image_ids = [id for id in image_ids if id in self.image_id_to_class] if num_images != len(image_ids): custom_print('Train/Test split: %d/%d images with unknown class!' % ( num_images - len(image_ids), num_images, )) # r = random.Random(random_seed) all_test = [] all_train = [] patients_dict = {} for class_idx, class_def in enumerate(test_images_def): if not isinstance(class_def['class'], list): assert (isinstance(class_def['class'], str)) class_def['class'] = [class_def['class']] # We want to split by patient ID # and reach a minimum number of images class_images = [ img_id for img_id in image_ids if self.image_id_to_class[img_id] in class_def['class'] ] class_images = sorted( class_images, key=self.image_id_to_patient_id.__getitem__, ) custom_print('Class %d [%s]' % ( class_idx, ','.join(class_def['class']), )) num_test_images = class_def['count'] if num_test_images == 0: test_images = [] train_images = class_images else: last_test_patient_id = self.image_id_to_patient_id[ class_images[num_test_images - 1]] while self.image_id_to_patient_id[ class_images[num_test_images]] == last_test_patient_id: num_test_images += 1 test_images = class_images[:num_test_images] train_images = class_images[num_test_images:] custom_print(' TRAIN: %d images / %d patients' % ( len(train_images), len( set([ self.image_id_to_patient_id[img_id] for img_id in train_images ])), )) custom_print(' VALID: %d images / %d patients' % ( len(test_images), len( set([ self.image_id_to_patient_id[img_id] for img_id in test_images ])), )) all_test += test_images all_train += train_images patients_dict.update( {'I%d' % id: class_idx for id in class_images}) pickle.dump( ['I%d' % i for i in all_test], open(os.path.join(self.path, '%stest.pkl' % pkl_prefix), 'wb'), ) pickle.dump( ['I%d' % i for i in all_train], open(os.path.join(self.path, '%strain.pkl' % pkl_prefix), 'wb'), ) pickle.dump( patients_dict, open(os.path.join(self.path, '%slabels.pkl' % pkl_prefix), 'wb'), )
def print_shape(self, text): if self.enable_print_shapes: custom_print(text)
def log(text): if text in UniqueLogger.printed: return UniqueLogger.printed.add(text) custom_print(text)
def _exec(self, cmd): custom_print('[Exec] ' + cmd) os.environ['FSLOUTPUTTYPE'] = 'NIFTI_GZ' os.environ['FSLDIR'] = '/local/fsl' subprocess.call(cmd, shell=True)