%% Transform Volumetric Data in Random Patch Extraction Datastore
% This example shows how to transform pairs of volumetric data in a random patch extraction datastore using a sample image preprocessing pipeline.
% Specify two sets of volumetric images saved at MAT files. Each set contains five volumetric images.

dir = fullfile(matlabroot,"toolbox","images","imdata","BrainMRILabeled");
filesVol1 = fullfile(dir,"images");
filesVol2 = fullfile(dir,"labels");

% Store each set of volumetric images in an image datastore. Specify a custom read function, matRead. This function is defined in the Supporting Functions section of this example. Use the default ReadSize of 1.

vol1DS = imageDatastore(filesVol1,"FileExtensions",".mat","ReadFcn",@(x) matRead(x));
vol2DS = imageDatastore(filesVol2,"FileExtensions",".mat","ReadFcn",@(x) matRead(x));

% Specify the input size of the network.

inputSize = [128 128];

% Create a random patch extraction datastore that extracts corresponding patches from the two datastores. Select three patches per image.

patchVolDS = randomPatchExtractionDatastore(vol1DS,vol2DS,inputSize,"PatchesPerImage",3);

% Preprocess the volumetric images in patchVolDS using the custom preprocessing pipeline defined in the preprocessVolumetricPatchDS supporting function.

dsTrain = transform(patchVolDS,@(x) preprocessVolumetricPatchDS(x));

% Read a batch of data.

minibatch = read(dsTrain)

%
%% Supporting Functions
% The matRead function loads volume data from the first variable of a MAT file.

function data = matRead(filename)
    inp = load(filename);
    f = fields(inp);
    data = inp.(f{1});
end

% The preprocessVolumetricPatchDS function performs the desired transformations of data read from the underlying random patch extraction datastore. The function must accept a table. The function transforms the data according to this preprocessing pipeline:
% Randomly select one of five augmentations.
% Apply the same augmentation to the data in both columns of the table.
% Return the augmented image pair in a table.

function batchOut = preprocessVolumetricPatchDS(batchIn)

numRows = size(batchIn,1); 
batchOut = batchIn;

% 5 augmentations: nil,rot90,fliplr,flipud,rot90(fliplr)
augType = {@(x) x,@rot90,@fliplr,@flipud,@(x) rot90(fliplr(x))};

for idx = 1:numRows

    img = batchIn{idx,1}{1};
    resp = batchIn{idx,2}{1};
    
    rndIdx = randi(5,1);
    imgAug = augType{rndIdx}(img);
    respAug = augType{rndIdx}(resp);
    
    batchOut(idx,:) = {imgAug,respAug};

end
end

% Copyright 2020 The MathWorks, Inc.