%% Generate Image from Segmentation Map Using Deep Learning
% This example shows how to generate a synthetic image of a scene from a semantic segmentation map using a Pix2PixHD conditional generative adversarial network (CGAN).
% Pix2PixHD [1] consists of two networks that are trained simultaneously to maximize the performance of both.
% The generator is an encoder-decoder style neural network that generates a scene image from a semantic segmentation map. A CGAN network trains the generator to generate a scene image that the discriminator misclassifies as real.
% The discriminator is a fully convolutional neural network that compares a generated scene image and the corresponding real image and attempts to classify them as fake and real, respectively. A CGAN network trains the discriminator to correctly distinguish between generated and real image.
% The generator and discriminator networks compete against each other during training. The training converges when neither network can improve further.
%
%% Download CamVid Data Set
% This example uses the CamVid data set [2] from the University of Cambridge for training. This data set is a collection of 701 images containing street-level views obtained while driving. The data set provides pixel labels for 32 semantic classes including car, pedestrian, and road.
% Download the CamVid data set from these URLs. The download time depends on your internet connection.

imageURL = 'http://web4.cs.ucl.ac.uk/staff/g.brostow/MotionSegRecData/files/701_StillsRaw_full.zip';
labelURL = 'http://web4.cs.ucl.ac.uk/staff/g.brostow/MotionSegRecData/data/LabeledApproved_full.zip';

dataDir = fullfile(tempdir,'CamVid'); 
downloadCamVidData(dataDir,imageURL,labelURL);
imgDir = fullfile(dataDir,"images","701_StillsRaw_full");
labelDir = fullfile(dataDir,'labels');

%
%% Preprocess Training Data
% Create an ImageDatastore to store the images in the CamVid data set.

imds = imageDatastore(imgDir);
imageSize = [576 768];

% Define the class names and pixel label IDs of the 32 classes in the CamVid data set using the helper function defineCamVid32ClassesAndPixelLabelIDs. This function is attached to the example as a supporting file.

numClasses = 32;
[classes,labelIDs] = defineCamVid32ClassesAndPixelLabelIDs;

% Create a PixelLabelDatastore to store the pixel label images.

pxds = pixelLabelDatastore(labelDir,classes,labelIDs);

% Preview a pixel label image and the corresponding ground truth scene image. Convert the labels from categorical labels to RGB colors by using the label2rgb function, then display the pixel label image and ground truth image in a montage.

im = preview(imds);
px = preview(pxds);
px = label2rgb(px);
montage({px,im})

%
% Partition the data into training and test sets using the helper function partitionCamVidForPix2PixHD. This function is attached to the example as a supporting file. The helper function splits the data into 648 training files and 32 test files.

[imdsTrain,imdsTest,pxdsTrain,pxdsTest] = partitionCamVidForPix2PixHD(imds,pxds,classes,labelIDs);

% Use the combine function to combine the pixel label images and ground truth scene images into a single datastore.

dsTrain = combine(pxdsTrain,imdsTrain);

%
% Augment the training and validation data by using the transform function with custom preprocessing operations specified by the helper function preprocessCamVidForPix2PixHD. This helper function is attached to the example as a supporting file.
% The preprocessCamVidForPix2PixHD function performs these operations:
% Scale the ground truth data to the range [-1,1]. This range matches the range of the tanhLayer in the generator network.
% Resize the image and labels to the output size of the network, 576-by-768 pixels, using nearest neighbor downsampling.
% Convert the single channel segmentation map to a 32-channel one-hot encoded segmentation map. This map represents each class using a unique binary mask, in which pixels of the class are 1 and all other pixels are 0. The map concatenates the masks of each class along the depth direction.
% Randomly flip image and pixel label pairs in the horizontal direction.

dsTrain = transform(dsTrain,@(x) preprocessCamVidForPix2PixHD(x,imageSize));

%
% Preview the channels of a one-hot encoded segmentation map in a montage.

map = preview(dsTrain);
montage(map{1},'Size',[4 8],'Bordersize',5,'BackgroundColor','b')

%
%% Create Generator Network
% Define a generator network that generates a scene image from a depth-wise one-hot encoded segmentation map. This input has same height and width as the original segmentation map and the same number of channels as classes.

