def predict(im, net, device, outScaling): ''' Process an image using our network. Parameters ---------- im: numpy array 2D image we want to process net: a pytorch model the network we want to use device: The device your network lives on, e.g. your GPU outScaling: float We found that scaling the output by a factor (default=10) helps to speedup training. Returns ---------- decoPred: numpy array Image containing the prediction. ''' stdTorch=torch.Tensor(np.array(net.std)).to(device) meanTorch=torch.Tensor(np.array(net.mean)).to(device) #im=(im-net.mean)/net.std inputs_raw= torch.zeros(1,1,im.shape[0],im.shape[1]) inputs_raw[0,:,:,:]=imgToTensor(im); # copy to GPU inputs_raw = inputs_raw.to(device) # normalize inputs = (inputs_raw-meanTorch)/stdTorch output=net(inputs) samples = (output).permute(1, 0, 2, 3)*outScaling #We found that this factor can speed up training samples = samples * stdTorch + meanTorch # denormalize decoPred = torch.mean(samples,dim=0,keepdim=True)[0,...] # Sum up over all samples decoPred = decoPred.cpu().detach().numpy() decoPred.shape = (output.shape[2],output.shape[3]) return decoPred, None
def trainingPred(my_train_data, net, dataCounter, size, bs, numPix, device, augment=True, supervised=True): ''' This function will assemble a minibatch and process it using the a network. Parameters ---------- my_train_data: numpy array Your training dataset, should be a stack of 2D images, i.e. a 3D numpy array net: a pytorch model the network we want to use dataCounter: int The index of the next image to be used. size: int Witdth and height of the training patches that are to be used. bs: int The batch patch_size. numPix: int The number of pixels that is to be manipulated/masked N2V style. augment: bool should the patches be randomy flipped and rotated? Returns ---------- samples: pytorch tensor The output of the network labels: pytorch tensor This is the tensor that was is used a target. It holds the raw unmanipulated patches. masks: pytorch tensor A tensor marking which pixels have been manipulated (value 1) and which not (value 0). In N2V or PN2V only these pixels should be used to calculate gradients. dataCounter: int The updated counter parameter, it is increased by one. When the counter reaches the end of the dataset, it is reset to zero and the dataset is shuffled. ''' # Init Variables inputs = torch.zeros(bs, 1, size, size) labels = torch.zeros(bs, size, size) masks = torch.zeros(bs, size, size) # Assemble mini batch for j in range(bs): im, l, m, dataCounter = randomCropFRI(my_train_data, size, numPix, counter=dataCounter, augment=augment, supervised=supervised) inputs[j, :, :, :] = utils.imgToTensor(im) labels[j, :, :] = utils.imgToTensor(l) masks[j, :, :] = utils.imgToTensor(m) # Move to GPU inputs_raw, labels, masks = inputs.to(device), labels.to(device), masks.to( device) # Move normalization parameter to GPU stdTorch = torch.Tensor(np.array(net.std)).to(device) meanTorch = torch.Tensor(np.array(net.mean)).to(device) # Forward step outputs = net( (inputs_raw - meanTorch) / stdTorch) * 10.0 # We found that this factor can speed up training samples = (outputs).permute(1, 0, 2, 3) # Denormalize samples = utils.denormalize(samples, meanTorch, stdTorch) return samples, labels, masks, dataCounter
def predict(im, net, noiseModel, device, outScaling): ''' Process an image using our network. Parameters ---------- im: numpy array 2D image we want to process net: a pytorch model the network we want to use noiseModel: NoiseModel The noise model to be used. device: The device your network lives on, e.g. your GPU outScaling: float We found that scaling the output by a factor (default=10) helps to speedup training. Returns ---------- means: numpy array Image containing the means of the predicted prior distribution. This is similar to normal N2V. mseEst: numpy array Image containing the MMSE prediction, computed using the prior and noise model. ''' stdTorch = torch.Tensor(np.array(net.std)).to(device) meanTorch = torch.Tensor(np.array(net.mean)).to(device) #im=(im-net.mean)/net.std inputs_raw = torch.zeros(1, 1, im.shape[0], im.shape[1]) inputs_raw[0, :, :, :] = imgToTensor(im) # copy to GPU inputs_raw = inputs_raw.to(device) # normalize inputs = (inputs_raw - meanTorch) / stdTorch output = net(inputs) samples = (output).permute( 1, 0, 2, 3) * outScaling #We found that this factor can speed up training # denormalize samples = samples * stdTorch + meanTorch means = torch.mean(samples, dim=0, keepdim=True)[0, ...] # Sum up over all samples means = means.cpu().detach().numpy() means.shape = (output.shape[2], output.shape[3]) if noiseModel is not None: # call likelihood using denormalized observations and samples likelihoods = noiseModel.likelihood(inputs_raw, samples) mseEst = torch.sum(likelihoods * samples, dim=0, keepdim=True)[0, ...] # Sum up over all samples mseEst /= torch.sum(likelihoods, dim=0, keepdim=True)[0, ...] # Normalize # Get data from GPU mseEst = mseEst.cpu().detach().numpy() mseEst.shape = (output.shape[2], output.shape[3]) return means, mseEst else: return means, None