%% End to end learning a raw camera pipeline
% DSLRs and many modern phone cameras offer the ability to save images as RAW, meaning that the original, untouched sensor data is saved directly to a file so that it can be processed later. This RAW data is the original counts measured by the camera sensor in the particular sensor layout of the device (e.g. many cameras use a particular arrangement of a Bayer pattern). Calling back to the analog film world, RAW images can conceptually be thought of as digital negatives, where they capture the original information at acquisition time that can later be processed into a final image using so-called digital darkroom software.
% As is discussed in excellent detail in [1], a RAW camera pipeline processes a RAW image into a final eye-pleasing RGB form in a multi-step, non-linear sequence of operations consistening of some combination of linearization, white-balancing, demosaicing, denoising, color space correction, and local/global brightness and contrast adjustment. Typically, these pipelines are designed as separate functional modules using traditional image processing concepts.
% Recently, deep learning literature has raised the possibility of end-to-end learning networks that can take as input RAW image inputs and yield as output a fully-formed RGB image. In [2], a novel approach to learning a pipeline for low-light images was demonstrated by collecting two images of the same scene, an RGB target that was well exposed and a RAW image of the same scene under-exposed by increasing shutter speed.
% In [3], it is shown that given RAW images from a lower end camera sensor and formed RGB images of the exact same scene from a higher end DSLR, a network can be trained to approximate the DSLR pipeline from the camera RAW data. This paper also releases a useful dataset, the Zurich RAW to RGB Dataset.
% This example demonstrates a simple U-net based deep learning approach to learning a camera pipeline using a combination of the concepts from [2] and [3].
%% Download Zurich RAW to RGB dataset
% The Zurich RAW to RGB Dataset can be downloaded here. The rest of this example will use the location of the data defined below.

zurichRAWToRGBDataPath = "/mathworks/devel/sandbox/ataylor/rawProcessing/ZurichRawDataset/";

%
%% Construct datastores for reading in dataset
% A typical workflow for managing large data in MATLAB is to construct datastores for reading and iterating through data.

dsInTrain = imageDatastore(zurichRAWToRGBDataPath + "train/huawei_raw");
dsOutTrain = imageDatastore(zurichRAWToRGBDataPath + "train/canon");

%
% In this example, the input and output images used for training are sub-images or patches taken from full-sized images which are also included in the dataset. The code below forms Datastores which contain full sized input and output images from the test set.

dsInFullTest = imageDatastore(zurichRAWToRGBDataPath + "test/huawei_full_resolution/");
targetFilesToInclude = extractBefore(extractAfter(string(dsInFullTest.Files),"huawei_full_resolution/"),".png");
dsOutFullTest = imageDatastore(zurichRAWToRGBDataPath + "full_resolution/canon");
dsOutFullTest.Files = dsOutFullTest.Files(contains(dsOutFullTest.Files,targetFilesToInclude));

%
% Below is an example of a fully formed RGB image in the dataset taken from a Canon DSLR

figure
imshow(preview(dsOutFullTest))

%
% The input datastore contains RAW CFA RGGB Bayer pattern data. The camera sensor is 10-bit data.
% The RAW data can also be visualized by demosaicing the CFA pattern and rescaling to the 10-bit range of the sensor. Because the demosaiced linear RGB data has not yet been white balanced, color adjusted, or contrast adjusted, the resulting image is dark and the colors are distorted.

exampleInput = preview(dsInFullTest);
rgbImageVis = single(demosaic(exampleInput,'rggb'))/2^10;
figure
imshow(rgbImageVis)

% As mentioned previously, in the training set, the data has been patched. The input RAW and output RGB patches in the training set can be visualized the same way. The input patch is rescaled to the range [0,1] to brighten it for visualization purposes, which introduces noticeable artifacts.

inputPatchExample = preview(dsInTrain);
inputPatchExample = single(demosaic(inputPatchExample,'rggb'))/2^10;
outputPatchExample = preview(dsOutTrain);
figure
montage({normalize(inputPatchExample,'range'),outputPatchExample})

% The last noteworthy thing about the Zurich RAW to RGB dataset is that the input RAW images in the training set are represented as both 8-bit and 16-bit PNG files. 8-bit images are used for patches where the sensor counts are in the range [0,255] for compact representation. No scaling  has been done on any of the RAW data.
%
%% Setup validation set
% The Zurich RAW to RGB dataset contains a large, separate test set. This example uses a small subset of patched images from the test set as a validation set.

dsInVal = imageDatastore(zurichRAWToRGBDataPath + "test/huawei_raw");
dsOutVal = imageDatastore(zurichRAWToRGBDataPath + "test/canon");
filesToInclude = randperm(dsInVal.numpartitions);
valSetSize = 200;
filesToInclude = filesToInclude(1:valSetSize);
dsInVal = subset(dsInVal,filesToInclude);
dsOutVal = subset(dsOutVal,filesToInclude);

%
%% Setup preprocessing pipline for input data
% Use the transform function of Datastore to define a pre-processing function that will be run for each image input image to the network. The function inputPipeline is shown in the Supporting Functions section at the bottom of  this example.