inputSize = [imageSize numClasses];

% Create the initial subnetwork layers. Specify the filter size and number of filters in first convolutional layer of the generator. reflectionPad2d is a custom layer implemented specifically for this example. This layer is attached to the example as a supporting file.

filterSize = [7,7];
numFiltersGenerator = 64;

layers = [ ...
    imageInputLayer(inputSize,'Normalization','none','Name','inputLayer') ...
    reflectionPad2dLayer(3,'iPad') ...
    convolution2dLayer(filterSize,numFiltersGenerator,'Name','iConv', ...
        "BiasLearnRateFactor",0,"WeightsInitializer",'he',"BiasInitializer",'zeros') ...
    batchNormalizationLayer('Name','iBn', ...
        'Scale',ones([1 1 numFiltersGenerator]),'ScaleLearnRateFactor',0, ...
        'Offset',zeros([1 1 numFiltersGenerator]),'OffsetLearnRateFactor',0)...
    reluLayer('Name','iRelu')
    ];

% Create the layer graph.

generator = layerGraph(layers);

%
% Create layers of the downsampling subnetwork. Specify the filter size for the downsampling layers. Use four downsampling layers.

filterSize = [3 3];
numDownSamplingLayers = 4;

for idx = 1:numDownSamplingLayers
    % Compute the number of filters in the next convolutional layer
    multiplier = 2^(idx - 1);
    numFilters = numFiltersGenerator*multiplier*2;
    
    layers = [ ...
        reflectionPad2dLayer(1,strcat("dPad",int2str(idx))) ...
        convolution2dLayer(filterSize,numFilters,"Name",strcat("dConv",int2str(idx)), ...
            "BiasLearnRateFactor",0,"WeightsInitializer",'he',"Stride",2)
        batchNormalizationLayer('Name',strcat("dBn",int2str(idx)), ...
            'OffsetLearnRateFactor',0,'ScaleLearnRateFactor',0) ...
        reluLayer('Name',strcat('dRelu',int2str(idx)))
    ];
    
    % Add the layers to the layer graph
    lg = addLayers(generator,layers);
    generator = connectLayers(lg,generator.Layers(end).Name,layers(1).Name);       
end

%
% Create layers of the residual subnetwork. Specify nine residual blocks in the generator.

numFiltersResidual = numFilters;
numResidualBlocks = 9;
for idx = 1:numResidualBlocks
    % Name of the layer which will act as the source of the residual connection
    res = generator.Layers(end).Name;

    % Specify the layer names
    s = int2str(idx);
    convLayer1Name = strcat("rConv",s,"_1");
    convLayer2Name = strcat("rConv",s,"_2");
    bnLayer1Name = strcat("rBn",s,"_1");
    bnLayer2Name = strcat("rBn",s,"_2");
    pad1Name = strcat("rPad",s,"_1");
    pad2Name = strcat("rPad",s,"_2");
    
    layers = [...
        reflectionPad2dLayer(1,pad1Name) ...
        convolution2dLayer(filterSize,numFiltersResidual,"Name",convLayer1Name, ...
            "BiasLearnRateFactor",0,"WeightsInitializer",'he') ...
        batchNormalizationLayer('Name',bnLayer1Name, ...
            'OffsetLearnRateFactor',0,'ScaleLearnRateFactor',0) ...
        reflectionPad2dLayer(1,pad2Name) ...
        convolution2dLayer(filterSize,numFiltersResidual,"Name",convLayer2Name, ...
            "BiasLearnRateFactor",0,"WeightsInitializer",'he') ...
        batchNormalizationLayer('Name',bnLayer2Name, ...
            'OffsetLearnRateFactor',0,'ScaleLearnRateFactor',0) ...
        reluLayer("Name",strcat("rRelu",s))...
        additionLayer(2,'Name',strcat("rAdd",s))
        ];
    
    % Add the layers to the layer graph
    lg = addLayers(generator,layers);
    generator = connectLayers(lg,generator.Layers(end).Name,layers(1).Name);
    
    % Link the residual connection.
    generator = connectLayers(generator,res,strcat("rAdd",s,"/in2"));
end

%
% Create layers of the upsampling subnetwork.

