%% Classify Large Multiresolution Images Using blockedImage and Deep Learning
%
% TODO:
% uncomment showlabels when available.
% spell check
% link to inference example at the end as the next step.
%
% This example shows how to train a Inception-v3 deep neural network using multiresolution whole slide images (WSIs), that do not fit in memory.
% Deep learning methods for tumor classification rely on digital pathology, in which whole tissue slides are imaged and digitized. The resulting WSIs have high resolution (~200k-by-100k). WSIs are frequently stored in a multiresolution format to facilitate efficient display, navigation, and processing of images.
% WSIs are best processed using out-of-core memory image processing techniques. blockedImage and blockedImageDatastore objects help access, process and store these types of large multiresolution images.
% This example shows how to use lower resolution image data to efficiently prepare training data from the finer levels. It then outlines an architecture to use block based processing to train large WSIs using blockedImage and blockedImageDatastore. The main focus of this example is in setting up the training infrastructure. An Inception-v3 based network is trained using transfer learning techniques to classify individual blocks as 'Normal' or 'Tumor'. The generated tumor heatmaps are used for classification.
%% Download Data
% This example uses WSIs from the Camelyon16 challenge [1]. The data from this challenge contains a total of 400 WSIs of lymph nodes from two independent sources, separated into 270 training images and 130 test images. The WSIs are stored as TIF files in a stripped format with an 11-level pyramid structure. The size of each image file is approximately 2 GB. The training data set consists of 159 WSIs of normal lymph nodes and 111 WSIs of lymph nodes with tumor and healthy tissue. Ground truth coordinates of the lesion boundaries accompany the tumor images.
% Create a root directory to store all the data. Training data takes ~451 GBs.

dataDir ='/mathworks/inside/files/dev/spc/ipt/Camelyon16/';

%
% To download the training data, go to the Camelyon17 website and click the first "CAMELYON16 data set" link.
% Open the "training" directory, then follow these steps:
% Download the "lesion_annotations.zip" file. Extract the files to the directory specified by the trainAnnotationDir variable.
% Open the "normal" directory. Download the images to the directory specified by the trainNormalDir variable.
% Open the "tumor" directory. Download the images to the directory specified by the trainTumorDir variable.

trainingDir = fullfile(dataDir, 'Train');
trainNormalDir = fullfile(trainingDir,'normal');
trainTumorDir = fullfile(trainingDir,'tumor');
trainAnnotationDir = fullfile(trainingDir,'lesion_annotations');
if ~exist(trainingDir,'dir')
    mkdir(trainingDir)
    mkdir(trainNormalDir)
    mkdir(trainTumorDir)
    mkdir(trainAnnotationDir)
end

% Final folder structure:
%
%
% Create a location for intermediate files. This location is used to store masks and the trained network. This location will need ~2 GBs of free space.

workingDir = '/tmp/Working_Data/';
%workingDir = '/ipgeodata/Deep_Learning_Examples_Data/Camelyon16/Working_Data'

%
%% Create blockedImage Objects For Training
% Create a file list of normal and tumor files

normalFileSet = matlab.io.datastore.FileSet([trainNormalDir,filesep,'normal*']);
normalDataFiles = normalFileSet.FileInfo.Filename;
tumorFileSet = matlab.io.datastore.FileSet([trainTumorDir,filesep,'tumor*']);
tumorDataFiles = tumorFileSet.FileInfo.Filename;

%
% Note that one of the training images of normal tissue, normal_144.tif is missing important metadata required to deduce physical extents of the image. Exclude this file.

normalDataFiles(contains(normalDataFiles,'normal_144')) = [];

%
% Create corresponding arrays of blockedImage objects. This loada image metadata from each file and creates a blockedImage object pointing to the image data on disk. This process takes ~2 mins [4].

normalImages = blockedImage(normalDataFiles);
tumorImages = blockedImage(tumorDataFiles);

%
%% Visually Inspect WSI Images
% To get a better understanding of the training data, pick an image at random to visualize.

tumorImage = tumorImages(10);

%
% Display the blockedImage using the bigimageshow function. bigimageshow automatically picks an appropriate resolution based on the screen size and current zoom level. Note that this image contains a lot of empty white space. The tissue occupies only a small portion of the image - a typical characteristic of WSIs.