dsInTrain = transform(dsInTrain,@inputPipeline);
dsInVal = transform(dsInVal,@inputPipeline);
dsInFullTest = transform(dsInFullTest,@inputPipeline);

% After preprocessing, the RAW image data is H/2-by-W/2-by-4 with one red, one blue, and two green channels and is cast to single in the floating point range [0,1].

inputExample = preview(dsInTrain);
whos inputExample

%
%% Setup preprocessing pipeline for output data
% The target RGB images are stored on disk as unsigned 8-bit data. To make the computation of metrics and some specifics of network design more convenient, the target images are cast to single and scaled to the range [0,1].

dsOutTrain = transform(dsOutTrain,@setupOutputPipeline);
dsOutVal = transform(dsOutVal,@setupOutputPipeline);

% Again, verify that transform function definition is working as expected.

outputExample = preview(dsOutTrain);
whos outputExample

% The output, example, is twice the width and height of the input patches in 3-channel RGB form as expected.

figure, imshow(outputExample)

%
%% Build training and validation sets
% First, set the ReadSize property of the UnderlyingDatastore in each input and output datastore to optimize file i/o costs when reading through data.

dsInTrain.UnderlyingDatastore.ReadSize = 16;
dsOutTrain.UnderlyingDatastore.ReadSize = 16;
dsInVal.UnderlyingDatastore.ReadSize = 16;
dsOutVal.UnderlyingDatastore.ReadSize = 16;

% Use combine to form datastores which are conceptually the horizontally concatenated data from each input datastore.

dsTrain = combine(dsInTrain,dsOutTrain);
dsVal = combine(dsInVal,dsOutVal);

% Use the preview function to inspect an example observation from the training set.

preview(dsTrain)

%
%% Data augmenation
% The last remaining item related to data management in this example is to add some randomized data augmentation to our training set. In this example, we use a simple scheme of randomized 90-degree rotations and reflections, both of which avoid the need for interpolation or fill values in the RAW or RGB image data. The function augmentData is shown below in the supporting functions section.

dsTrainAug = transform(dsTrain,@augmentData);

% One last time, display and visualize the output of preview on the datastore to make sure the pre-processing pipeline is setup correctly. The input is visualized by looking at the first channel, blue, rescaled to [0,1].

exampleAug = preview(dsTrainAug);
disp(exampleAug)
exampleInput = exampleAug{1,1};
exampleOutput = exampleAug{1,2};
figure
montage({rescale(exampleInput(:,:,1)),exampleOutput})

% The simple rescaling of the RAW blue channel results in a lot of noise shown in the left hand side image, but the point of this visualization is to check that the input RAW and output RGB have been augmented correctly, which it appears they have from the common orientation in the input and output.
%
%% Compute training set statistics for input normalization
% Use tall to compute per-channel mean reduction across training set. The mean statistics will be used in the input layer in the model at both training and test time to perform mean centering of inputs. The training set is quite large (46839 images), so the mean statistics may take several minutes to compute depending on parallel pool settings and CPU hardware.

dsIn = copy(dsInTrain);
dsIn.UnderlyingDatastore.ReadSize = 1;
t = tall(dsIn);
perChannelMean = gather(mean(t,[1 2]))

%
%% Network Architecture
% Influenced by [1], this example uses a simple U-net architecture. The buildUNet function is included below in the supporting functions section.

lgraph = buildUNet(perChannelMean);
net = dlnetwork(lgraph);
plot(lgraph)

% The network uses a 4-level U-net encoder/decoder structure with 32 convolutional filters at the first encoder stage. Instance normalization is applied to each channel activation following convolution operations starting at the second encoder level and is symmetrically applied up to the penultimate stage of the decoder. In the decoder stage transposedConv2dLayer is used with the same 3-by-3 filter size as the corresponding convolution2dLayer layers in the encoder with the same padding options. The second to last layer in the network uses pixelShuffle upsampling to move from the H/2-by-W/2-by-12 channel size of the activations from the final convolution to H-by-W-by 3 channel activations. The final layer in the network uses a hyperbolic tangent function to encourage outputs in the desired [0,1] range. These final two operations are described in detail in the secitons below.
%% Pixel Shuffle
% Pixel shuffle upsampling is used instead of transposed convolution in the final upsampling stage of the decoder. In [6], it is shown that convolution followed by pixel shuffle can be used to define sub-pixel convolution for super resolution applications to avoid the checkboard artifacts that can arise from transposed convolution. Since the model needs to map H/2-by-W/2-by-4 RAW inputs to a W-by-H-by-3 RGB outputs, the final upsampling stage of the model can be thought of similarly to super resolution where the number of spatial samples grows from the input to the output.
% The code below provides a visualization to provide an intuition of how pixel shuffle upsampling works. The pixelShuffle function is provided below in the supporting functions section. Start by creating a 2x2x4 input with unique values where the first two dimensions are spatial dimensions that will be resampled and the 3rd dimension is a channel dimension.