for idx=1:numDownSamplingLayers
    s = int2str(idx);
    % Compute the number of filters in the next convolutional layer
    multiplier = 2^(numDownSamplingLayers + 1 - idx);
    numFilters = (numFiltersGenerator * multiplier) / 2;
    
    layers = [ ...
        transposedConv2dLayer(filterSize,numFilters,"Name",strcat("uConv",s), ...
            "Stride",2,"BiasLearnRateFactor",0,"WeightsInitializer",'he',"Cropping","Same") ...
        batchNormalizationLayer('Name',strcat("uBn",s), ...
            'OffsetLearnRateFactor',0,'ScaleLearnRateFactor',0) ...
        reluLayer('Name',strcat('uRelu',s));
    ];

    % Add the layers to the layer graph
    lg = addLayers(generator,layers);
    generator = connectLayers(lg,generator.Layers(end).Name,layers(1).Name);
end

%
% Create the final subnetwork layers. Specify the filter size and number of filters of the final convolutional layer of the generator.

filterSize = [7,7];
numFilters = 3;

layers = [ ...
    reflectionPad2dLayer(3,'fPad') ...
    convolution2dLayer(filterSize,numFilters,'Name','fConv', ...
        "WeightsInitializer",'he',"BiasInitializer",'narrow-normal') ...
    tanhLayer('Name','ftanh')];

% Add the subnetwork layers to the layer graph.

lg = addLayers(generator,layers);
lgraphGenerator = connectLayers(lg,generator.Layers(end).Name,layers(1).Name);

% To train the network with a custom training loop and to enable automatic differentiation, convert the layer graph to a dlnetwork object.

dlnetGenerator = dlnetwork(lgraphGenerator);

% Visualize the generator network in a plot.

plot(lgraphGenerator)
title("Generator")

%
%% Create Discriminator Network
% Define a discriminator network that classifies an input image as either real (1) or fake (0).
% The input to the discriminator is the depth-wise concatenation of the one-hot encoded segmentation maps and the scene image to be classified. Specify the number of channels input to the discriminator as the total number of labeled classes and image color channels.

numImageChannels = 3;
numChannelsDiscriminator = numClasses + numImageChannels;
inputSize = [imageSize numChannelsDiscriminator];

% Specify the number of filters in the first convolutional layer of the discriminator.

numFilters = 64;

% Specify the filter size of the first convolutional layer of the discriminator.

filterSize  = [4 4];

% Define the discriminator layers.

discriminator = [
    imageInputLayer(inputSize,"Name","inputLayer","Normalization","none")
    convolution2dLayer(filterSize,numFilters, ...
        "Name","iConv","BiasInitializer","narrow-normal", ...
        "Padding",2,"Stride",2,"WeightsInitializer","he")
    leakyReluLayer(0.2,"Name","lrelu1")
    convolution2dLayer(filterSize,numFilters*2,"Name","dConv1", ...
        "BiasInitializer","narrow-normal", ...
        "Padding",2,"Stride",2,"WeightsInitializer","he")
    batchNormalizationLayer("Name","dBn1", ...
        "OffsetLearnRateFactor",0,"ScaleLearnRateFactor",0)
    leakyReluLayer(0.2,"Name","lrelu2")
    convolution2dLayer(filterSize,numFilters*4,"Name","dConv2", ...
        "BiasInitializer","narrow-normal", ...
        "Padding",2,"Stride",2,"WeightsInitializer","he");
    batchNormalizationLayer("Name","dBn2", ...
        "OffsetLearnRateFactor",0,"ScaleLearnRateFactor",0)
    leakyReluLayer(0.2,"Name","lrelu3")
    convolution2dLayer(filterSize,numFilters*8,"Name","dConv3", ...
        "BiasInitializer","narrow-normal", ...
        "Padding",2,"WeightsInitializer","he")
    batchNormalizationLayer("Name","dBn3", ...
        "OffsetLearnRateFactor",0,"ScaleLearnRateFactor",0)
    leakyReluLayer(0.2,"Name","lrelu4")
    convolution2dLayer(filterSize,1,"Name","fConv", ...
        "BiasInitializer","narrow-normal", ...
        "Padding",2,"WeightsInitializer","he")
    ];

% Create the layer graph.