h = bigimageshow(tumorImage);
title("Resolution Level: " + h.ResolutionLevel)

%
% Zoom into an area of interest and note that the resolution level automatically increases to show more detail:

xlim([33600, 35000])
ylim([104600, 106000])
title("Resolution Level: " + h.ResolutionLevel)

%
%% Set Correct Spatial Extents
% Inspect the dimensions of the blockedImage at each resolution level. Level 1 has the most pixels and is the finest resolution level. Level 10 has the fewest pixels and is the coarsest resolution level. Notice that the aspect ratio is not consistent, this usually indicates that the levels do not span the same world area. It is important to ensure that correct spatial extents are set before processing this data.


levelSizeInfo = table((1:length(tumorImage.Size))', ...
    tumorImage.Size(:,1), ...
    tumorImage.Size(:,2), ...
    tumorImage.Size(:,1)./tumorImage.Size(:,2), ...
    'VariableNames',["Resolution Level" "Image Width" "Image Height" "Aspect Ratio"]);
disp(levelSizeInfo)

%
% The metadata in these image files have additional spatial referencing information saved in the ImageDescription TIFF tag. To ensure regions line up across different levels, extract this metadata and update the WorldStart and WorldEnd properties of the blockedImages.
% This is an important step that allows information (e.g. masks) extracted from coarse levels to be correctly used at finer levels.
% setSpatialReferencingForCamelyon16 is a helper function which loads coordinate metadata from the TIFF headers and sets WorldStart and WorldEnd properties.

tumorImages = setSpatialReferencingForCamelyon16(tumorImages);
normalImages = setSpatialReferencingForCamelyon16(normalImages);

%
% See Set Spatial Referencing for Blocked Images for more information.
%% Create Tissue Masks
% Majority of a typical WSI consists of background pixels. Region of Interest (ROI) masks can be efficiently created from coarser levels, and then used to limit computation to corresponding regions at finer levels. For more information, see Process Big Images Efficiently Using Masks.
% Two important factors while picking a mask level are:
% The image size at the chosen mask level. Smaller sizes can be processed faster.
% Fidelity of ROI boundaries at chosen level. Intricate boundary features might need larger mask sizes to capture the boundary accurately. This comes with a cost of increased processing time.
% Use bigimageshow to sample and visualize the training data to determine a suitable mask level. This example uses Level 8.

normalMaskLevel = 8;

% If the chosen level is small enough to fit in memory, use the browseBlockedImage helper function to launch the Image Browser App to visually inspect the blockedImage array.

%browseBlockedImages(normalImages, normalMaskLevel)

%
%
% The background is uniformly light. A simple thresholding function can be used to segment out the tissue. Use the apply function to apply this thresholding in a parallel, block-wise manner on all normal training images. Also specify an OutputLocation to save the result to disk. This process takes ~2 mins [4].

trainNormalMaskDir = fullfile(workingDir,['normal_tissue_mask_at_level' num2str(normalMaskLevel)]);

if ~isfolder(trainNormalMaskDir)
    normalMasks = apply(normalImages, @(bs)imclose(im2gray(bs.Data)<150, ones(5)), ...
        'BlockSize', [512 512],...
        'Level', normalMaskLevel, ...
        'UseParallel', canUseParallelPool,...
        'OutputLocation', trainNormalMaskDir);
    save([trainNormalMaskDir, filesep, 'normalMasks.mat'], 'normalMasks')
else
    % Load previously saved data
    load([trainNormalMaskDir, filesep, 'normalMasks.mat'], 'normalMasks');
end

%
% Inspect generated masks to visually confirm. Note that these masks have only one level. Mask levels are expected to be small enough to fit in memory.

%browseBlockedImages(normalMasks, 1)

%
% Based on this above overview, pick individual images to further assess mask accuracy. Display the original image and layer the mask on top of it.

idx = 42;
hNormal = bigimageshow(normalImages(idx));
%showlabels(hNormal, normalMasks(idx),...
%'AlphaData', normalMasks(idx), 'Alphamap',[0.3 0])
title('Green region is masked OUT')
% Zoom in to inspect area of interest.
xlim([47540 62563])
ylim([140557 155581])

