def generate_csv(self, batch): # Get patient ID and make a prediction # if self.val_dataloader().batch_size != 1: # raise Exception("Batch size for validation must be 1!") pat_id = np.squeeze(batch['patient_list']) pat_path = np.squeeze(batch['patient_path_list']).tolist() image = batch['ct'] generated = self.generator(image) # if not self.opt.no_scaling: # generated = 40.0*generated + 40.0 # Scale back dose to 0 - 80 dose_pred_gy = generated * batch['possible_dose_mask'] dose_pred_gy = dose_pred_gy.view(1, 1, 128, 128, 128, 1) # Prepare the dose to save dose_pred_gy = np.squeeze(dose_pred_gy) dose_pred_gy = dose_pred_gy.cpu().numpy() dose_to_save = sparse_vector_function(dose_pred_gy) dose_df = pd.DataFrame(data=dose_to_save['data'].squeeze(), index=dose_to_save['indices'].squeeze(), columns=['data']) file_name = '{}/{}.csv'.format(self.prediction_dir, pat_id) dose_df.to_csv(file_name) return pat_path, file_name
def predict_dose(self, epoch=1): """Predicts the dose for the given epoch number, this will only work if the batch size of the data loader is set to 1. :param epoch: The epoch that should be loaded to make predictions """ # Define new models, or load most recent model if model already exists self.generator = load_model('{}{}.h5'.format(self.model_path_template, epoch)) os.makedirs(self.prediction_dir, exist_ok=True) # Use generator to predict dose number_of_batches = self.data_loader.number_of_batches() print('Predicting dose') for idx in tqdm.tqdm(range(number_of_batches)): image_batch = self.data_loader.get_batch(idx) # Get patient ID and make a prediction pat_id = image_batch['patient_list'][0] dose_pred_gy = self.generator.predict( [image_batch['ct'], image_batch['structure_masks']]) dose_pred_gy = dose_pred_gy * image_batch['possible_dose_mask'] # Prepare the dose to save dose_pred_gy = np.squeeze(dose_pred_gy) dose_to_save = sparse_vector_function(dose_pred_gy) dose_df = pd.DataFrame(data=dose_to_save['data'].squeeze(), index=dose_to_save['indices'].squeeze(), columns=['data']) dose_df.to_csv('{}/{}.csv'.format(self.prediction_dir, pat_id))
def conversion_step(self, batch, batch_idx): ct, mask, structure, pat_id = batch['ct'], batch[ 'possible_dose_mask'], batch['structure_masks'], batch[ 'patient_list'][0][0] pred = self(torch.cat([ct, structure], dim=1).float()) * mask * 100 dose_pred_gy = np.transpose( pred.squeeze(0).squeeze(0).cpu().numpy(), [1, 2, 0]) dose_to_save = sparse_vector_function(dose_pred_gy) dose_df = pd.DataFrame(data=dose_to_save['data'].squeeze(), index=dose_to_save['indices'].squeeze(), columns=['data']) dose_df.to_csv('{}/{}.csv'.format('checkpoints/results/', pat_id)) #dose_score = (torch.sum(torch.abs((dose * 100) - (pred * 100))) / torch.sum(mask)) #log_dict = {'val_loss': dose_score, 'train_dose_score': dose_score} return #{'val_loss': dose_score, 'log': log_dict, 'progress_bar': log_dict}