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