%
%% Create Tumor Masks
% For tumor images, use manually drawn lesion boundaries that accompany the tumor images to create masks. showTumorAnnotations is a helper function which uses readstruct to read polygon coordinates from supplied ground truth XML file and displays them as free hand ROIS. Notice that normal regions (shown in green) can occur inside cancer regions.

showTumorAnnotations(tumorImages(64), trainAnnotationDir);
xlim([77810 83602])
ylim([139971 145763])

%
% These regions are hand drawn at the finest level. Higher resolution masks may be used (compared to normal images) to get better fidelity if desired. Note that finer levels have more pixels and will take longer to process.

tumorMaskLevel = 8;

%
% Steps involved in creating a tumor mask for each image are:
% Initialize a writable blockedImage based on the tumor image size at the required resolution.
% Read lesion annotations from corresponding XML file
% Create bounding boxes for each ROI instance
% Loop through each output block dimensions, if any of the ROI bounding boxes intersect the bounds of the current block, use inpolygon to compute the mask for the block
% Save results to disk for reuse.
% This logic is implemented in the createMaskForTumorTissue helper function.

trainTumorMaskDir = fullfile(workingDir,['tumor_mask_level' num2str(tumorMaskLevel)]);
if ~isfolder(trainTumorMaskDir)
    % Create masks
    tumorMasks = createMaskForTumorTissue(tumorImages,trainAnnotationDir, trainTumorMaskDir,tumorMaskLevel);
    save([trainTumorMaskDir, filesep, 'tumorMasks.mat'], 'tumorMasks')
else
    % Load previously saved data
    load([trainTumorMaskDir, filesep, 'tumorMasks.mat'], 'tumorMasks');
end

%
% Visually inspect all the generated masks.

%browseBlockedImages(tumorMasks, 1)

%
%
% Closely inspect images which have intricate regions to ensure that the generated mask at the chosen level is adequate.

idx = 64;
hTumor = bigimageshow(tumorImages(idx));
%showlabels(hTumor, tumorMasks(idx),'AlphaData', tumorMasks(idx), 'Alphamap',[0.3 0]);
title('Green region is masked OUT')
xlim([77810 83602]), ylim([139971 145763])

%
%% Select Blocks For Training
% The network can be trained using data from any image level. Finer levels offer more uniform class distribution within a chosen block size. Coarser levels, which include more spatial region for the same block size, will have more surrounding context. For this dataset, training at the finest level was found to be most performant.

trainingLevel = 1;

%
% This example uses Inception-v3 which requires the following block size

networkBlockSize = [299,299,3];

%
% Use selectBlockLocations and the generated masks to create a set of Normal and Tumor blocks. This function helps pick blocks which are within the specfied mask area. A subset of all available blocks is chosen balancing two important criterion:
% Using as much of the training data as possible to help generalize the network during training - while ensuring a good class representation balance in the selected block set.
% The total training budget available. This includes hardware and training time available. More blocks will result in more training time.
% Start with conservative numbers resulting in a small set of blocks to initially focus on setting up the full infrastructure. Once training and validation are in-place, the whole pipeline can then be easily repeated with larger number of blocks to study the effect of various tradeoffs in setting BlockOffsets and InclusionThrehold values. These are two parameters available to control the number of blocks picked for each class:
% BlockOffsets, the stride size. This controls the spacing between two successive blocks. For a block size of [M, N], any stride larger than [M, N] results in blocks having spaces between them (undersampling). While strides smaller than [M, N] results in overlapping blocks (oversampling).
% InclusionThreshold, the percentage area of a block that has to be fully inside a mask to be chosen. A larger value will ensure a more homogeneous block. This could potentially help the training converge faster, likely at the expense of lower accuracy along class boundaries. These values can be manually tuned to ensure a balanced sampling.

% Increasing these two values decreases the number of blocks:
tumorOffsetFactor = 1;
tumorInclusionThreshold = 0.90;
blsTumorData = selectBlockLocations(tumorImages, ...
    "BlockSize",networkBlockSize(1:2),"BlockOffsets",round(networkBlockSize(1:2)*tumorOffsetFactor), ...
    "Levels", trainingLevel,...
    "Masks",tumorMasks,"InclusionThreshold", tumorInclusionThreshold,...
    "ExcludeIncompleteBlocks",true,...
    "UseParallel", canUseParallelPool); % ~8 mins[4]
disp(blsTumorData)