lgraphDiscriminator = layerGraph(discriminator);

% To train the network with a custom training loop and to enable automatic differentiation, convert the layer graph to a dlnetwork object.

dlnetDiscriminator = dlnetwork(lgraphDiscriminator);

%
% Visualize the discriminator network in a plot.

plot(lgraphDiscriminator)
title("Discriminator")

%
%% Define Model Gradients and Loss Functions
% The helper function modelGradients calculates the gradients and adversarial loss for the generator and discriminator. The function also calculates the feature matching loss and VGG loss for the generator. This function is defined in Supporting Functions section of this example.
%
%% Generator Loss
% The objective of the generator is to generate images that the discriminator classifies as real (1). The generator loss consists of three losses.
% The adversarial loss is computed as the squared difference between a vector of ones and the discriminator predictions on the generated image. \hat{Y}_{generated} are discriminator predictions on the image generated by the generator. This loss is implemented using part of the pix2pixhdAdversarialLoss helper function defined in the Supporting Functions section of this example.
%                            $lossAdversarialGenerator = (1 - \hat{Y}_{generated})^2
% The feature matching loss penalises the L^1 distance between the real and generated feature maps obtained as predictions from the discriminator network. T is total number of discriminator feature layers. Y_{real} and \hat{Y}_{generated} are the ground truth images and generated images, respectively. This loss is implemented using the pix2pixhdFeatureMatchingLoss helper function defined in the Supporting Functions section of this example
%                             $$lossFeatureMatching = \sum_{i=1}^{T} ||Y_{real}  -  \hat{Y}_{generated}||_1
% The perceptual loss penalises the L^1 distance between real and generated feature maps obtained as predictions from a feature extraction network. T is total number of feature layers. Y_{VggReal} and \hat{Y}_{VggGenerated} are network predictions for ground truth images and generated images, respectively. This loss is implemented using the pix2pixhdVggLoss helper function defined in the Supporting Functions section of this example. The feature extraction network is created in Load Feature Extraction Network.
%                             $$lossVgg = \sum_{i=1}^{T} ||Y_{VggReal}  -  \hat{Y}_{VggGenerated}||_1
% The overall generator loss is a weighted sum of all three losses. \lambda_1, \lambda_2, and \lambda_3 are the weight factors for adversarial loss, feature matching loss, and perceptual loss, respectively.
%              lossGenerator = \lambda_1* lossAdversarialGenerator + \lambda_2 * lossFeatureMatching
% + \lambda_3 * lossPerceptual
% Note that the adversarial loss and feature matching loss for the generator are computed for two different scales.
%
%% Discriminator Loss
% The objective of the discriminator is to correctly distinguish between ground truth images and generated images. The discriminator loss is a sum of two components:
% The squared difference between a vector of ones and the predictions of the discriminator on real images
% The squared difference between a vector of zeros and the predictions of the discriminator on generated images
%           lossDiscriminator = (1 - Y_{real})^2 + (0 - \hat{Y}_{generated})^2
% The discriminator loss is implemented using part of the pix2pixhdAdversarialLoss helper function defined in the Supporting Functions section of this example. Note that adversarial loss for the discriminator is computed for two different scales.
%
%% Load Feature Extraction Network
% This example modifies a pretrained VGG-19 deep neural network to extract the features of the real and generated images at various layers. These multilayer features are used to compute the perceptual loss of the generator.
% To get a pretrained VGG-19 network, install Deep Learning Toolbox™ Model for VGG-19 Network. If you do not have the required support packages installed, then the software provides a download link.

net = vgg19;

% Use the analyzeNetwork function to display an interactive visualization of the network architecture.

analyzeNetwork(net)

% To make the VGG-19 network suitable for feature extraction, keep the layers up to 'pool5' and remove all of the fully connected layers from the network. The resulting network is a fully convolutional network.

netLg = layerGraph(net.Layers(1:38));

% Create a new image input layer with no normalization. Replace the original image input layer with the new layer.

inp = imageInputLayer([imageSize 3],"Normalization","None","Name","Input");
netLg = replaceLayer(netLg,"input",inp);
netVGG = dlnetwork(netLg);