input = reshape(1:16,2,2,4);
disp(input);

% Given an upsampling factor of S, pixel shuffle takes a H-by-W-by-C input and yields a S*H-by-S*W-by-\frac{C}{S^2 } output.

output = pixelShuffle(input,2);
disp(output)

% The pixel shuffle function grows the spatial dimensions of the output by mapping information from channel dimensions at a given spatial location into SxS sized spatial blocks in the output in which each channel contributes to a consistent spatial position relative to its neighbors during upsampling.
%% Hyperbolic tangent final activation
% The final layer in the model is defined using hyperbolic tangent. The motivation of this layer is to encourage but not strictly enforce a [0,1] range in the RGB outputs from the network. Similar to the idea of leaky relu vs. relu, this function is scaled and shifted to encourage gradient propagation for input values outside [0,1] during the early stages of training. The function scaledAndShiftedTanh is shown in the supporting functions below.

x = -3:0.01:3;
plot(x,scaledAndShiftedTanh(x))

%
%% Loss function
% This example uses a blend of three different losses. A basic L1 loss is used on a per sample basis between the output image from the network and the target image. As shown in [4], L1 is often a better choice than L2 for image processing applications because it can help reduce blurring artifacts. Influenced by [3], the content loss portion of perceptual loss from [5] is used with equal weighting of L1 losses of low and higher level features extracted from VGG16 with the outputs of the network and the target images. The content loss helps the network learn both accurate structural information from the higher level features and lower level edge and color information.  Weighting factors for additively combining the separately losses are defined emperically by computing the separate losses with training data with the randomly initialized network and choosing weights that roughly balance the contribution of each separate loss, as shown below. The loss definitions are included in the supporting functions below.

s = load('/sandbox/ataylor/rawProcessing/code/vgg16.mat');
vggNet = dlnetwork(layerGraph(s.net.Layers(1:31)));

trainingData = preview(dsTrainAug);
rawInput = dlarray(gpuArray(trainingData{1,1}),'SSC');
targetOutput = dlarray(gpuArray(trainingData{1,2}),'SSC');
networkOutput = forward(net,rawInput);

l1LossExample = maeLoss(networkOutput,targetOutput);
contentLossExample = contentLoss(vggNet,networkOutput,targetOutput);

weightContent = l1LossExample/contentLossExample;

overallLossExample = overallLoss(l1LossExample,contentLossExample,weightContent)

%
%% Construct minibatchqueue for batching training and validation set during training
% This example uses a custom training loop. When using custom training loops, minibatchqueue is a useful object for managing the minibatching of observations, along with moving data to the GPU and casting to dlarray, which is the MATLAB datatype that enables auto differentiation in deep learning applications.

miniBatchSize = 12;
valBatchSize = 10;
trainingQueue = minibatchqueue(dsTrainAug,'MiniBatchSize',miniBatchSize,'PartialMiniBatch','discard','MiniBatchFormat','SSCB');
validationQueue = minibatchqueue(dsVal,'MiniBatchSize', valBatchSize,'MiniBatchFormat','SSCB');

% Within a custom training loop, the next function of minibatchqueue yields the next minibatch of data.

[raw,targetRGB] = next(trainingQueue);
whos raw
whos targetRGB

% The outputs from next are of class dlarray, already cast to gpuArray so these 4-D arrays are already on the GPU and ready for training.
%
%% Define overall function for computing loss and gradients
% As part of our training loop, define a function which will return the gradients for each learnable parameter in the model along with the current loss. The function modelGradients does this and is included in the supporting functions section. The gradients and loss for each batch of data are computed in a custom training loop using dlfeval. In the code below, indexing is used to only compute gradients for a single observation to reduce GPU memory requirements.

[grad,loss] = dlfeval(@modelGradients,net,vggNet,raw(:,:,:,1),targetRGB(:,:,:,1),weightContent)

% The gradient information is returned as a table which includes the layer, parameter name and value for each learnable parameter in the model. This gradient information is used by the solver, which is adam in this example, to update the model in each iteration.
%
%% Define custom training loop
% Start by defining the training options that will be used within the custom training loop to control aspects of training

learnRate = 5e-5;
numEpochs = 20;
checkpointDir = fullfile(tempdir,'checkpoints',filesep);
fullSizedImageDir = fullfile(tempdir,'fullSizedOutputImages',filesep);
if ~exist(checkpointDir,'dir')
    mkdir(checkpointDir);
    mkdir(fullSizedImageDir)
end

% Define a full sized input RAW image that can be used to visually monitor the progress of training as the model trains. This data will be written to fullSizedImageDir defined above as a .png file.

fullSizedRaw = dlarray(gpuArray(preview(dsInFullTest)),'SSCB');

% The complete custom training loop for training the model is included below. To run the code, set the variable trainModel below to true. A pre-trained version of the model is included and used in subsequent sections. The network was trained on a NVIDIA Titan RTX, taking approximately 36 hours to train.

trainNetwork = false;