normalOffsetFactor = 3.5;
normalInclusionThreshold = .97;
blsNormalData = selectBlockLocations(normalImages, ...
    "BlockSize",networkBlockSize(1:2),"BlockOffsets",round(networkBlockSize(1:2)*normalOffsetFactor), ...
    "Levels", trainingLevel,...
    "Masks",normalMasks,"InclusionThreshold",normalInclusionThreshold,...
    "ExcludeIncompleteBlocks",true,...
    "UseParallel", canUseParallelPool); % ~2 mins[4]
disp(blsNormalData)

%
% Create a single blockedImageDatastore by merging normal and tumor blocks

allImages = [normalImages, tumorImages];
blsAll = blockLocationSet(...
    [blsNormalData.ImageNumber; blsTumorData.ImageNumber + numel(normalImages)],...
    [blsNormalData.BlockOrigin; blsTumorData.BlockOrigin],...
    networkBlockSize, trainingLevel);
dsAllBlocks = blockedImageDatastore(allImages, 'BlockLocationSet', blsAll);

%
% Shuffle the blocks to reduce chances of overfitting and to increase generalization in the training process.

% Reset random stream to make process repeatable.
stream = RandStream.getGlobalStream;
reset(stream);
dsAllBlocks = shuffle(dsAllBlocks);

%
% Split the blocks into two sets, training and validation. Retain 99% for training, and set the remaining 1% as validation. The initial set is already shuffled, so the split sets would have randomized blocks from different images.

numericBlockLabels = [zeros(size(blsNormalData.ImageNumber)); ones(size(blsTumorData.ImageNumber))];
blockLabels = categorical(numericBlockLabels,[0,1],{'normal','tumor'});
idx  = splitlabels(blockLabels, 0.99, 'randomized');
dsTrain = subset(dsAllBlocks, idx{1});
dsVal = subset(dsAllBlocks, idx{2});

%
% The training process needs labelled blocks. A read on a blockedImageDatastore returns a block of data, and additional metadata in an info structure. The ImageNumber field corresponds to the array index of the blockedImage from which the corresponding block was read from. This index can be used to label the block as "Normal" or "Tumor". Use a transform function to perform this look up.

type labelBlocks

numericImageLabels = [zeros(size(normalImages)), ones(size(tumorImages))];
imageLabels = categorical(numericImageLabels,[0,1],{'normal','tumor'});
dsTrainLabelled = transform(dsTrain,...
    @(block, blockInfo) labelBlocks(block, blockInfo, imageLabels), 'IncludeInfo', true);
dsValLabelled = transform(dsVal,...
    @(block, blockInfo) labelBlocks(block, blockInfo, imageLabels), 'IncludeInfo', true);

%
% Augment these blocks with another transform call. This augmentation includes simple geometric transformations (e.g 90 degree rotation, horizontal and vertical flips). Additional augmentation techniques may be added to improve the ability of the network to generalize. For example, a common augmentation technique for WSI images is 'stain normalization'. Stain normalization reduces the variability in color and intensity present in stained images from different sources.

type augmentBlocks

dsTrainLabelled = transform(dsTrainLabelled, @augmentBlocks);

%
%% Set Up Inception-v3 Network Layers For Transfer Learning
% This example uses the Inception-v3 network [2], a convolutional neural network that is trained on more than a million images from the ImageNet database [3]. The network is 48 layers deep and can classify images into 1,000 object categories, such as keyboard, mouse, pencil, and many animals.
% The inceptionv3 function returns a pretrained Inception-v3 network. Inception-v3 requires the Deep Learning Toolbox™ Model for Inception-v3 Network support package. If this support package is not installed, then the function provides a download link.

net = inceptionv3;

%
% The convolutional layers of the network extract image features that the last learnable layer and the final classification layer use to classify the input image. These two layers contain information on how to combine the features that the network extracts into class probabilities, a loss value, and predicted labels. To retrain a pretrained network to classify new images, replace these two layers with new layers adapted to the new data set. For more information, see Train Deep Learning Network to Classify New Images.

% Extract the layer graph from the trained network.
lgraph = layerGraph(net);

%
% Find the names of the two layers to replace by using the supporting function findLayersToReplace. In Inception-v3, these two layers are named 'predictions' and 'ClassificationLayer_predictions'.