%
%% Specify Training Options
% Specify the options for Adam optimization. Train for 60 epochs. Specify identical options for the generator and discriminator networks.
% Start with an equal learning rate of 0.0002
% Initialize the trailing average gradient and trailing average gradient-square decay rates with [].
% Use a gradient decay factor of 0.5 and a squared gradient decay factor of 0.999.

numEpochs = 60;
learningRate = 0.0002;
trailingAvgGenerator = [];
trailingAvgSqGenerator = [];
trailingAvgDiscriminator = [];
trailingAvgSqDiscriminator = [];
gradientDecayFactor = 0.5;
squaredGradientDecayFactor = 0.999;

%
%% Train Network
% Train the model by using the trainNetwork function in a custom training loop. Loop over the training data and update the network parameters each iteration. This example uses a learning rate scheduling policy where the learning rate is constant for the first half of the number of epochs and then decays linearly to zero for remaining half of the number of epochs.
% For each iteration:
% Convert both segmentation map and the ground truth image data to dlarray objects with the underlying type single and spatial dimension labels SSCB (spatial, spatial, channel, batch).
% For GPU training, convert the data to gpuArray objects.
% Evaluate the model gradients using the dlfeval function and the modelGradients helper function.
% Update the network parameters using the adamupdate function.
% Update the training progress plot for every iteration and display various computed losses.
% By default, this example downloads a pretrained version of the Pix2PixHD network for the CamVid data set by using the helper function downloadTrainedPix2PixHDNet. The helper function is attached to the example as a supporting file. The pretrained network enables you to run the entire example without waiting for training to complete.
% Train on a GPU if possible. Use of a CUDA-capable NVIDIA™ GPU with compute capability 3.0 or higher is highly recommended for training Pix2PixHD network (requires Parallel Computing Toolbox™). Training takes about five and a half days on a 12 GB NVIDIA™ Titan X and can take even longer depending on your GPU hardware. Training the network requires at least 14 MB of memory. If you want to train the network with less memory, try reducing the size of the input images to 480-by-640 pixels by specifying the imageSize variable in the Preprocess Training Data section as [480 640] and rerunning the example.

doTraining = false;
if doTraining
    fig = figure;
    lossPlotter = configureTrainingProgressPlotter(fig);
    iteration = 0;

    % Loop over epochs.
    for epoch = 1:numEpochs
        
        % Reset and shuffle the datastore.
        reset(dsTrain);
        dsTrain = shuffle(dsTrain);
        
        % Loop over each image.
        while hasdata(dsTrain)
            iteration = iteration + 1;
            
            % Read data.
            data = read(dsTrain);
            segMap = data{1,1};
            realImage = data{1,2};
            
            % Convert data to dlarray specify the dimension labels 'SSCB'
            % (spatial, spatial, channel, batch).
            dlRealImage = dlarray(realImage,'SSCB');
            dlInputSegMap = dlarray(segMap,'SSCB');
            
            % If training on a GPU, then convert data to gpuArray.
            if canUseGPU
                dlRealImage = gpuArray(dlRealImage);
                dlInputSegMap = gpuArray(dlInputSegMap);
            end
            
            % Evaluate the model gradients and the generator state using
            % dlfeval and the GANLoss function listed at the end of the
            % example.
            [gradParamsG,gradParamsD,losses] = ...
                dlfeval(@modelGradients,dlInputSegMap,dlRealImage,dlnetGenerator,dlnetDiscriminator,vgg19FeatureNet);
            
            % Update the generator parameters.
            [dlnetGenerator,trailingAvgGenerator,trailingAvgSqGenerator] = ...
                adamupdate(dlnetGenerator,gradParamsG, ...
                trailingAvgGenerator,trailingAvgSqGenerator,iteration, ...
                learningRate,gradientDecayFactor,squaredGradientDecayFactor);
            
            % Update the discriminator parameters.
            [dlnetDiscriminator,trailingAvgDiscriminator,trailingAvgSqDiscriminator] = ...
                adamupdate(dlnetDiscriminator,gradParamsD, ...
                trailingAvgDiscriminator,trailingAvgSqDiscriminator, iteration, ...
                learningRate,gradientDecayFactor,squaredGradientDecayFactor);
            
            % Plot and display various losses.
            lossPlotter = updateTrainingProgressPlotter(lossPlotter,iteration,losses);
        end
    end
    save('trainedPix2PixHDNet.mat','dlnetGenerator');