if trainNetwork
    % Initialize training plot
    [hFig,batchLine,validationLine] = initializeTrainingPlot(); %#ok<UNRCH> 
    
    % Initialize Adam solver state
    [averageGrad,averageSqGrad] = deal([]);
    iteration = 0;
    
    start = tic;
    for epoch = 1:numEpochs
        reset(trainingQueue);
        shuffle(trainingQueue);
        while hasdata(trainingQueue)
            [raw,targetRGB] = next(trainingQueue);  
            
            [grad,loss] = dlfeval(@modelGradients,net,vggNet,raw,targetRGB,weightContent);
            
            iteration = iteration + 1;
            
            [net,averageGrad,averageSqGrad] = adamupdate(net,grad,averageGrad,averageSqGrad,iteration,...
                learnRate);
              
            updatePlot(batchLine,validationLine,iteration,loss,start,epoch,...
                validationQueue,valSetSize,valBatchSize,net,vggNet,weightContent);
        end
        % Write out current full sized image output of network and save
        % checkpoint of network state.
        fullSizedOutput = im2uint8(gather(extractdata(forward(net,fullSizedRaw))));
        imwrite(fullSizedOutput,fullSizedImageDir + "epoch" + epoch + ".png");        
        save(checkpointDir + "epoch" + epoch,'net');
    end
    
    save(checkpointDir + "finalModelState.mat", 'net');
    
end

if ~trainNetwork
    imshow('trainingPlot.png')
end

%
%% Compute full sized output images from test set
% Once the network is trained, either by running the training loop above or by using the included pre-trained network, the next step is to evaluate the network on full sized RAW input images from the test set.

if ~trainNetwork
    s = load('finalModelState.mat');
    net = s.net;
end

% Use a minibatchqueue to iterate through the test set.

testQueue = minibatchqueue(dsInFullTest,'MiniBatchSize',1,'MiniBatchFormat','SSCB');

% Iterate through the queue of test RAW images, calling forward on the network to predict the output sRGB form learned by the network RAW pipeline and storing the result in a cell array.

outputSize = 2*size(preview(dsInFullTest),[1 2]);

outputImages = zeros([outputSize,3,dsInFullTest.numpartitions],'uint8');
idx = 1;
while hasdata(testQueue)
   raw = next(testQueue);
   rgbOut = forward(net,raw);
   rgbOut = im2uint8(gather(extractdata(rgbOut)));
   outputImages(:,:,:,idx) = rgbOut;
   idx = idx+1;
end

% Get a sense of the overall output by first looking at a montage view.

montage(outputImages)

% Recall that in this example, the purpose is to train a network to learn how to produce the final sRGB output of a target dSLR ISP given RAW sensor data from a phone camera. Read in the target sRGB images from the test set from the high end DSLR.

targetRGB = readall(dsOutFullTest);

% Display an example image predicted by the network

hIm = imshow(outputImages(:,:,:,1));
hIm.Interpolation = 'bilinear';

% And the corresponding target image

hImTarget = imshow(targetRGB{1});
hImTarget.Interpolation = 'bilinear';

% The images produced by the network in this example display similar characteristics throughout the test set, fairly accurate reproduction of structural content, but noticeably different color reproduction than the target DSLR ISP. The network produces colors which are a bit more saturated.
%  As described in [3], an ideal network architecture for learning a camera pipeline makes use of multiple scales from more global features which describe color and contrast and uses this features as part of learning finer spatial scale features. Though the colors from the simple U-net architecture are not the same as the DSLR targets, the images are still qualitively pleasing in many cases.
% Note that all of the images above are downsampled as part of displaying them in a live script. To get the best sense of the network output and target images, it is best to view the images full size outside of the context of live script.
%
%% Image quality metrics
% The previous section uses a qualitive judgement of how similar the network RGB outputs are to the DSLR targets. A quantitive way of accessing this is to use a reference-based quality metric like MSSIM or PSNR like in [3]. Because of sensor differences between the camera and DSLR used to collect the full resolution input and target images in the test set, the scenes are not registered and are not the same size, making reference-based comparison of the full resolution images from the network and the DSLR ISP difficult.
% This example uses the patched images in the test set which are same size and registered to judge image quality. The psnr and multissim functions in MATLAB require grayscale inputs because these metrics aren't well defined for RGB images. RGB approximations of these metrics are included in the supporting functions section which take the per-channel mean of the metrics on the R,G and B channels separately, which is a common approximation.

dsInPatchTest = imageDatastore(zurichRAWToRGBDataPath + "test/huawei_raw");
dsInPatchTest = transform(dsInPatchTest,@inputPipeline);
dsOutPatchTest = imageDatastore(zurichRAWToRGBDataPath + "test/canon");
patchTestSet = combine(dsInPatchTest,dsOutPatchTest);
testPatchQueue = minibatchqueue(patchTestSet,'MiniBatchSize',16,'MiniBatchFormat','SSCB');

% Iterate through the test set and compute mean mssim and mean psnr over the test set. This may take a couple minutes to compute over the test set.

