class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# volume is a numpy array of shape [X,Y,Z] and I will slice X axis
slices = []
# create mask for each slice across the X (0th) dimension.
# put all slices into a 3D Numpy array
for ix in range(0, volume.shape[0]):
slice_tensor = torch.from_numpy(volume[ix, :, :].astype(
np.single)).unsqueeze(0).unsqueeze(0)
pred = self.model(slice_tensor.to(self.device))
mask = torch.argmax(np.squeeze(pred.cpu().detach()), dim=0)
slices.append(mask)
return np.dstack(slices).transpose(2, 0, 1)
def single_volume_inference_unpadded(self, volume, patch_size):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
volume = med_reshape(volume, (volume.shape[0], patch_size, patch_size))
return self.single_volume_inference(volume)
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self, parameter_file_path='', model=None, device="cpu", patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
raise NotImplementedError
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes
mask = np.zeros(volume.shape)
for i in range(0, volume.shape[0]):
ind_slice = torch.from_numpy(volume[i, : , : ].astype(np.single)).unsqueeze(0).unsqueeze(0)
pred = self.model(ind_slice.to(self.device))
mask[i,:, :] = torch.argmax(np.squeeze(pred.cpu().detach()), dim = 0)
return mask
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self, parameter_file_path='', model=None, device="cpu", patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
raise NotImplementedError
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
# <YOUR CODE HERE>
outputs = self.model(torch.tensor(volume).type(torch.cuda.FloatTensor).unsqueeze(1).to(self.device)).cpu().detach()
_, slices = torch.max(outputs, 1)
return slices.numpy()
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self, parameter_file_path='', model=None, device="cpu", patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
# Working through this: this scenario is the same as running `single_volume_inference`, but we have to
# pad it first (with the `med_reshape` function.
# First, reshape. Previous implementations look like this:
# med_reshape(image, new_shape=(image.shape[0], y_shape, z_shape))
# I thought we could pull this. For now, hard code
patch_size = 64
print("New shape ", volume.shape[0], patch_size, patch_size)
reshaped_volume = med_reshape(volume, new_shape=(volume.shape[0], patch_size, patch_size))
# Now we run the `singled_volume_inference` function on this reshaped volume
# We have precedent for this too:
# inference_agent.single_volume_inference(x["image"])
# I think we can run direclty on this reshaped_volume...
return self.single_volume_inference(reshaped_volume)
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = np.zeros(volume.shape)
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
# <YOUR CODE HERE>
for slice_position in range(volume.shape[0]):
# Note that this is different than the min max scaler used previously...
# I freely acknowledge that this is not ideal (but comparable).
norm_slice = volume[slice_position, :, :].astype(np.single) / 255.0
pred = self.model(torch.from_numpy(norm_slice).unsqueeze(0).unsqueeze(0).to(self.device))
pred_values = np.squeeze(pred.cpu().detach())
slices[slice_position, :, :] = torch.argmax(pred_values, dim=0)
return slices
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
# def single_volume_inference_unpadded(self, volume):
# """
# Runs inference on a single volume of arbitrary patch size,
# padding it to the conformant size first
# Arguments:
# volume {Numpy array} -- 3D array representing the volume
# Returns:
# 3D NumPy array with prediction mask
# """
# self.model.eval()
# # Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
# slices = []
# volume = med_reshape(volume, new_shape=(volume.shape[0], 64, 64))
# slices = np.zeros(volume.shape)
# def inference(img):
# tsr_test = torch.from_numpy(img.astype(np.single)/np.max(img)).unsqueeze(0).unsqueeze(0)
# print(tsr_test.shape)
# pred = self.model(tsr_test.to(self.device))
# return np.squeeze(pred.cpu().detach())
# for slc_ix in range(volume.shape[0]):
# pred = inference(volume[slc_ix,:,:])
# slices[slc_ix,:,:] = torch.argmax(pred, dim=0)
# return slices
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
new_volume = np.moveaxis(volume, -1, 0)
print("new volume", new_volume.shape)
reshaped_volume = med_reshape(new_volume,
new_shape=(new_volume.shape[0],
self.patch_size,
self.patch_size))
pred = self.single_volume_inference(reshaped_volume.astype(np.float32))
reshaped_pred = np.zeros_like(new_volume)
print("hey", reshaped_volume.shape, pred.shape, reshaped_pred.shape)
resize_dim = [reshaped_pred.shape[2], reshaped_pred.shape[1]]
print("lala", resize_dim)
for i in range(reshaped_pred.shape[0]):
reshaped_pred[i] = Image.fromarray(pred[i].astype(
np.uint8)).resize(resize_dim)
return np.moveaxis(reshaped_pred, 0, -1)
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
# <YOUR CODE HERE>
# slicing the x axis
#test_tensor = torch.tensor(volume).unsqueeze(1)
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
# outputs = self.model(test_tensor.float().to(self.device)).cpu().detach()
# _, slices = torch.max(outputs, 1)
# slices = slices.numpy
input_tensor = torch.tensor(volume).unsqueeze(1)
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
outputs = self.model(input_tensor.float().to(
self.device)).cpu().detach()
_, slices = torch.max(outputs, 1)
return slices.numpy()
class UNetExperiment:
"""
This class implements the basic life cycle for a segmentation task with UNet(https://arxiv.org/abs/1505.04597).
The basic life cycle of a UNetExperiment is:
run():
for epoch in n_epochs:
train()
validate()
test()
"""
def __init__(self, config, split, dataset):
self.n_epochs = config.n_epochs
self.split = split
self._time_start = ""
self._time_end = ""
self.epoch = 0
self.name = config.name
# Create output folders
dirname = f'{time.strftime("%Y-%m-%d_%H%M", time.gmtime())}_{self.name}'
self.out_dir = os.path.join(config.test_results_dir, dirname)
os.makedirs(self.out_dir, exist_ok=True)
# Create data loaders
# TASK: SlicesDataset class is not complete. Go to the file and complete it.
# Note that we are using a 2D version of UNet here, which means that it will expect
# batches of 2D slices.
self.train_loader = DataLoader(SlicesDataset(dataset[split["train"]]),
batch_size=config.batch_size,
shuffle=True,
num_workers=0)
self.val_loader = DataLoader(SlicesDataset(dataset[split["val"]]),
batch_size=config.batch_size,
shuffle=True,
num_workers=0)
# we will access volumes directly for testing
self.test_data = dataset[split["test"]]
# Do we have CUDA available?
if not torch.cuda.is_available():
print(
"WARNING: No CUDA device is found. This may take significantly longer!"
)
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
# Configure our model and other training implements
# We will use a recursive UNet model from German Cancer Research Center,
# Division of Medical Image Computing. It is quite complicated and works
# very well on this task. Feel free to explore it or plug in your own model
self.model = UNet(num_classes=3)
self.model.to(self.device)
# We are using a standard cross-entropy loss since the model output is essentially
# a tensor with softmax'd prediction of each pixel's probability of belonging
# to a certain class
self.loss_function = torch.nn.CrossEntropyLoss()
# We are using standard SGD method to optimize our weights
self.optimizer = optim.Adam(self.model.parameters(),
lr=config.learning_rate)
# Scheduler helps us update learning rate automatically
self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
self.optimizer, 'min')
# Set up Tensorboard. By default it saves data into runs folder. You need to launch
self.tensorboard_train_writer = SummaryWriter(comment="_train")
self.tensorboard_val_writer = SummaryWriter(comment="_val")
def train(self):
"""
This method is executed once per epoch and takes
care of model weight update cycle
"""
print(f"Training epoch {self.epoch}...")
self.model.train()
# Loop over our minibatches
for i, batch in enumerate(self.train_loader):
self.optimizer.zero_grad()
# TASK: You have your data in batch variable. Put the slices as 4D Torch Tensors of
# shape [BATCH_SIZE, 1, PATCH_SIZE, PATCH_SIZE] into variables data and target.
# Feed data to the model and feed target to the loss function
#
# data = <YOUR CODE HERE>
# target = <YOUR CODE HERE>
data = batch["image"].to(self.device, dtype=torch.float)
target = batch["seg"].to(self.device)
prediction = self.model(data)
# We are also getting softmax'd version of prediction to output a probability map
# so that we can see how the model converges to the solution
prediction_softmax = F.softmax(prediction, dim=1)
loss = self.loss_function(prediction, target[:, 0, :, :])
# TASK: What does each dimension of variable prediction represent?
# ANSWER: Dimensions represent: batch_size, classes, coronal data, axial data
loss.backward()
self.optimizer.step()
if (i % 10) == 0:
# Output to console on every 10th batch
print(
f"\nEpoch: {self.epoch} Train loss: {loss}, {100*(i+1)/len(self.train_loader):.1f}% complete"
)
counter = 100 * self.epoch + 100 * (i / len(self.train_loader))
# You don't need to do anything with this function, but you are welcome to
# check it out if you want to see how images are logged to Tensorboard
# or if you want to output additional debug data
log_to_tensorboard(self.tensorboard_train_writer, loss, data,
target, prediction_softmax, prediction,
counter)
print(".", end='')
print("\nTraining complete")
def validate(self):
"""
This method runs validation cycle, using same metrics as
Train method. Note that model needs to be switched to eval
mode and no_grad needs to be called so that gradients do not
propagate
"""
print(f"Validating epoch {self.epoch}...")
# Turn off gradient accumulation by switching model to "eval" mode
self.model.eval()
loss_list = []
with torch.no_grad():
for i, batch in enumerate(self.val_loader):
# TASK: Write validation code that will compute loss on a validation sample
# <YOUR CODE HERE>
data = batch["image"].to(self.device, dtype=torch.float)
target = batch["seg"].to(self.device)
prediction = self.model(data)
prediction_softmax = F.softmax(prediction, dim=1)
loss = self.loss_function(prediction, target[:, 0, :, :])
print(f"Batch {i}. Data shape {data.shape} Loss {loss}")
# We report loss that is accumulated across all of validation set
loss_list.append(loss.item())
self.scheduler.step(np.mean(loss_list))
log_to_tensorboard(self.tensorboard_val_writer, np.mean(loss_list),
data, target, prediction_softmax, prediction,
(self.epoch + 1) * 100)
print(f"Validation complete")
def save_model_parameters(self):
"""
Saves model parameters to a file in results directory
"""
path = os.path.join(self.out_dir, "model.pth")
torch.save(self.model.state_dict(), path)
def load_model_parameters(self, path=''):
"""
Loads model parameters from a supplied path or a
results directory
"""
if not path:
model_path = os.path.join(self.out_dir, "model.pth")
else:
model_path = path
if os.path.exists(model_path):
self.model.load_state_dict(torch.load(model_path))
else:
raise Exception(f"Could not find path {model_path}")
def run_test(self):
"""
This runs test cycle on the test dataset.
Note that process and evaluations are quite different
Here we are computing a lot more metrics and returning
a dictionary that could later be persisted as JSON
"""
print("Testing...")
self.model.eval()
# In this method we will be computing metrics that are relevant to the task of 3D volume
# segmentation. Therefore, unlike train and validation methods, we will do inferences
# on full 3D volumes, much like we will be doing it when we deploy the model in the
# clinical environment.
# TASK: Inference Agent is not complete. Go and finish it. Feel free to test the class
# in a module of your own by running it against one of the data samples
inference_agent = UNetInferenceAgent(model=self.model,
device=self.device)
out_dict = {}
out_dict["volume_stats"] = []
dc_list = []
jc_list = []
# for every in test set
for i, x in enumerate(self.test_data):
pred_label = inference_agent.single_volume_inference(x["image"])
# We compute and report Dice and Jaccard similarity coefficients which
# assess how close our volumes are to each other
# TASK: Dice3D and Jaccard3D functions are not implemented.
# Complete the implementation as we discussed
# in one of the course lessons, you can look up definition of Jaccard index
# on Wikipedia. If you completed it
# correctly (and if you picked your train/val/test split right ;)),
# your average Jaccard on your test set should be around 0.80
dc = Dice3d(pred_label, x["seg"])
jc = Jaccard3d(pred_label, x["seg"])
dc_list.append(dc)
jc_list.append(jc)
# STAND-OUT SUGGESTION: By way of exercise, consider also outputting:
# * Sensitivity and specificity (and explain semantic meaning in terms of
# under/over segmenting)
# * Dice-per-slice and render combined slices with lowest and highest DpS
# * Dice per class (anterior/posterior)
out_dict["volume_stats"].append({
"filename": x['filename'],
"dice": dc,
"jaccard": jc
})
print(
f"{x['filename']} Dice {dc:.4f}. {100*(i+1)/len(self.test_data):.2f}% complete"
)
out_dict["overall"] = {
"mean_dice": np.mean(dc_list),
"mean_jaccard": np.mean(jc_list)
}
print("\nTesting complete.")
return out_dict
def run(self):
"""
Kicks off train cycle and writes model parameter file at the end
"""
self._time_start = time.time()
print("Experiment started.")
# Iterate over epochs
for self.epoch in range(self.n_epochs):
self.train()
self.validate()
# save model for inferencing
self.save_model_parameters()
self._time_end = time.time()
print(
f"Run complete. Total time: {time.strftime('%H:%M:%S', time.gmtime(self._time_end - self._time_start))}"
)
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
volume_padded = med_reshape(volume,
new_shape=(self.patch_size,
self.patch_size,
self.patch_size))
prediction_mask = self.single_volume_inference(volume_padded)
return prediction_mask
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
# <YOUR CODE HERE>
for i in range(volume.shape[0]):
test_slice = torch.from_numpy(volume[i, :, :].astype(np.single) /
np.max(volume[i, :, :]))
test_slice = test_slice.unsqueeze(0).unsqueeze(0).to(self.device)
pred = self.model(test_slice)
# print(pred.shape) # [1, 3, 64, 64]
cpu_pred = pred.cpu()
# class 0 is no hippocampus and 1,2 are the two segments of hippocampi
result = cpu_pred.detach().numpy()[
0] # does .data instead of .detach() work ?
result = np.argmax(
result, axis=0
) # as fast as torch.argmax and doesn't return tensor but np array
slices.append(result)
return np.stack(slices)
class UNetExperiment:
"""
This class implements the basic life cycle for a segmentation task with UNet(https://arxiv.org/abs/1505.04597).
The basic life cycle of a UNetExperiment is:
run():
for epoch in n_epochs:
train()
validate()
test()
"""
def __init__(self, config, split, dataset):
self.n_epochs = config.n_epochs
self.split = split
self._time_start = ""
self._time_end = ""
self.epoch = 0
self.name = config.name
# Create output folders
dirname = f'{time.strftime("%Y-%m-%d_%H%M", time.gmtime())}_{self.name}'
self.out_dir = os.path.join(config.test_results_dir, dirname)
os.makedirs(self.out_dir, exist_ok=True)
# Create data loaders
self.train_loader = DataLoader(SlicesDataset(dataset[split["train"]]),
batch_size=config.batch_size, shuffle=True, num_workers=0)
self.val_loader = DataLoader(SlicesDataset(dataset[split["val"]]),
batch_size=config.batch_size, shuffle=True, num_workers=0)
# access volumes directly for testing
self.test_data = dataset[split["test"]]
if not torch.cuda.is_available():
print("WARNING: No CUDA device is found. This may take significantly longer!")
#self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.device = torch.device('cpu')
# use a recursive UNet model from German Cancer Research Center, Division of Medical Image Computing
self.model = UNet()
self.model.to(self.device)
# use a standard cross-entropy loss since the model output is essentially
# a tensor with softmax prediction of each pixel's probability of belonging to a certain class
self.loss_function = torch.nn.CrossEntropyLoss()
# use standard SGD method to optimize the weights
self.optimizer = optim.Adam(self.model.parameters(), lr=config.learning_rate)
# Scheduler helps to update learning rate automatically
self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, 'min')
# Set up Tensorboard. By default it saves data into runs folder. You need to launch
# self.tensorboard_train_writer = SummaryWriter(comment="_train")
# self.tensorboard_val_writer = SummaryWriter(comment="_val")
def train(self):
"""
This method is executed once per epoch and takes
care of model weight update cycle
"""
print(f"Training epoch {self.epoch}...")
self.model.train()
# Loop over the minibatches
for i, batch in enumerate(self.train_loader):
self.optimizer.zero_grad()
# Feed data to the model and feed target to the loss function
data = batch['image'].float()
target = batch['seg']
prediction = self.model(data.to(self.device))
prediction_softmax = F.softmax(prediction, dim=1)
loss = self.loss_function(prediction_softmax, target[:, 0, :, :].to(self.device))
# What does each dimension of variable prediction represent?
# batch_size, 3 classes, coronal, axial
loss.backward()
self.optimizer.step()
if (i % 10) == 0:
# Output to console on every 10th batch
print(f"\nEpoch: {self.epoch} Train loss: {loss}, {100*(i+1)/len(self.train_loader):.1f}% complete")
counter = 100*self.epoch + 100*(i/len(self.train_loader))
# log_to_tensorboard(
# self.tensorboard_train_writer,
# loss,
# data,
# target,
# prediction_softmax,
# prediction,
# counter)
print(".", end='')
print("\nTraining complete")
def validate(self):
"""
This method runs validation cycle, using same metrics as
Train method. Note that model needs to be switched to eval
mode and no_grad needs to be called so that gradients do not
propagate
"""
print(f"Validating epoch {self.epoch}...")
# Turn off gradient accumulation by switching model to "eval" mode
self.model.eval()
loss_list = []
with torch.no_grad():
for i, batch in enumerate(self.val_loader):
data = batch['image'].float()
target = batch['seg']
prediction = self.model(data.to(self.device))
prediction_softmax = F.softmax(prediction, dim=1)
loss = self.loss_function(prediction_softmax, target[:, 0, :, :].to(self.device))
print(f"Batch {i}. Data shape {data.shape} Loss {loss}")
# We report loss that is accumulated across all of validation set
loss_list.append(loss.item())
self.scheduler.step(np.mean(loss_list))
# log_to_tensorboard(
# self.tensorboard_val_writer,
# np.mean(loss_list),
# data,
# target,
# prediction_softmax,
# prediction,
# (self.epoch+1) * 100)
print(f"Validation complete")
def save_model_parameters(self):
"""
Saves model parameters to a file in results directory
"""
path = os.path.join(self.out_dir, "model.pth")
torch.save(self.model.state_dict(), path)
def load_model_parameters(self, path=''):
"""
Loads model parameters from a supplied path or a
results directory
"""
if not path:
model_path = os.path.join(self.out_dir, "model.pth")
else:
model_path = path
if os.path.exists(model_path):
self.model.load_state_dict(torch.load(model_path))
else:
raise Exception(f"Could not find path {model_path}")
def run_test(self):
"""
This runs test cycle on the test dataset.
Note that process and evaluations are quite different
Here we are computing a lot more metrics and returning
a dictionary that could later be persisted as JSON
"""
print("Testing...")
self.model.eval()
inference_agent = UNetInferenceAgent(model=self.model, device=self.device)
out_dict = {}
out_dict["volume_stats"] = []
dc_list = []
jc_list = []
# for every in test set
for i, x in enumerate(self.test_data):
pred_label = inference_agent.single_volume_inference(x["image"])
# We compute and report Dice and Jaccard similarity coefficients which
# assess how close our volumes are to each other
dc = Dice3d(pred_label, x["seg"])
jc = Jaccard3d(pred_label, x["seg"])
dc_list.append(dc)
jc_list.append(jc)
# STAND-OUT SUGGESTION: By way of exercise, consider also outputting:
# * Sensitivity and specificity (and explain semantic meaning in terms of
# under/over segmenting)
# * Dice-per-slice and render combined slices with lowest and highest DpS
# * Dice per class (anterior/posterior)
out_dict["volume_stats"].append({
"filename": x['filename'],
"dice": dc,
"jaccard": jc
})
print(f"{x['filename']} Dice {dc:.4f}. {100*(i+1)/len(self.test_data):.2f}% complete")
out_dict["overall"] = {
"mean_dice": np.mean(dc_list),
"mean_jaccard": np.mean(jc_list)}
print("\nTesting complete.")
return out_dict
def run(self):
"""
Kicks off train cycle and writes model parameter file at the end
"""
self._time_start = time.time()
print("Experiment started.")
# Iterate over epochs
for self.epoch in range(self.n_epochs):
self.train()
self.validate()
# save model for inferencing
self.save_model_parameters()
self._time_end = time.time()
print(f"Run complete. Total time: {time.strftime('%H:%M:%S', time.gmtime(self._time_end - self._time_start))}")
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
# normalize the data volume
image = (volume.astype(np.single) - np.min(volume)) / (np.max(volume) -
np.min(volume))
# reshape the image volume to the same patch size used for training
img_reshaped = med_reshape(image,
new_shape=(self.patch_size, self.patch_size,
image.shape[2]))
# create a new 3d mask to store predicted results
mask3d = np.zeros(img_reshaped.shape)
# iterate over the image array and predict the all the slices
for slc_idx in range(img_reshaped.shape[2]):
# compute for each slice
slc = torch.from_numpy(img_reshaped[:, :, slc_idx].astype(
np.single)).unsqueeze(0).unsqueeze(0)
# make prediction
pred = self.model(slc.to(self.device))
pred = np.squeeze(pred.cpu().detach())
# store predicted data
mask3d[:, :, slc_idx] = torch.argmax(pred, dim=0)
# return the predicted volume
return mask3d
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
# normalize
image = (volume.astype(np.single) - np.min(volume)) / (np.max(volume) -
np.min(volume))
new_image = med_reshape(image,
new_shape=(self.patch_size, self.patch_size,
image.shape[2]))
mask3d = np.zeros(new_image.shape)
for slc_ix in range(new_image.shape[2]):
tsr_test = torch.from_numpy(new_image[:, :, slc_ix].astype(
np.single)).unsqueeze(0).unsqueeze(0)
#image = torch.from_numpy(self.data[slc[0]]["image"][:,:,slc[1]]).unsqueeze(0)
#tsr_test = torch.from_numpy(slc.astype(np.single)).unsqueeze(0).unsqueeze(0)
pred = self.model(tsr_test.to(self.device))
pred = np.squeeze(pred.cpu().detach())
mask3d[:, :, slc_ix] = torch.argmax(pred, dim=0)
return mask3d
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
raise NotImplementedError
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
volume_shape = volume.shape
data = np.zeros((volume_shape[0], 1, volume_shape[1], volume_shape[2]))
data[:, 0, :, :] = volume / 255.0
data = torch.from_numpy(data).float().to(self.device)
prediction = self.model(data)
pred = prediction.cpu().detach().numpy()
result = pred.argmax(axis=1)
return result
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
PATCH_SIZE = 64
reshaped_vol = med_reshape(volume,
(volume.shape[0], PATCH_SIZE, PATCH_SIZE))
pred_vol = self.single_volume_inference(reshaped_vol)
return pred_vol[:volume.shape[0], :volume.shape[1], :volume.shape[2]]
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
volume = volume.astype(np.single) / np.amax(volume)
pred_vol = np.zeros(volume.shape)
for slc_i in range(volume.shape[0]):
img = volume[slc_i, :, :]
img_tensor = torch.from_numpy(img).unsqueeze(0).unsqueeze(0).to(
self.device, dtype=torch.float)
pred = self.model(img_tensor)
mask = torch.argmax(np.squeeze(pred.cpu().detach()), dim=0)
mask_np = mask.numpy()
pred_vol[slc_i, :, :] = mask_np
return pred_vol
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
raise NotImplementedError
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = np.zeros(volume.shape)
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
with torch.no_grad():
for i, slice_img in enumerate(volume):
# Convert slice to torch tensor
slice_img_tensor = torch.from_numpy(
slice_img).float().unsqueeze(0).unsqueeze(0).to(
self.device)
# Get prediction from the model
predicion = self.model(slice_img_tensor)
predicion = np.squeeze(predicion.cpu().detach())
# Add the result to our slices
slices[i, :, :] = torch.argmax(predicion, dim=0)
return slices
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
volume = med_reshape(volume,
new_shape=(volume.shape[0], self.patch_size,
self.patch_size))
volume = self.single_volume_inference(volume)
return volume
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# Put all slices into a 3D Numpy array.
def inference(img):
img = torch.from_numpy(img.astype(np.single) /
np.max(img)).unsqueeze(0).unsqueeze(0)
pred = self.model(img.to(self.device))
return np.squeeze(pred.cpu().detach())
for idx in range(volume.shape[0]):
pred = inference(volume[idx, :, :])
slices.append(torch.argmax(pred, dim=0).numpy())
slices = np.array([slc for slc in slices])
return slices
class UNetExperiment:
"""
This class implements the basic life cycle for a segmentation task with UNet(https://arxiv.org/abs/1505.04597).
The basic life cycle of a UNetExperiment is:
run():
for epoch in n_epochs:
train()
validate()
test()
"""
def __init__(self, config, split, dataset):
self.n_epochs = config.n_epochs
self.split = split
self._time_start = ""
self._time_end = ""
self.epoch = 0
self.name = config.name
# Create output folders
dirname = f'{time.strftime("%Y-%m-%d_%H%M", time.gmtime())}_{self.name}'
self.out_dir = os.path.join(config.test_results_dir, dirname)
os.makedirs(self.out_dir, exist_ok=True)
# Create data loaders
self.train_loader = DataLoader(SlicesDataset(dataset[split["train"]]),
batch_size=config.batch_size,
shuffle=True,
num_workers=0)
self.val_loader = DataLoader(SlicesDataset(dataset[split["val"]]),
batch_size=config.batch_size,
shuffle=True,
num_workers=0)
# we will access volumes directly for testing
self.test_data = dataset[split["test"]]
# Do we have CUDA available?
if not torch.cuda.is_available():
print(
"WARNING: No CUDA device is found. This may take significantly longer!"
)
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
# Configure our model and other training implements
self.model = UNet(num_classes=3)
self.model.to(self.device)
# Cross entropy loss
self.loss_function = torch.nn.CrossEntropyLoss()
# We are using standard SGD method to optimize our weights
self.optimizer = optim.Adam(self.model.parameters(),
lr=config.learning_rate)
# Scheduler helps us update learning rate automatically
self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
self.optimizer, 'min')
# Set up Tensorboard. By default it saves data into runs folder.
self.tensorboard_train_writer = SummaryWriter(comment="_train")
self.tensorboard_val_writer = SummaryWriter(comment="_val")
def train(self):
"""
This method is executed once per epoch and takes
care of model weight update cycle
"""
print(f"Training epoch {self.epoch}...")
self.model.train()
# Loop over our minibatches
for i, batch in enumerate(self.train_loader):
self.optimizer.zero_grad()
# Put the slices as 4D Torch Tensors of
# shape [BATCH_SIZE, 1, PATCH_SIZE, PATCH_SIZE] into variables data and target.
# Feed data to the model and feed target to the loss function
data = batch['image'].float().to(self.device)
target = batch['seg'].to(self.device)
prediction = self.model(data)
# We are also getting softmax'd version of prediction to output a probability map
prediction_softmax = F.softmax(prediction, dim=1)
loss = self.loss_function(prediction, target[:, 0, :, :].long())
loss.backward()
self.optimizer.step()
if (i % 10) == 0:
# Output to console on every 10th batch
print(
f"\nEpoch: {self.epoch} Train loss: {loss}, {100*(i+1)/len(self.train_loader):.1f}% complete"
)
counter = 100 * self.epoch + 100 * (i / len(self.train_loader))
log_to_tensorboard(self.tensorboard_train_writer, loss, data,
target, prediction_softmax, prediction,
counter)
print(".", end='')
print("\nTraining complete")
def validate(self):
"""
This method runs validation cycle, using same metrics as
Train method. Note that model needs to be switched to eval
mode and no_grad needs to be called so that gradients do not
propagate
"""
print(f"Validating epoch {self.epoch}...")
# Turn off gradient accumulation by switching model to "eval" mode
self.model.eval()
loss_list = []
with torch.no_grad():
for i, batch in enumerate(self.val_loader):
# Compute loss on a validation sample
data = batch["image"].float().to(self.device)
target = batch["seg"].to(self.device)
prediction = self.model(data)
prediction_softmax = F.softmax(prediction, dim=1)
loss = self.loss_function(prediction, target[:,
0, :, :].long())
print(f"Batch {i}. Data shape {data.shape} Loss {loss}")
# We report loss that is accumulated across all of validation set
loss_list.append(loss.item())
self.scheduler.step(np.mean(loss_list))
log_to_tensorboard(self.tensorboard_val_writer, np.mean(loss_list),
data, target, prediction_softmax, prediction,
(self.epoch + 1) * 100)
print(f"Validation complete")
def save_model_parameters(self):
"""
Saves model parameters to a file in results directory
"""
path = os.path.join(self.out_dir, "model.pth")
torch.save(self.model.state_dict(), path)
def load_model_parameters(self, path=''):
"""
Loads model parameters from a supplied path or a
results directory
"""
if not path:
model_path = os.path.join(self.out_dir, "model.pth")
else:
model_path = path
if os.path.exists(model_path):
self.model.load_state_dict(torch.load(model_path))
else:
raise Exception(f"Could not find path {model_path}")
def run_test(self):
"""
This runs test cycle on the test dataset.
Note that process and evaluations are quite different
Here we are computing a lot more metrics and returning
a dictionary that could later be persisted as JSON
"""
print("Testing...")
self.model.eval()
# In this method we will be computing metrics that are relevant to the task of 3D volume
# segmentation. Therefore, unlike train and validation methods, we will do inferences
# on full 3D volumes, much like we will be doing it when we deploy the model in the
# clinical environment.
# Inference Agent is not complete.
inference_agent = UNetInferenceAgent(model=self.model,
device=self.device)
out_dict = {}
out_dict["volume_stats"] = []
dc_list = []
jc_list = []
# for every in test set
for i, x in enumerate(self.test_data):
pred_label = inference_agent.single_volume_inference(x["image"])
# Dice3D and Jaccard3D functions are not implemented.
dc = Dice3d(pred_label, x["seg"])
jc = Jaccard3d(pred_label, x["seg"])
dc_list.append(dc)
jc_list.append(jc)
out_dict["volume_stats"].append({
"filename": x['filename'],
"dice": dc,
"jaccard": jc
})
print(
f"{x['filename']} Dice {dc:.4f}. {100*(i+1)/len(self.test_data):.2f}% complete"
)
out_dict["overall"] = {
"mean_dice": np.mean(dc_list),
"mean_jaccard": np.mean(jc_list)
}
print("\nTesting complete.")
return out_dict
def run(self):
"""
Kicks off train cycle and writes model parameter file at the end
"""
self._time_start = time.time()
print("Experiment started.")
# Iterate over epochs
for self.epoch in range(self.n_epochs):
self.train()
self.validate()
# save model for inferencing
self.save_model_parameters()
self._time_end = time.time()
print(
f"Run complete. Total time: {time.strftime('%H:%M:%S', time.gmtime(self._time_end - self._time_start))}"
)
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
patch_size = 64
volume = med_reshape(volume,
new_shape=(volume.shape[0], patch_size,
patch_size))
return self.single_volume_inference(volume)
raise NotImplementedError
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
# <YOUR CODE HERE>
print('volume shape: ', volume.shape)
# print('volume[0] shape: ', volume[0].shape)
# print('volume[0:1] shape: ', volume[0:1].shape)
# for i in range(volume.shape[0]):
# slices.append(volume[0:1])
# for i in range(volume.shape[0]):
# slices.append(i)
#slices.append(volume[0:1])
# arr = np.zeros(volume.shape, dtype=np.float32)
mask3d = np.zeros(volume.shape)
# trial = volume[11, :, :]
# a = torch.from_numpy(np.asarray(trial).astype(np.single) / 0xff).unsqueeze(0).unsqueeze(0)
# prediction = self.model(a.to(self.device))
# print(prediction.shape)
# mask3d[slc_idx, :, :] = torch.argmax(prediction, dim = 0)
for slc_idx in range(volume.shape[0]):
sliced = volume[slc_idx, :, :]
a = torch.from_numpy(sliced.astype(np.single) /
0xff).unsqueeze(0).unsqueeze(0)
prediction = self.model(a.to(self.device))
# print(prediction.shape)
mask3d[slc_idx, :, :] = torch.argmax(np.squeeze(prediction.cpu()),
dim=0)
# slices.append(prediction.argmax(1).cpu().numpy().squeeze())
# slices.append(torch.argmax(np.squeeze(prediction.cpu()), dim = 0))
# prediction_softmax = F.softmax(prediction, dim=1)
# print(prediction_softmax)
# print('len(slices): ', len(slices))
# print('slices[0]: ', slices[0])
# a = torch.from_numpy(slices).type(torch.FloatTensor).unsqueeze(0)
# print(a.shape)
# for idx, label in enumerate(slices):
# print(f'idx: {idx}')
# print(f'label: {label}')
# print(label)
# if label is not np.nan:
#label = label.split(" ")
#mask = np.zeros(volume.shape[1] * volume.shape[2], dtype=np.uint8)
#posit = map(int, label[0::2])
#leng = map(int, label[1::2])
#for p, l in zip(posit, leng):
# mask[p:(p+l)] = 1
# arr[:, :, idx] = mask.reshape(volume.shape[1], volume.shape[1], order='F')
# slices = np.asarray(slices)
# slices.reshape((-1, slices.shape[0], slices.shape[1]))
return mask3d #slices #np.array(slices)
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
raise NotImplementedError
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
# <YOUR CODE HERE>
#Initialize the mask3d array
mask3d = np.zeros(volume.shape)
#cycle for each slice x
for ix in range(volume.shape[0]):
#get the image slice and normalize it
slice = volume[ix, :, :].astype(np.single) / np.max(
volume[ix, :, :])
#inference on the slice
mytensor = torch.from_numpy(slice).unsqueeze(0).unsqueeze(0)
pred = self.model(mytensor.to(self.device))
#apply the torch.argmax
mask3d[ix, :, :] = torch.argmax(np.squeeze(pred.cpu().detach()),
dim=0)
return mask3d
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
slices = np.zeros(volume.shape)
for i in range(volume.shape[0]):
s = volume[i, :, :]
s = s.astype(np.single)
s = s / 255.0
s = torch.from_numpy(s).unsqueeze(0).unsqueeze(0)
pred = self.model(s.to(self.device))
pred = np.squeeze(pred.cpu().detach())
slices[i, :, :] = torch.argmax(pred, dim=0)
return slices
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
patch_size = 64
volume = med_reshape(volume,
new_shape=(volume.shape[0], patch_size,
patch_size))
return single_volume_inference(volume)
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
print(parameter_file_path)
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
x, y, z = volume.shape
padded_shape = (x, self.patch_size, self.patch_size)
padded_volume = np.zeros(padded_shape)
padded_volume[:x, :y, :z] = volume
return self.single_volume_inference(padded_volume)
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
# <YOUR CODE HERE>
slices = np.zeros(volume.shape)
for slc_ix in range(volume.shape[0]):
img = volume[slc_ix, :, :]
img = img.astype(np.single)
img_norm = img / 255.0
img_t = torch.from_numpy(img_norm).unsqueeze(0).unsqueeze(0)
pred = self.model(img_t.to(self.device))
pred = np.squeeze(pred.cpu().detach())
slices[slc_ix, :, :] = torch.argmax(pred, dim=0)
return slices
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self, parameter_file_path='', model=None, device="cpu", patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
# raise NotImplementedError
volume = med_reshape(volume, new_shape=(volume.shape[0], self.patch_size, self.patch_size))
return self.single_volume_inference(volume)
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
x = volume.shape[0]
mask = np.zeros(volume.shape)
for i in range(x):
_slice = volume[i, :, :]
if _slice.min() != _slice.max():
_slice = (_slice - _slice.min()) / (_slice.max() - _slice.min()) # change the range of _slice to [0, 1]
_slice_torch = torch.from_numpy(_slice).type(torch.float).unsqueeze(0).unsqueeze(0).to(self.device)
pred = self.model(_slice_torch)
pred = np.squeeze(pred.cpu().detach())
mask[i, :, :] = torch.argmax(pred, dim=0)
return mask
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
patch_size = 64
volume = med_reshape(volume,
new_shape=(volume.shape[0], patch_size,
patch_size))
return volume
#raise NotImplementedError
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
#slices = []
slices = np.zeros(volume.shape)
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
# <YOUR CODE HERE>
#tsr_test = torch.from_numpy(img.astype(np.single)/np.max(img)).unsqueeze(0).unsqueeze(0)
for slice_idx in range(volume.shape[0]):
# normalize the image
slice0 = volume[slice_idx, :, :]
slice0 = slice0.astype(np.single)
slice0_norm = slice0 / 255.0
# convert to pytorch tensor
slice0_t = torch.from_numpy(slice0_norm).unsqueeze(0).unsqueeze(0)
# predict the label
pred = self.model(slice0_t.to(self.device))
# normal image slice
pred = np.squeeze(pred.cpu().detach())
slices[slice_idx, :, :] = torch.argmax(pred, dim=0)
return slices
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
raise NotImplementedError
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
vol_tensor = torch.from_numpy(volume).type(torch.FloatTensor).to(
self.device)
vol_tensor = vol_tensor.unsqueeze(1)
# v now has shape (num_sagittal_slices,1,patch_size,patch_size)
# 0th index is being used as batch index. so the entire volume is being treated as
# one big batch of slices
with torch.no_grad():
prediction = F.softmax(self.model(vol_tensor), dim=1).cpu().numpy()
# prediction has shape (num_sagittal_slices,3,patch_size,patch_size)
return prediction.argmax(axis=1)
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
patch_size = 64
volume = (volume - volume.min()) / (volume.max() - volume.min())
volume = med_reshape(volume,
new_shape=(volume.shape[0], patch_size,
patch_size))
masks = np.zeros(volume.shape)
for slice_idx in range(masks.shape[0]):
# normalize the image
slice_0 = volume[slice_idx, :, :]
#slice0_norm = (slice0-slice0.min())/(slice0.max()-slice0.min())
data = torch.from_numpy(slice_0).unsqueeze(0).unsqueeze(
0).float().to(self.device)
pred = self.model(data)
pred = np.squeeze(pred.cpu().detach())
pred = pred.argmax(axis=0)
masks[slice_idx, :, :] = pred
return masks
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
# <YOUR CODE HERE>
masks = np.zeros(volume.shape)
for slice_idx in range(masks.shape[0]):
slice_0 = volume[slice_idx, :, :]
data = torch.from_numpy(slice_0).unsqueeze(0).unsqueeze(
0).float().to(self.device)
pred = self.model(data)
pred = np.squeeze(pred.cpu().detach())
pred = pred.argmax(axis=0)
masks[slice_idx, :, :] = pred
return masks
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
volume = med_reshape(volume,
new_shape=(self.patch_size, self.patch_size,
self.patch_size))
return self.single_volume_inference(volume)
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# Create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array.
for slice in volume:
if (np.count_nonzero(slice) == 0):
slices.append(np.zeros((self.patch_size, self.patch_size)))
continue
slice_input = torch.tensor(slice[None, None, :, :],
dtype=torch.float).to(self.device)
prediction = np.squeeze(self.model(slice_input).cpu().detach())
mask = torch.argmax(prediction, dim=0).numpy()
# print(f"SVI slice count nonzero: {np.count_nonzero(mask)}" )
slices.append(mask)
return np.array(slices)
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
reshaped = med_reshape(
volume, (volume.shape[0], self.patch_size, self.patch_size))
return self.single_volume_inference(reshaped)
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
for slc in volume:
indata = torch.tensor(slc).unsqueeze(0).unsqueeze(0).type(
torch.FloatTensor).to(self.device)
rawp = self.model(indata)
amax = torch.argmax(rawp, dim=1)
pred = amax.squeeze().type(torch.LongTensor).cpu()
slices.append(pred.numpy())
return np.array(slices)
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
#print('testing init...')
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
print('testing init end...')
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
#print('single_volume_inference_unpadded')
raise NotImplementedError
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
print('single_volume_inference init...')
self.model.eval()
print('volume', volume.shape)
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = np.zeros(volume.shape)
for i in range(volume.shape[0]):
tsr_test = torch.from_numpy(
volume[i, :, :].astype(np.single) /
np.max(volume[i, :, :])).unsqueeze(0).unsqueeze(0)
#print('tsr_test',tsr_test.shape)
pred = self.model(tsr_test)
slices[i, :, :] = torch.argmax(np.squeeze(pred.cpu().detach()),
dim=0)
#print('slices.slices',(slices[0]))
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
# <YOUR CODE HERE>
return slices
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self, parameter_file_path='', model=None, device="cpu", patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
raise NotImplementedError
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
# <YOUR CODE HERE>
slices.append(volume[0:1])
arr = np.zeros(volume.shape, dtype=np.float32)
for idx, label in enumerate(slices):
if label is not np.nan:
label = label.split(" ")
mask = np.zeros(volume.shape[1] * volume.shape[2], dtype=np.uint8)
posit = map(int, label[0::2])
leng = map(int, label[1::2])
for p, l in zip(posit, leng):
mask[p:(p+l)] = 1
arr[:, :, idx] = mask.reshape(volume.shape[1], volume.shape[1], order='F')
slices = slices.asarray()
slices.reshape((-1, slices.shape[0], slices.shape[1]))
return slices
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
volume = med_reshape(volume, (volume.shape[0], patch_size, patch_size))
raise self.single_volume_inference(volume)
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
# <YOUR CODE HERE>
for i in range(0, volume.shape[0]):
slc_tensor = torch.from_numpy(volume[i, :, :].astype(
np.single)).unsqueeze(0).unsqueeze(0)
pred = self.model(slc_tensor.to(self.device))
mask = torch.argmax(np.squeeze(pred.cpu().detach()), dim=0)
slices.append(mask)
return np.dstack(slices).transpose(2, 0, 1)
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
x, y, z = volume.shape
new_shape = (x, self.patch_size, self.patch_size)
new_volume = med_reshape(volume, new_shape)
slices = np.zeros(new_volume.shape)
for slice_idx in range(new_volume.shape[0]):
slc = new_volume[slice_idx, :, :]
slice = torch.from_numpy(slc.astype(np.single) /
np.max(slc)).unsqueeze(0).unsqueeze(0).to(
self.device)
prediction = self.model(slice)
prediction = np.squeeze(prediction).cpu().detach()
slices[slice_idx, :, :] = torch.argmax(prediction, dim=0)
return slices
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
# slices = []
# Create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array.
slices = np.zeros(volume.shape)
for slice_index in range(volume.shape[0]):
slc = volume[slice_index, :, :]
slc = slc.astype(np.single) / np.max(slc)
slice = torch.from_numpy(slc).unsqueeze(0).unsqueeze(0).to(
self.device)
prediction = self.model(slice)
prediction = np.squeeze(prediction.cpu().detach())
slices[slice_index, :, :] = torch.argmax(prediction, dim=0)
return slices
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self, parameter_file_path='', model=None, device="cpu", patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(torch.load(parameter_file_path, map_location=self.device))
self.model.to(self.device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
raise NotImplementedError
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# TASK: Write code that will create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. You can verify if your method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
img_ = med_reshape(volume, new_shape=(volume.shape[0], self.patch_size, self.patch_size))
for i in range(img_.shape[0]):
image = torch.from_numpy(img_[i] / np.max(img_[i])).unsqueeze(0).unsqueeze(0)
output = self.model(image.to(self.device, dtype=torch.float))
output = np.squeeze(output.cpu().detach())
slices.append(torch.argmax(output, dim=0).numpy())
return np.asarray(slices)
class UNetInferenceAgent:
"""
Stores model and parameters and some methods to handle inferencing
"""
def __init__(self,
parameter_file_path='',
model=None,
device="cpu",
patch_size=64):
self.model = model
self.patch_size = patch_size
self.device = device
if model is None:
self.model = UNet(num_classes=3)
if parameter_file_path:
self.model.load_state_dict(
torch.load(parameter_file_path, map_location=self.device))
self.model.to(device)
def single_volume_inference_unpadded(self, volume):
"""
Runs inference on a single volume of arbitrary patch size,
padding it to the conformant size first
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
# Set model to eval no gradients
self.model.eval()
# Initialise slices with zeros
slices = np.zeros(volume.shape)
# Reshape volume to conform to required model patch size
volume = med_reshape(volume,
new_shape=(volume.shape[0], self.patch_size,
self.patch_size))
# For each x slice in the volume
for x_index in range(volume.shape[0]):
# Get the x slice
x_slice = volume[x_index, :, :].astype(np.single)
# Convert to tensor
tensor_x_slice = torch.from_numpy(x_slice).unsqueeze(0).unsqueeze(
0).to(self.device)
# Pass slice through model to get predictions
predictions = self.model(tensor_x_slice)
# Resize predictions
pred_resized = np.squeeze(predictions.cpu().detach())
# Append predictions
slices[x_index, :, :] = torch.argmax(pred_resized, dim=0)
# Return volume of predictions
return slices
def single_volume_inference(self, volume):
"""
Runs inference on a single volume of conformant patch size
Arguments:
volume {Numpy array} -- 3D array representing the volume
Returns:
3D NumPy array with prediction mask
"""
self.model.eval()
# Assuming volume is a numpy array of shape [X,Y,Z] and we need to slice X axis
slices = []
# Create mask for each slice across the X (0th) dimension. After
# that, put all slices into a 3D Numpy array. We can verify if the method is
# correct by running it on one of the volumes in your training set and comparing
# with the label in 3D Slicer.
# Set model to eval no gradients
self.model.eval()
# Initialise slices with zeros
slices = np.zeros(volume.shape)
# For each x slice in the volume
for x_index in range(volume.shape[0]):
# Get the x slice
x_slice = volume[x_index, :, :].astype(np.single)
# Convert to tensor
tensor_x_slice = torch.from_numpy(x_slice).unsqueeze(0).unsqueeze(
0).to(self.device)
# Pass slice through model to get predictions
predictions = self.model(tensor_x_slice)
# Resize predictions
pred_resized = np.squeeze(predictions.cpu().detach())
# Append predictions
slices[x_index, :, :] = torch.argmax(pred_resized, dim=0)
# Return volume of predictions
return slices