else
    trainedPix2PixHDNet_url = 'https://www.mathworks.com/supportfiles/vision/data/trainedPix2PixHDNet.mat';
    netDir = fullfile(tempdir,'CamVid');
    downloadTrainedPix2PixHDNet(trainedPix2PixHDNet_url,netDir);
    load(fullfile(netDir,'trainedPix2PixHDNet.mat'));
end

%
%% Generate Images from Test Data
% Get a ground truth scene image from the test data. Resize the image.

gtImage = read(imdsTest);
gtImage = imresize(gtImage,imageSize,"nearest");

% Get the corresponding pixel label image from the test data. Resize the pixel label image.

segMap = read(pxdsTest);
segMap = segMap{1};
segMap = imresize(single(segMap),imageSize,"nearest");

% For display, convert the labels from categorical labels to RGB colors by using the label2rgb function.

coloredSegMap = label2rgb(segMap);

% Convert the pixel label image to a multichannel one-hot segmentation map by using the helper function oneHotEncodeSegMap. The helper function is attached to the example as a supporting file.

segMap1Hot = oneHotEncodeSegMap(segMap,numClasses);

% Create a dlarray object to input data to the generator. For GPU inference, convert data to a gpuArray object.

dlSegMap = dlarray(segMap1Hot,'SSCB');

if canUseGPU
    dlSegMap = gpuArray(dlSegMap);
end

% Generate a scene image from the generator and one-hot segmentation map. Use the forward function because the generator network has non-learnable batch normalization layers. Using the predict function yields incorrect results.

[dlGeneratedImage,~] = forward(dlnetGenerator,dlSegMap);
generatedImage = extractdata(gather(dlGeneratedImage));

% The final layer of the generator network is a tanh layer, which produces activations in the range [-1, 1]. For display, rescale the activations to the range [0, 1].

generatedImage = rescale(generatedImage);

%  Display the RGB pixel label image, generated scene image, and ground truth scene image in a montage.

montage({coloredSegMap generatedImage gtImage},'Size',[1 3])

%
%% Generate Images from Custom Pixel Label Images
% To evaluate how well the network generalizes to pixel label images outside the CamVid data set, generate scene images from custom pixel label images. This example uses five pixel label images that were created using the Image Labeler app. The pixel label images are attached to the example as supporting files.
% Create a pixel label datastore to read and process the pixel label images in the current example directory.

cpxds = pixelLabelDatastore(pwd,classes,labelIDs);

% For each test pixel label image, generate and display a scene image.

for idx = 1:length(cpxds.Files)

    segMap = readimage(cpxds,idx);
    segMap = imresize(single(segMap),imageSize,"nearest");
    
    % Convert the pixel label image to a multichannel one-hot image by using
    % the helper function oneHotEncodeSegMap. The helper function is attached
    % to the example as a supporting file.
    segMap1Hot = oneHotEncodeSegMap(segMap,numClasses);
    
    % Convert data to dlarray specify the dimension labels 'SSCB'
    % (spatial, spatial, channel, batch).
    dlSegMap = dlarray(segMap1Hot,'SSCB');
    
    % If training on a GPU, then convert data to gpuArray.
    if canUseGPU
        dlSegMap = gpuArray(dlSegMap);
    end
    
    % Generate the image.
    [generated,~] = forward(dlnetGenerator,dlSegMap);
    generated = extractdata(gather(generated));
    
    % Since the final layer in the generator is a tanh layer, rescale
    % the image to the range [0, 1].
    generated = rescale(generated);
    
    % Display the pixel label image and generated scene image in a montage.
    figure
    montage({label2rgb(segMap) generated})
    title(['Pixel Label Image ',num2str(idx),' and Generated Scene Image'])

end

%
%% Supporting Functions
%% Model Gradients Function
% The modelGradients helper function calculates the gradients and adversarial loss for the generator and discriminator. The function also calculates the feature matching loss and VGG loss for the generator.