totalMSSIM = 0;
totalPSNR = 0;
numObservations = dsOutPatchTest.numpartitions;
while hasdata(testPatchQueue)
    [raw,targetRGB] = next(testPatchQueue);
    outputRGB = forward(net,raw);
    [mssimOut,psnrOut] = computeQualityMetrics(outputRGB,targetRGB);
    totalMSSIM = totalMSSIM + mssimOut;
    totalPSNR = totalPSNR + psnrOut;
end

meanMSSIM = totalMSSIM / numObservations
meanPSNR = totalPSNR / numObservations

% This result is consistent with the similar U-net approach from [3] for mean MSSIM and competitive with the PyNet approach in [3] in mean PSNR. The differences in loss functions and use of pixel shuffle upsampling compared to [3] likely account for these differences.
%
%% Conclusion
% This example demonstrated that it is possible to train a convolutional neural network to approximate the ISP of a high-end DSLR given RAW input data from a phone camera sensor. To be viable in realistic settings, the network size and inference time would need to be reduced to make it viable for deployment. Also, some additional work would need to be done to correct the color accuracy issues noted in the approach in this example. That said, the output of this relatively simple approach is promising and the overall takeway of this example is that it is possible to end-to-end learn a RAW camera pipeline using a straight foward deep learning approach.
% Copyright 2020 The MathWorks, Inc.
%
%% References
% 1) Processing RAW Images in MATLAB, R Sumner, 2014, https://rcsumner.net/raw_guide/RAWguide.pdf
% 2) Learning to See in the Dark, Chen Chen, Qifeng Chen, Jia Xu, Vladlen Koltun, https://arxiv.org/abs/1805.01934
% 3) Replacing Mobile Camera ISP with a Single Deep Learning Model, Andrey Ignatov, Luc Van Gool, Radu Timofte, https://arxiv.org/abs/2002.05509, Project Website
% 4) Loss Functions for Neural Networks for Image Processing, Hang Zhao, Orazio Gallo, Iuri Frosio, Jan Kautz,  https://arxiv.org/abs/1511.08861
% 5) Perceptual Losses for Real-Time Style Transfer and Super-Resolution, Justin Johnson, Alexandre Alahi, Li Fei-Fei, https://arxiv.org/abs/1603.08155
% 6) Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network, Wenzhe Shi, Jose Caballero, Ferenc Huszár, Johannes Totz, Andrew P. Aitken, Rob Bishop, Daniel Rueckert, Zehan Wang, https://arxiv.org/abs/1609.05158
%
%% Supporting functions

function out = inputPipeline(in)

if ~iscell(in)
    out = packRawImage(in);
else
    out = cell(size(in));
    for idx = 1:size(in,1)
        out{idx} = packRawImage(in{idx});
    end
end
end

function out = packRawImage(cfa)
% This function takes input CFA bayer pattern images and extracts blue,
% red, and two green channels from the bayer pattern and forms a
% H/2-by-W/2-by-4 image with each channel represented as a separate plane.
% The RAW counts are normalized to [0,1] according to the range of the
% 10-bit sensor.

% RGGB order
blue = cfa(2:2:end,2:2:end);
green1 = cfa(1:2:end,2:2:end);
red = cfa(1:2:end,1:2:end);
green2 = cfa(2:2:end,1:2:end);

rawCombined = cat(3,blue,green1,red,green2);
out = single(rawCombined) / (4*255); % 10-bit sensor

end

function out = setupOutputPipeline(in)

if iscell(in)
    out = cell(size(in));
    for idx = 1:size(in,1)
        out{idx} = normalizeTargetImages(in{idx});
    end
else
    out = normalizeTargetImages(in);
end

end

function out = normalizeTargetImages(in)
out = single(in)./255;
end

function dataOut = augmentData(dataIn)
% Randomized 90 degree rotation + horizontal reflection

dataOut = cell(size(dataIn));
for idx = 1:size(dataIn,1)
    
    % Randomized rotation
    rotAmount = randi(4)-1;
    dataOut(idx,:) = {rot90(dataIn{idx,1},rotAmount),rot90(dataIn{idx,2},rotAmount)};
    
    % Randomized horizontal reflection
    xReflection = rand > 0.5;
    if xReflection
        dataOut(idx,:) = {fliplr(dataOut{idx,1}),fliplr(dataOut{idx,2})};
    end
end
end

function out = scaledAndShiftedTanh(x)
    out = 0.58*tanh(x) + 0.5;
end

function output = pixelShuffle(input, scale_factor)

sizeIn = size(input);

outChannels = sizeIn(3) / (scale_factor^2);
outHeight = sizeIn(1) * scale_factor;
outWidth = sizeIn(2) * scale_factor;
batchSize = size(input,4);

temp = reshape(input,[sizeIn(1:2),scale_factor,scale_factor,outChannels,batchSize]);
temp = permute(temp,[3 1 4 2 5 6]);
output = reshape(temp,[outHeight,outWidth,outChannels,size(input,4)]);

end

function lgraph = buildUNet(inputMeans)

lgraph = layerGraph();

