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