[learnableLayer,classLayer] = findLayersToReplace(lgraph);

%
% The goal of this example is to perform binary segmentation between two classes, tumor and normal regions. Create a new fully connected layer for two classes. Replace the original final fully connected layer with the new layer.

numClasses = 2;
newLearnableLayer = fullyConnectedLayer(numClasses,'Name','predictions');
lgraph = replaceLayer(lgraph,learnableLayer.Name,newLearnableLayer);

%
% Create a new classification layer for two classes. Replace the original final classification layer with the new layer.

newClassLayer = classificationLayer('Name','ClassificationLayer_predictions');
lgraph = replaceLayer(lgraph,classLayer.Name,newClassLayer);

%
%% Specify Training Options
% Train the network using the rmsprop optimization solver. This solver automatically adjusts the learning rate and momentum for faster convergence. Specify other hyperparameter settings by using the trainingOptions function. Reduce MaxEpochs to a small number because the large amount of training data enables the network to reach convergence sooner. The MiniBatchSize is constrained by available GPU memory at the upper end. While larger size can make the training faster, larger sizes are known to reduce the ability of the network to generalize. Validation is expensive, so reduce the validation frequency to once ever 250 iterations. This network normalizes input data using rescale-symmetric option, with input range expected to be in between [0 255]. This matches the current data set, hence set ResetInputNormalization to false to prevent a full read of the training data to compute normalization stats. Enable mult-gpu training to spread training workload across multiple GPUs, if present.

checkpointsDir = fullfile(workingDir,'checkpoints');
if ~exist(checkpointsDir,'dir')
    mkdir(checkpointsDir);
end

options = trainingOptions('rmsprop', ...
    'MaxEpochs',1, ...
    'MiniBatchSize',256, ...
    'Shuffle','every-epoch',...
    'ValidationData',dsValLabelled, ...
    'ValidationFrequency', 250,...
    'InitialLearnRate',1e-4, ...
    'SquaredGradientDecayFactor',0.99, ...
    'ResetInputNormalization', false,...
    'ExecutionEnvironment','multi-gpu',...  % Change this based on the training environment.
    'Plots','training-progress', ...
    'CheckpointPath',checkpointsDir);

%
%% Train or Download the Network
% By default, this example downloads a pretrained version of the trained Inception-v3 network using the helper function downloadTrainedCamelyonNet. The pretrained network can be used to run the entire example without waiting for training to complete.
% To train the network, set the doTraining variable in the following code to true. Train the network using the trainNetwork function.
% Train on one or more GPUs, if available. Using a GPU requires Parallel Computing Toolbox™ and a CUDA® enabled NVIDIA® GPU. For more information, see GPU Support by Release. Training takes about 250 mins [4].

% Change this to true to perform the full training
doTraining = false; 
if doTraining
    trainedNet = trainNetwork(dsTrainLabelled,lgraph,options);
    modelDateTime = string(datetime('now','Format',"yyyy-MM-dd-HH-mm-ss"));    
    save(workingDir + "trainedCamelyonNet.mat",'trainedNet');
    % Save a time-stamped backup.
    save(workingDir + "trainedCamelyonNet-" + modelDateTime + ".mat",'trainedNet');
end

% Training plot:
%
% To continue with using this network and quantifying its performance, see 'Inference and Quantification Of a Multiresolution Deep Learning Network"
%
%
%% References
% [1] Ehteshami B. B., et al. "Diagnostic Assessment of Deep Learning Algorithms for Detection of Lymph Node Metastases in Women With Breast Cancer." Journal of the American Medical Association. Vol. 318, No. 22, 2017, pp. 2199–2210. doi:10.1001/jama.2017.14585
% [2] Szegedy, C., V. Vanhoucke, S. Ioffe, J. Shlens, and Z. Wojna. "Rethinking the Inception Architecture for Computer Vision." In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2818–2826. Las Vegas, NV: IEEE, 2016.
% [3] ImageNet. http://www.image-net.org.
% [4] All reported times were measured on a PC running Debian GNU/Linux 10, with dual socket Intel® Xeon® Gold 5222 CPUs, 376 GB RAM and 7 TITAN RTX GPUs. Original data was stored on an NFS drive. Intermediate data was stored on a local SSD.
% Copyright 2020-2021 The MathWorks, Inc.