tempLayers = [
    imageInputLayer([256 256 4],"Name","ImageInputLayer",'Normalization','zerocenter','Mean',inputMeans)
    convolution2dLayer([3 3],32,"Name","Encoder-Stage-1-Conv-1","Padding","same","WeightsInitializer", 'narrow-normal')
    leakyReluLayer(0.2,"Name","Encoder-Stage-1-ReLU-1")
    convolution2dLayer([3 3],32,"Name","Encoder-Stage-1-Conv-2","Padding","same","WeightsInitializer", 'narrow-normal')
    leakyReluLayer(0.2,"Name","Encoder-Stage-1-ReLU-2")];
lgraph = addLayers(lgraph,tempLayers);

tempLayers = [
    maxPooling2dLayer([2 2],"Name","Encoder-Stage-1-MaxPool","Stride",[2 2])
    convolution2dLayer([3 3],64,"Name","Encoder-Stage-2-Conv-1","Padding","same","WeightsInitializer", 'narrow-normal')
    groupNormalizationLayer("channel-wise","Name",'cn1')
    leakyReluLayer(0.2,"Name","Encoder-Stage-2-ReLU-1")
    convolution2dLayer([3 3],64,"Name","Encoder-Stage-2-Conv-2","Padding","same","WeightsInitializer", 'narrow-normal')
    groupNormalizationLayer("channel-wise","Name",'cn2')
    leakyReluLayer(0.2,"Name","Encoder-Stage-2-ReLU-2")];
lgraph = addLayers(lgraph,tempLayers);

tempLayers = [
    maxPooling2dLayer([2 2],"Name","Encoder-Stage-2-MaxPool","Stride",[2 2])
    convolution2dLayer([3 3],128,"Name","Encoder-Stage-3-Conv-1","Padding","same","WeightsInitializer", 'narrow-normal')
    groupNormalizationLayer("channel-wise","Name",'cn3')
    leakyReluLayer(0.2,"Name","Encoder-Stage-3-ReLU-1")
    convolution2dLayer([3 3],128,"Name","Encoder-Stage-3-Conv-2","Padding","same","WeightsInitializer", 'narrow-normal')
    groupNormalizationLayer("channel-wise","Name",'cn4')
    leakyReluLayer(0.2,"Name","Encoder-Stage-3-ReLU-2")];
lgraph = addLayers(lgraph,tempLayers);

tempLayers = [
    maxPooling2dLayer([2 2],"Name","Encoder-Stage-3-MaxPool","Stride",[2 2])
    convolution2dLayer([3 3],256,"Name","Encoder-Stage-4-Conv-1","Padding","same","WeightsInitializer", 'narrow-normal')
    groupNormalizationLayer("channel-wise","Name",'cn5')
    leakyReluLayer(0.2,"Name","Encoder-Stage-4-ReLU-1")
    convolution2dLayer([3 3],256,"Name","Encoder-Stage-4-Conv-2","Padding","same","WeightsInitializer", 'narrow-normal')
    groupNormalizationLayer("channel-wise","Name",'cn6')
    leakyReluLayer(0.2,"Name","Encoder-Stage-4-ReLU-2")];
lgraph = addLayers(lgraph,tempLayers);

tempLayers = [
    maxPooling2dLayer([2 2],"Name","Encoder-Stage-4-MaxPool","Stride",[2 2])
    convolution2dLayer([3 3],512,"Name","Bridge-Conv-1","Padding","same","WeightsInitializer", 'narrow-normal')
    groupNormalizationLayer("channel-wise","Name",'cn7')
    leakyReluLayer(0.2,"Name","Bridge-ReLU-1")
    convolution2dLayer([3 3],512,"Name","Bridge-Conv-2","Padding","same","WeightsInitializer", 'narrow-normal')
    groupNormalizationLayer("channel-wise","Name",'cn8')
    leakyReluLayer(0.2,"Name","Bridge-ReLU-2")
    transposedConv2dLayer([3 3],256,"Name","Decoder-Stage-1-UpConv","Stride",[2 2],"WeightsInitializer", 'narrow-normal','Cropping','same')
    leakyReluLayer(0.2,"Name","Decoder-Stage-1-UpReLU")];
lgraph = addLayers(lgraph,tempLayers);

tempLayers = [
    depthConcatenationLayer(2,"Name","Decoder-Stage-1-DepthConcatenation")
    convolution2dLayer([3 3],256,"Name","Decoder-Stage-1-Conv-1","Padding","same","WeightsInitializer", 'narrow-normal')
    groupNormalizationLayer("channel-wise","Name",'cn9')
    leakyReluLayer(0.2,"Name","Decoder-Stage-1-ReLU-1")
    convolution2dLayer([3 3],256,"Name","Decoder-Stage-1-Conv-2","Padding","same","WeightsInitializer", 'narrow-normal')
    groupNormalizationLayer("channel-wise","Name",'cn10')
    leakyReluLayer(0.2,"Name","Decoder-Stage-1-ReLU-2")
    transposedConv2dLayer([3 3],128,"Name","Decoder-Stage-2-UpConv","Stride",[2 2],"WeightsInitializer", 'narrow-normal','Cropping','same')
    leakyReluLayer(0.2,"Name","Decoder-Stage-2-UpReLU")];
