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}