function [gradParamsG,gradParamsD,lossGGAN,lossGFM,lossGVGG,lossD] = modelGradients(inputSegMap,realImage,generator,discriminator,netVGG,lossPlot,iterNum)
  
    % Compute the image generated by the generator given the input semantic map
    generatedImage = forward(generator,inputSegMap);
    
    % Define the loss weights.
    lambdaDiscriminator = 1;
    lambdaGenerator = 1;
    lambdaFeatureMatching = 5;
    lambdaVGG = 5;
    
    % Concatenate the image to be classified and the semantic map.
    inpDiscriminatorReal = cat(3,inputSegMap,realImage);
    inpDiscriminatorGenerated = cat(3,inputSegMap,generatedImage);
    
    % Compute the adversarial loss for the discriminator and the generator.
    [DLossScale1,GLossScale1,realPredScale1D,fakePredScale1G] = pix2pixHDAdverserialLoss(inpDiscriminatorReal,inpDiscriminatorGenerated,discriminator);
            
    % Scale the generated image, the real image, and the input semantic map to
    % half size using average pooling.
    resizedRealImage = avgpool(realImage,[3 3],"Stride",[2 2],"Padding",[1 1]);
    resizedGeneratedImage = avgpool(generatedImage,[3 3],"Stride",[2 2],"Padding",[1 1]);
    resizedinputSegMap = avgpool(inputSegMap,[3 3],"Stride",[2 2],"Padding",[1 1]);
    
    % Concatenate the image to be classified and the semantic map.
    inpDiscriminatorReal = cat(3,resizedinputSegMap,resizedRealImage);
    inpDiscriminatorGenerated = cat(3,resizedinputSegMap,resizedGeneratedImage);
    
    % Compute the adversarial loss for the discriminator and the generator
    [DLossScale2,GLossScale2,realPredScale2D,fakePredScale2G] = pix2pixHDAdverserialLoss(inpDiscriminatorReal,inpDiscriminatorGenerated,discriminator);
    
    % Compute the feature matching loss for scale 1
    FMLossScale1 = pix2pixHDFeatureMatchingLoss(realPredScale1D,fakePredScale1G);
    FMLossScale1 = FMLossScale1 * lambdaFeatureMatching;
    
    % Compute the feature matching loss for scale 2
    FMLossScale2 = pix2pixHDFeatureMatchingLoss(realPredScale2D,fakePredScale2G);
    FMLossScale2 = FMLossScale2 * lambdaFeatureMatching;
    
    % Compute the VGG loss.
    VGGLoss = pix2pixHDVGGLoss(realImage,generatedImage,netVGG);
    VGGLoss = VGGLoss * lambdaVGG;
    
    % Compute the combined generator loss
    lossGCombined = GLossScale1 + GLossScale2 + FMLossScale1 + FMLossScale2 + VGGLoss;
    lossGCombined = lossGCombined * lambdaGenerator;
    
    % Compute gradients for the generator
    gradParamsG = dlgradient(lossGCombined,generator.Learnables);
    
    % Compute the combined discriminator loss
    lossDCombined = (DLossScale1 + DLossScale2)/2 * lambdaDiscriminator;
    
    % Compute gradients for the discriminator
    gradParamsD = dlgradient(lossDCombined,discriminator.Learnables);
    
    % Log the values for displaying later
    lossD = gather(extractdata(lossDCombined));
    lossGGAN = gather(extractdata(GLossScale1 + GLossScale2));
    lossGFM  = gather(extractdata(FMLossScale1 + FMLossScale2));
    lossGVGG = gather(extractdata(VGGLoss));
    
    % Plot the losses
    addpoints(lossPlot.dline,iterNum,double(lossD));
    addpoints(lossPlot.gline,iterNum,double(lossGGAN));
    addpoints(lossPlot.fmline,iterNum,double(lossGFM));
    addpoints(lossPlot.vggline,iterNum,double(lossGVGG));
    drawnow

end

%% Adversarial Loss Function
% The helper function pix2pixHDAdverserialLoss computes the adversarial loss gradients for the generator and the discriminator. The function also returns feature maps of the real image and synthetic images.