lgraph = addLayers(lgraph,tempLayers);

tempLayers = [
    depthConcatenationLayer(2,"Name","Decoder-Stage-2-DepthConcatenation")
    convolution2dLayer([3 3],128,"Name","Decoder-Stage-2-Conv-1","Padding","same","WeightsInitializer", 'narrow-normal')
    groupNormalizationLayer("channel-wise","Name",'cn11')
    leakyReluLayer(0.2,"Name","Decoder-Stage-2-ReLU-1")
    convolution2dLayer([3 3],128,"Name","Decoder-Stage-2-Conv-2","Padding","same","WeightsInitializer", 'narrow-normal')
    groupNormalizationLayer("channel-wise","Name",'cn12')
    leakyReluLayer(0.2,"Name","Decoder-Stage-2-ReLU-2")
    transposedConv2dLayer([3 3],64,"Name","Decoder-Stage-3-UpConv","Stride",[2 2],"WeightsInitializer", 'narrow-normal','Cropping','same')
    leakyReluLayer(0.2,"Name","Decoder-Stage-3-UpReLU")];
lgraph = addLayers(lgraph,tempLayers);

tempLayers = [
    depthConcatenationLayer(2,"Name","Decoder-Stage-3-DepthConcatenation")
    convolution2dLayer([3 3],64,"Name","Decoder-Stage-3-Conv-1","Padding","same","WeightsInitializer", 'narrow-normal')
    groupNormalizationLayer("channel-wise","Name",'cn13')
    leakyReluLayer(0.2,"Name","Decoder-Stage-3-ReLU-1")
    convolution2dLayer([3 3],64,"Name","Decoder-Stage-3-Conv-2","Padding","same","WeightsInitializer", 'narrow-normal')
    groupNormalizationLayer("channel-wise","Name",'cn14')
    leakyReluLayer(0.2,"Name","Decoder-Stage-3-ReLU-2")
    transposedConv2dLayer([3 3],32,"Name","Decoder-Stage-4-UpConv","Stride",[2 2],"WeightsInitializer", 'narrow-normal','Cropping','same')
    leakyReluLayer(0.2,"Name","Decoder-Stage-4-UpReLU")];
lgraph = addLayers(lgraph,tempLayers);

tempLayers = [
    depthConcatenationLayer(2,"Name","Decoder-Stage-4-DepthConcatenation")
    convolution2dLayer([3 3],32,"Name","Decoder-Stage-4-Conv-1","Padding","same","WeightsInitializer", 'narrow-normal')
    leakyReluLayer(0.2,"Name","Decoder-Stage-4-ReLU-1")
    convolution2dLayer([3 3],32,"Name","Decoder-Stage-4-Conv-2","Padding","same","WeightsInitializer", 'narrow-normal')
    leakyReluLayer(0.2,"Name","Decoder-Stage-4-ReLU-2")
    convolution2dLayer([3 3],12,"Name","Decoder-Stage-4-Conv-3","Padding","same","WeightsInitializer", 'narrow-normal')
    functionLayer("pixelShuffle",@(x) pixelShuffle(x,2))
    functionLayer("tanhLayer",@(x) scaledAndShiftedTanh(x))];

lgraph = addLayers(lgraph,tempLayers);

lgraph = connectLayers(lgraph,"Encoder-Stage-1-ReLU-2","Encoder-Stage-1-MaxPool");
lgraph = connectLayers(lgraph,"Encoder-Stage-1-ReLU-2","Decoder-Stage-4-DepthConcatenation/in2");
lgraph = connectLayers(lgraph,"Encoder-Stage-2-ReLU-2","Encoder-Stage-2-MaxPool");
lgraph = connectLayers(lgraph,"Encoder-Stage-2-ReLU-2","Decoder-Stage-3-DepthConcatenation/in2");
lgraph = connectLayers(lgraph,"Encoder-Stage-3-ReLU-2","Encoder-Stage-3-MaxPool");
lgraph = connectLayers(lgraph,"Encoder-Stage-3-ReLU-2","Decoder-Stage-2-DepthConcatenation/in2");
lgraph = connectLayers(lgraph,"Encoder-Stage-4-ReLU-2",'Encoder-Stage-4-MaxPool');
lgraph = connectLayers(lgraph,"Encoder-Stage-4-ReLU-2","Decoder-Stage-1-DepthConcatenation/in2");
lgraph = connectLayers(lgraph,"Decoder-Stage-1-UpReLU","Decoder-Stage-1-DepthConcatenation/in1");
lgraph = connectLayers(lgraph,"Decoder-Stage-2-UpReLU","Decoder-Stage-2-DepthConcatenation/in1");
lgraph = connectLayers(lgraph,"Decoder-Stage-3-UpReLU","Decoder-Stage-3-DepthConcatenation/in1");
lgraph = connectLayers(lgraph,"Decoder-Stage-4-UpReLU","Decoder-Stage-4-DepthConcatenation/in1");

