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 self.patch_size = config.patch_size # 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 setup(self): pkl_dir = self.config.split_dir with open(os.path.join(pkl_dir, "splits.pkl"), 'rb') as f: splits = pickle.load(f) tr_keys = splits[self.config.fold]['train'] val_keys = splits[self.config.fold]['val'] test_keys = splits[self.config.fold]['test'] self.device = torch.device( self.config.device if torch.cuda.is_available() else "cpu") self.train_data_loader = NumpyDataSet( self.config.data_dir, target_size=self.config.patch_size, batch_size=self.config.batch_size, keys=tr_keys) self.val_data_loader = NumpyDataSet(self.config.data_dir, target_size=self.config.patch_size, batch_size=self.config.batch_size, keys=val_keys, mode="val", do_reshuffle=False) self.test_data_loader = NumpyDataSet( self.config.data_test_dir, target_size=self.config.patch_size, batch_size=self.config.batch_size, keys=test_keys, mode="test", do_reshuffle=False) self.model = UNet(num_classes=self.config.num_classes, in_channels=self.config.in_channels) self.model.to(self.device) # We use a combination of DICE-loss and CE-Loss in this example. # This proved good in the medical segmentation decathlon. self.dice_loss = SoftDiceLoss( batch_dice=True) # Softmax für DICE Loss! self.ce_loss = torch.nn.CrossEntropyLoss( ) # Kein Softmax für CE Loss -> ist in torch schon mit drin! self.optimizer = optim.Adam(self.model.parameters(), lr=self.config.learning_rate) self.scheduler = ReduceLROnPlateau(self.optimizer, 'min') # If directory for checkpoint is provided, we load it. if self.config.do_load_checkpoint: if self.config.checkpoint_dir == '': print( 'checkpoint_dir is empty, please provide directory to load checkpoint.' ) else: self.load_checkpoint(name=self.config.checkpoint_dir, save_types=("model")) self.save_checkpoint(name="checkpoint_start") self.elog.print('Experiment set up.')
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 # 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"]] # USE CUDA 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. self.model = UNet(num_classes=3) self.model.to(self.device) self.loss_function = torch.nn.CrossEntropyLoss() self.optimizer = optim.Adam(self.model.parameters(), lr=config.learning_rate) # Update learning rate automatically self.scheduler = optim.lr_scheduler.ReduceLROnPlateau( self.optimizer, 'min') # Set up Tensorboard. By default it saves data into runs folder. # To check the training results, run the tensorboard. # Like how we use the jupyter notebook remotely. self.tensorboard_train_writer = SummaryWriter(comment="_train") self.tensorboard_val_writer = SummaryWriter(comment="_val")
def segment_single_image(self, data): self.model = UNet(num_classes=self.config.num_classes, in_channels=self.config.in_channels) self.device = torch.device( self.config.device if torch.cuda.is_available() else "cpu") # a model must be present and loaded in here if self.config.model_dir == '': print( 'model_dir is empty, please provide directory to load checkpoint.' ) else: self.load_checkpoint(name=self.config.model_dir, save_types=("model")) self.elog.print("=====SEGMENT_SINGLE_IMAGE=====") self.model.eval() self.model.to(self.device) # Desired shape = [b, c, w, h] # split into even chunks (lets use size) with torch.no_grad(): ###### # When working entirely on CPU and in memory, the following lines replace the split/concat method # mr_data = data.float().to(self.device) # pred = self.model(mr_data) # pred_argmax = torch.argmax(pred.data.cpu(), dim=1, keepdim=True) ###### ###### # for CUDA (also works on CPU) split into batches blocksize = self.config.batch_size # number_of_elements = round(data.shape[0]/blocksize+0.5) # make blocks large enough to not lose any slices chunks = [ data[i:i + blocksize, ::, ::, ::] for i in range(0, data.shape[0], blocksize) ] pred_list = [] for data_batch in chunks: mr_data = data_batch.float().to(self.device) pred_dict = self.model(mr_data) pred_list.append(pred_dict.cpu()) pred = torch.Tensor(np.concatenate(pred_list)) pred_argmax = torch.argmax(pred, dim=1, keepdim=True) # detach result and put it back to cpu so that we can work with, create a numpy array result = pred_argmax.short().detach().cpu().numpy() return result
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 __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 __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") # 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')
def setup(self): pkl_dir = self.config.split_dir with open(os.path.join(pkl_dir, "splits.pkl"), 'rb') as f: splits = pickle.load(f) tr_keys = splits[self.config.fold]['train'] val_keys = splits[self.config.fold]['val'] keys = tr_keys + val_keys test_keys = splits[self.config.fold]['test'] self.device = torch.device(self.config.device if torch.cuda.is_available() else 'cpu') # self.model = UNet(num_classes=self.config.num_classes, num_downs=3) self.model.to(self.device) self.data_loader = NumpyDataSet(self.config.data_dir, target_size=256, batch_size=self.config.batch_size, keys=keys, mode='test', do_reshuffle=False) self.data_16_loader = NumpyDataSet(self.config.scaled_image_32_dir, target_size=32, batch_size=self.config.batch_size, keys=keys, mode='test', do_reshuffle=False) # We use a combination of DICE-loss and CE-Loss in this example. # This proved good in the medical segmentation decathlon. self.dice_loss = SoftDiceLoss(batch_dice=True) # Softmax für DICE Loss! # weight = torch.tensor([1, 30, 30]).float().to(self.device) self.ce_loss = torch.nn.CrossEntropyLoss() # Kein Softmax für CE Loss -> ist in torch schon mit drin! # self.dice_pytorch = dice_pytorch(self.config.num_classes) self.optimizer = optim.Adam(self.model.parameters(), lr=self.config.learning_rate) # self.optimizer = optim.SGD(self.model.parameters(), lr=self.config.learning_rate) self.scheduler = ReduceLROnPlateau(self.optimizer, 'min') # If directory for checkpoint is provided, we load it. if self.config.do_load_checkpoint: if self.config.checkpoint_dir == '': print('checkpoint_dir is empty, please provide directory to load checkpoint.') else: self.load_checkpoint(name=self.config.checkpoint_dir, save_types=("model"))
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 #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 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 = 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 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): 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 """ 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(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 """ 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 UNetExperiment(PytorchExperiment): """ The UnetExperiment is inherited from the PytorchExperiment. It implements the basic life cycle for a segmentation task with UNet(https://arxiv.org/abs/1505.04597). It is optimized to work with the provided NumpyDataLoader. The basic life cycle of a UnetExperiment is the same s PytorchExperiment: setup() (--> Automatically restore values if a previous checkpoint is given) prepare() for epoch in n_epochs: train() validate() (--> save current checkpoint) end() """ def setup(self): pkl_dir = self.config.split_dir with open(os.path.join(pkl_dir, "splits.pkl"), 'rb') as f: splits = pickle.load(f) tr_keys = splits[self.config.fold]['train'] val_keys = splits[self.config.fold]['val'] test_keys = splits[self.config.fold]['test'] self.device = torch.device( self.config.device if torch.cuda.is_available() else "cpu") self.train_data_loader = NumpyDataSet( self.config.data_dir, target_size=self.config.patch_size, batch_size=self.config.batch_size, keys=tr_keys) self.val_data_loader = NumpyDataSet(self.config.data_dir, target_size=self.config.patch_size, batch_size=self.config.batch_size, keys=val_keys, mode="val", do_reshuffle=False) self.test_data_loader = NumpyDataSet( self.config.data_test_dir, target_size=self.config.patch_size, batch_size=self.config.batch_size, keys=test_keys, mode="test", do_reshuffle=False) self.model = UNet(num_classes=self.config.num_classes, in_channels=self.config.in_channels) self.model.to(self.device) # We use a combination of DICE-loss and CE-Loss in this example. # This proved good in the medical segmentation decathlon. self.dice_loss = SoftDiceLoss( batch_dice=True) # Softmax for DICE Loss! self.ce_loss = torch.nn.CrossEntropyLoss( ) # No softmax for CE Loss -> is implemented in torch! self.optimizer = optim.Adam(self.model.parameters(), lr=self.config.learning_rate) self.scheduler = ReduceLROnPlateau(self.optimizer, 'min') # If directory for checkpoint is provided, we load it. if self.config.do_load_checkpoint: if self.config.checkpoint_dir == '': print( 'checkpoint_dir is empty, please provide directory to load checkpoint.' ) else: self.load_checkpoint(name=self.config.checkpoint_dir, save_types=("model")) self.save_checkpoint(name="checkpoint_start") self.elog.print('Experiment set up.') def train(self, epoch): self.elog.print('=====TRAIN=====') self.model.train() data = None batch_counter = 0 for data_batch in self.train_data_loader: self.optimizer.zero_grad() # Shape of data_batch = [1, b, c, w, h] # Desired shape = [b, c, w, h] # Move data and target to the GPU data = data_batch['data'][0].float().to(self.device) target = data_batch['seg'][0].long().to(self.device) pred = self.model(data) pred_softmax = F.softmax( pred, dim=1 ) # We calculate a softmax, because our SoftDiceLoss expects that as an input. The CE-Loss does the softmax internally. #loss = self.dice_loss(pred_softmax, target.squeeze()) + self.ce_loss(pred, target.squeeze()) loss = self.ce_loss(pred, target.squeeze()) loss.backward() self.optimizer.step() # Some logging and plotting if (batch_counter % self.config.plot_freq) == 0: self.elog.print('Epoch: {0} Loss: {1:.4f}'.format( self._epoch_idx, loss)) self.add_result( value=loss.item(), name='Train_Loss', tag='Loss', counter=epoch + (batch_counter / self.train_data_loader.data_loader.num_batches)) self.clog.show_image_grid(data.float().cpu(), name="data", normalize=True, scale_each=True, n_iter=epoch) self.clog.show_image_grid(target.float().cpu(), name="mask", title="Mask", n_iter=epoch) self.clog.show_image_grid(torch.argmax(pred.cpu(), dim=1, keepdim=True), name="unt_argmax", title="Unet", n_iter=epoch) self.clog.show_image_grid(pred.cpu()[:, 1:2, ], name="unt", normalize=True, scale_each=True, n_iter=epoch) batch_counter += 1 assert data is not None, 'data is None. Please check if your dataloader works properly' def validate(self, epoch): self.elog.print('VALIDATE') self.model.eval() data = None loss_list = [] with torch.no_grad(): for data_batch in self.val_data_loader: data = data_batch['data'][0].float().to(self.device) target = data_batch['seg'][0].long().to(self.device) pred = self.model(data) pred_softmax = F.softmax( pred, dim=1 ) # We calculate a softmax, because our SoftDiceLoss expects that as an input. The CE-Loss does the softmax internally. #loss = self.dice_loss(pred_softmax, target.squeeze()) + self.ce_loss(pred, target.squeeze()) loss = self.ce_loss(pred, target.squeeze()) loss_list.append(loss.item()) assert data is not None, 'data is None. Please check if your dataloader works properly' self.scheduler.step(np.mean(loss_list)) self.elog.print('Epoch: %d Loss: %.4f' % (self._epoch_idx, np.mean(loss_list))) self.add_result(value=np.mean(loss_list), name='Val_Loss', tag='Loss', counter=epoch + 1) self.clog.show_image_grid(data.float().cpu(), name="data_val", normalize=True, scale_each=True, n_iter=epoch) self.clog.show_image_grid(target.float().cpu(), name="mask_val", title="Mask", n_iter=epoch) self.clog.show_image_grid(torch.argmax(pred.data.cpu(), dim=1, keepdim=True), name="unt_argmax_val", title="Unet", n_iter=epoch) self.clog.show_image_grid(pred.data.cpu()[:, 1:2, ], name="unt_val", normalize=True, scale_each=True, n_iter=epoch) def test(self): from evaluation.evaluator import aggregate_scores, Evaluator from collections import defaultdict self.elog.print('=====TEST=====') self.model.eval() pred_dict = defaultdict(list) gt_dict = defaultdict(list) batch_counter = 0 with torch.no_grad(): for data_batch in self.test_data_loader: print('testing...', batch_counter) batch_counter += 1 # Get data_batches mr_data = data_batch['data'][0].float().to(self.device) mr_target = data_batch['seg'][0].float().to(self.device) pred = self.model(mr_data) pred_argmax = torch.argmax(pred.data.cpu(), dim=1, keepdim=True) fnames = data_batch['fnames'] for i, fname in enumerate(fnames): pred_dict[fname[0]].append( pred_argmax[i].detach().cpu().numpy()) gt_dict[fname[0]].append( mr_target[i].detach().cpu().numpy()) test_ref_list = [] for key in pred_dict.keys(): test_ref_list.append( (np.stack(pred_dict[key]), np.stack(gt_dict[key]))) scores = aggregate_scores(test_ref_list, evaluator=Evaluator, json_author=self.config.author, json_task=self.config.name, json_name=self.config.name, json_output_file=self.elog.work_dir + "/{}_".format(self.config.author) + self.config.name + '.json') print("Scores:\n", scores) def segment_single_image(self, data): self.model = UNet(num_classes=self.config.num_classes, in_channels=self.config.in_channels) self.device = torch.device( self.config.device if torch.cuda.is_available() else "cpu") # a model must be present and loaded in here if self.config.model_dir == '': print( 'model_dir is empty, please provide directory to load checkpoint.' ) else: self.load_checkpoint(name=self.config.model_dir, save_types=("model")) self.elog.print("=====SEGMENT_SINGLE_IMAGE=====") self.model.eval() self.model.to(self.device) # Desired shape = [b, c, w, h] # split into even chunks (lets use size) with torch.no_grad(): ###### # When working entirely on CPU and in memory, the following lines replace the split/concat method # mr_data = data.float().to(self.device) # pred = self.model(mr_data) # pred_argmax = torch.argmax(pred.data.cpu(), dim=1, keepdim=True) ###### ###### # for CUDA (also works on CPU) split into batches blocksize = self.config.batch_size # number_of_elements = round(data.shape[0]/blocksize+0.5) # make blocks large enough to not lose any slices chunks = [ data[i:i + blocksize, ::, ::, ::] for i in range(0, data.shape[0], blocksize) ] pred_list = [] for data_batch in chunks: mr_data = data_batch.float().to(self.device) pred_dict = self.model(mr_data) pred_list.append(pred_dict.cpu()) pred = torch.Tensor(np.concatenate(pred_list)) pred_argmax = torch.argmax(pred, dim=1, keepdim=True) # detach result and put it back to cpu so that we can work with, create a numpy array result = pred_argmax.short().detach().cpu().numpy() 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 """ # 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
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 """ # 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 """ 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 """ 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 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(PytorchExperiment): """ The UnetExperiment is inherited from the PytorchExperiment. It implements the basic life cycle for a segmentation task with UNet(https://arxiv.org/abs/1505.04597). It is optimized to work with the provided NumpyDataLoader. The basic life cycle of a UnetExperiment is the same s PytorchExperiment: setup() (--> Automatically restore values if a previous checkpoint is given) prepare() for epoch in n_epochs: train() validate() (--> save current checkpoint) end() """ def setup(self): pkl_dir = self.config.split_dir with open(os.path.join(pkl_dir, "splits.pkl"), 'rb') as f: splits = pickle.load(f) tr_keys = splits[self.config.fold]['train'] val_keys = splits[self.config.fold]['val'] test_keys = splits[self.config.fold]['test'] self.device = torch.device( self.config.device if torch.cuda.is_available() else "cpu") self.train_data_loader = NumpyDataSet( self.config.data_dir, target_size=self.config.patch_size, batch_size=self.config.batch_size, keys=tr_keys) self.val_data_loader = NumpyDataSet(self.config.data_dir, target_size=self.config.patch_size, batch_size=self.config.batch_size, keys=val_keys, mode="val", do_reshuffle=False) self.test_data_loader = NumpyDataSet( self.config.data_test_dir, target_size=self.config.patch_size, batch_size=self.config.batch_size, keys=test_keys, mode="test", do_reshuffle=False) self.model = UNet(num_classes=self.config.num_classes, in_channels=self.config.in_channels) self.model.to(self.device) # We use a combination of DICE-loss and CE-Loss in this example. # This proved good in the medical segmentation decathlon. self.dice_loss = SoftDiceLoss( batch_dice=True) # Softmax für DICE Loss! self.ce_loss = torch.nn.CrossEntropyLoss( ) # Kein Softmax für CE Loss -> ist in torch schon mit drin! self.optimizer = optim.Adam(self.model.parameters(), lr=self.config.learning_rate) self.scheduler = ReduceLROnPlateau(self.optimizer, 'min') # If directory for checkpoint is provided, we load it. if self.config.do_load_checkpoint: if self.config.checkpoint_dir == '': print( 'checkpoint_dir is empty, please provide directory to load checkpoint.' ) else: self.load_checkpoint(name=self.config.checkpoint_dir, save_types=("model")) self.save_checkpoint(name="checkpoint_start") self.elog.print('Experiment set up.') def train(self, epoch): self.elog.print('=====TRAIN=====') self.model.train() data = None batch_counter = 0 for data_batch in self.train_data_loader: self.optimizer.zero_grad() # Shape of data_batch = [1, b, c, w, h] # Desired shape = [b, c, w, h] # Move data and target to the GPU data = data_batch['data'][0].float().to(self.device) target = data_batch['seg'][0].long().to(self.device) pred = self.model(data) pred_softmax = F.softmax( pred, dim=1 ) # We calculate a softmax, because our SoftDiceLoss expects that as an input. The CE-Loss does the softmax internally. loss = self.dice_loss(pred_softmax, target.squeeze()) + self.ce_loss( pred, target.squeeze()) # loss = self.ce_loss(pred, target.squeeze()) loss.backward() self.optimizer.step() # Some logging and plotting if (batch_counter % self.config.plot_freq) == 0: self.elog.print('Epoch: %d Loss: %.4f' % (self._epoch_idx, loss)) self.add_result( value=loss.item(), name='Train_Loss', tag='Loss', counter=epoch + (batch_counter / self.train_data_loader.data_loader.num_batches)) self.clog.show_image_grid(data.float(), name="data", normalize=True, scale_each=True, n_iter=epoch) self.clog.show_image_grid(target.float(), name="mask", title="Mask", n_iter=epoch) self.clog.show_image_grid(torch.argmax(pred.cpu(), dim=1, keepdim=True), name="unt_argmax", title="Unet", n_iter=epoch) self.clog.show_image_grid(pred.cpu()[:, 1:2, ], name="unt", normalize=True, scale_each=True, n_iter=epoch) batch_counter += 1 assert data is not None, 'data is None. Please check if your dataloader works properly' def validate(self, epoch): self.elog.print('VALIDATE') self.model.eval() data = None loss_list = [] with torch.no_grad(): for data_batch in self.val_data_loader: data = data_batch['data'][0].float().to(self.device) target = data_batch['seg'][0].long().to(self.device) pred = self.model(data) pred_softmax = F.softmax( pred ) # We calculate a softmax, because our SoftDiceLoss expects that as an input. The CE-Loss does the softmax internally. loss = self.dice_loss(pred_softmax, target.squeeze()) + self.ce_loss( pred, target.squeeze()) loss_list.append(loss.item()) assert data is not None, 'data is None. Please check if your dataloader works properly' self.scheduler.step(np.mean(loss_list)) self.elog.print('Epoch: %d Loss: %.4f' % (self._epoch_idx, np.mean(loss_list))) self.add_result(value=np.mean(loss_list), name='Val_Loss', tag='Loss', counter=epoch + 1) self.clog.show_image_grid(data.float(), name="data_val", normalize=True, scale_each=True, n_iter=epoch) self.clog.show_image_grid(target.float(), name="mask_val", title="Mask", n_iter=epoch) self.clog.show_image_grid(torch.argmax(pred.data.cpu(), dim=1, keepdim=True), name="unt_argmax_val", title="Unet", n_iter=epoch) self.clog.show_image_grid(pred.data.cpu()[:, 1:2, ], name="unt_val", normalize=True, scale_each=True, n_iter=epoch) def test(self): # TODO print('TODO: Implement your test() method here')
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 """ 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