function [DLoss,GLoss,realPredFtrsD,genPredFtrsG] = pix2pixHDAdverserialLoss(inpReal,inpGenerated,discriminator)

    % Discriminator layer names containing feature maps
    featureNames = {'lrelu1','lrelu2','lrelu3','lrelu4','fConv'};
    
    % Get the feature maps for the real image from the discriminator    
    realPredFtrsD = cell(size(featureNames));
    [realPredFtrsD{:}] = forward(discriminator,inpReal,"Outputs",featureNames);
    
    % Get the feature maps for the generated image from the discriminator    
    genPredFtrsD = cell(size(featureNames));
    [genPredFtrsD{:}] = forward(discriminator,inpGenerated,"Outputs",featureNames);
    
    % Get the feature map from the final layer to compute the loss
    realPredD = realPredFtrsD{end};
    genPredD = genPredFtrsD{end};
    
    % Compute the discriminator loss
    DLoss = (1 - realPredD).^2 + (genPredD).^2;
    DLoss = mean(DLoss,"all");
    
    % Compute the generator loss
    genPredFtrsG = cell(size(featureNames));
    [genPredFtrsG{:}] = forward(discriminator,inpGenerated,"Outputs",featureNames);
    genPredG = genPredFtrsG{end};
    GLoss = (1 - genPredG).^2;
    GLoss = mean(GLoss,"all");
end

%% Feature Matching Loss Function
% The helper function pix2pixHDFeatureMatchingLoss computes the feature matching loss between a real image and a synthetic image generated by the generator.

function featureMatchingLoss = pix2pixHDFeatureMatchingLoss(realPredFtrs,genPredFtrs)

    % Number of features
    numFtrsMaps = numel(realPredFtrs);
    
    % Initialize the feature matching loss
    featureMatchingLoss = 0;
    
    for i = 1:numFtrsMaps
        % Get the feature maps of the real image
        a = extractdata(realPredFtrs{i});
        % Get the feature maps of the synthetic image
        b = genPredFtrs{i};
        
        % Compute the feature matching loss
        featureMatchingLoss = featureMatchingLoss + mean(abs(a - b),"all");
    end
end

%% Perceptual VGG Loss Function
% The helper function pix2pixHDVGGLoss computes the perceptual VGG loss between a real image and a synthetic image generated by the generator.

function vggLoss = pix2pixHDVGGLoss(realImage,generatedImage,netVGG)

    featureWeights = [1.0/32 1.0/16 1.0/8 1.0/4 1.0];
    
    % Initialize the VGG loss
    vggLoss = 0;
    
    % Specify the names of the layers with desired feature maps
    featureNames = ["relu1_1","relu2_1","relu3_1","relu4_1","relu5_1"];
    
    % Extract the feature maps for the real image
    activReal = cell(size(featureNames));
    [activReal{:}] = forward(netVGG,realImage,"Outputs",featureNames);
    
    % Extract the feature maps for the synthetic image
    activGenerated = cell(size(featureNames));
    [activGenerated{:}] = forward(netVGG,generatedImage,"Outputs",featureNames);
    
    % Compute the VGG loss
    for i = 1:numel(featureNames)
        vggLoss = vggLoss + featureWeights(i)*mean(abs(activReal{i} - activGenerated{i}),"all");
    end
end

%% References
% [1] Wang, Ting-Chun, Ming-Yu Liu, Jun-Yan Zhu, Andrew Tao, Jan Kautz, and Bryan Catanzaro. "High-Resolution Image Synthesis and Semantic Manipulation with Conditional GANs." In 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, 8798–8807, 2018. https://doi.org/10.1109/CVPR.2018.00917.
% [2] Brostow, Gabriel J., Julien Fauqueur, and Roberto Cipolla. "Semantic Object Classes in Video: A High-Definition Ground Truth Database." Pattern Recognition Letters. Vol. 30, Issue 2, 2009, pp 88-97.
% [3] Isola, Phillip, Jun-Yan Zhu, Tinghui Zhou, and Alexei A. Efros. "Image-to-Image Translation with Conditional Adversarial Networks." In 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 5967–76, 2017. https://doi.org/10.1109/CVPR.2017.632.
% [4] Johnson, Justin, Alexandre Alahi, and Fei-Fei Li. "Perceptual Losses for Real-Time Style Transfer and Super-Resolution." In 2016 European Conference on Computer Vision (ECCV), 694–711, 2016.
% Copyright 2019 The MathWorks, Inc.