end

function loss = maeLoss(Y,T)
loss = mean(abs(Y-T),'all');
end

function loss = mseLoss(Y,T)
loss = mean((Y-T).^2,'all');
end

function loss = contentLoss(net,Y,T)

layers = ["relu1_1","relu1_2","relu2_1","relu2_2","relu3_1","relu3_2","relu3_3","relu4_1"];
[T1,T2,T3,T4,T5,T6,T7,T8] = forward(net,T,'Outputs',layers);
[X1,X2,X3,X4,X5,X6,X7,X8] = forward(net,Y,'Outputs',layers);

l1 = mseLoss(X1,T1);
l2 = mseLoss(X2,T2);
l3 = mseLoss(X3,T3);
l4 = mseLoss(X4,T4);
l5 = mseLoss(X5,T5);
l6 = mseLoss(X6,T6);
l7 = mseLoss(X7,T7);
l8 = mseLoss(X8,T8);

% Weight each layer activation roughly equally
loss = l1 + 0.0449*l2 + 0.0107*l3 + .0023*l4 + 6.9445e-04*l5 +...
            2.0787e-04*l6 + 2.0118e-04*l7 + 6.4759e-04*l8;

end

function loss = overallLoss(lossMAE,lossContent,weightContent)
loss = lossMAE + weightContent.*lossContent;
end

function [gradients, loss] = modelGradients(dlnet,vggNet,Xpatch,Target,weightContent)
Y = forward(dlnet,Xpatch);
lossMAE = maeLoss(Y,Target);
lossContent = contentLoss(vggNet,Y,Target);
loss = overallLoss(lossMAE,lossContent,weightContent);
gradients = dlgradient(loss,dlnet.Learnables);
end

function updatePlot(batchLine,validationLine,iteration,loss,start,epoch,validationQueue,...
    valSetSize,valBatchSize,net,vggNet,contentWeight) %#ok<DEFNU> 

% Use to control frequency of update to plots
miniBatchLossFrequency = 5;
validationFrequency = 200;

if ~mod(iteration,miniBatchLossFrequency) || iteration == 1
    updateLossLine(batchLine,iteration,loss,start,epoch);
end

if ~mod(iteration,validationFrequency) || iteration == 1
    valLoss = computeValidationLoss(net,validationQueue,vggNet,valSetSize,valBatchSize,contentWeight);
    updateLossLine(validationLine,iteration,valLoss,start,epoch);
end

end

function valLoss = computeValidationLoss(net,validationQueue,vggNet,valSetSize,valBatchSize,contentWeight)

lossTotal = 0.0;
while hasdata(validationQueue)
    [Xpatch,Yexp] = next(validationQueue);
    Yexp = gpuArray(dlarray(Yexp,'SSCB'));
    Ypred = forward(net,Xpatch);
    
    valLoss = overallLoss(maeLoss(Yexp,Ypred),contentLoss(vggNet,Yexp,Ypred),contentWeight);
    lossTotal = lossTotal + valLoss;
end

valLoss = lossTotal / (valSetSize / valBatchSize);
reset(validationQueue);

end

function [hFig,batchLine,validationLine] = initializeTrainingPlot %#ok<DEFNU> 

hFig = figure;
hAx = axes;
batchLine = animatedline('Color','red','Parent',hAx);
validationLine = animatedline('Color','blue','LineStyle','--','Parent',hAx);
legend(hAx,'MiniBatchLoss','ValidationLoss');

grid on
grid minor

end

function updateLossLine(hLine,t,lossVal,start,epoch)

lossVal = double(gather(extractdata(lossVal)));
addpoints(hLine,t,lossVal);
D = duration(0,0,toc(start),'Format','hh:mm:ss');
title("Epoch" + epoch + ", Elapsed: " + string(D));
drawnow limitrate

end

function [mssimSum,psnrSum] = computeQualityMetrics(I,Iref)
% Compute psnr and mssim, mean reducing over channel dim and sum reducing
% over observation dimension

I = gather(extractdata(I));
I = im2uint8(I); % Network predicts [0,1]. Scale to [0,255] and cast to uint8.
Iref = uint8(gather(extractdata(Iref)));

numObservations = size(I,4);

mssimSum = 0;
psnrSum = 0;
for idx = 1:numObservations
    mssimRGB = mean([multissim(I(:,:,1,idx),Iref(:,:,1,idx)),...
                 multissim(I(:,:,2,idx),Iref(:,:,2,idx)),...
                 multissim(I(:,:,3,idx),Iref(:,:,3,idx))]);
    psnrRGB = mean([psnr(I(:,:,1,idx),Iref(:,:,1,idx)),...
                 psnr(I(:,:,2,idx),Iref(:,:,2,idx)),...
                 psnr(I(:,:,3,idx),Iref(:,:,3,idx))]);
    mssimSum = mssimSum + mssimRGB;
    psnrSum = psnrSum + psnrRGB;
end

end