classdef Mixup < handle
    %MIXUP Mixup data augmentation
    %   mixer = Mixup(augDatastore) creates an object that can mix features
    %   at a randomly set ratio and then probabilistically set the output
    %   label as one of the two original signals.
    %
    %   Mixup Properties:
    %   MixProbability - Mix probability
    %   AugDatastore   - Augmentation datastore
    %
    %   Mixup Methods:
    %   mix            - Apply mixup
    %

    % Copyright 2021 The MathWorks, Inc.

    properties (SetAccess=public,GetAccess=public)
        %MixProbability Mix probability
        % Specify the probability that mixing is applied as a scalar in the
        % range [0,1]. If unspecified, MixProbability defaults to 1/3.
        MixProbability (1,1) {mustBeNumeric} = 1/3;
    end
    properties (SetAccess=immutable,GetAccess=public)
        %AUGDATASTORE Augmentation datastore
        % Specify a datastore from which to get the mixing signals. The
        % datastore must contain a label in the info returned from reading.
        % This property is immutable, meaning it cannot be changed after
        % construction.
        AugDatastore
    end

    methods
        function obj = Mixup(augDatastore)
            obj.AugDatastore = augDatastore;
        end

        function [dataOut,infoOut] = mix(obj,x,infoIn)
            %MIX Apply mixup
            % [dataOut,infoOut] = mix(mixer,x,infoIn) probabilistically mix
            % the input, x, and its associated label contained in infoIn
            % with a signal randomly drawn from the augmentation datastore.
            % The output, dataOut, is a cell array with two columns. The
            % first column contains the features and the second column
            % contains the label.

            if rand > obj.MixProbability % Only mix ~1/3 the dataset

                % Randomly set mixing coefficient. Draw from a normal
                % distribution with mean 0.5 and contained within [0,1].
                lambda = max(min((randn./10)+0.5,1),0);

                % Read one file from the augmentation datastore.
                subDS = subset(obj.AugDatastore,randi([1,numel(obj.AugDatastore.UnderlyingDatastores{1}.Files)]));
                [y,yInfo] = read(subDS);

                % Mix the features element-by-element according to lambda.
                dataOut = lambda*x + (1-lambda)*y;

                % Set the output label probabilistically based on the mixing coefficient.
                if lambda < rand
                    labelOut = yInfo.Label;
                    infoOut.Label = labelOut;
                else
                    labelOut = infoIn.Label;
                end
                infoOut.Label = labelOut;

                % Combine the output data and labels.
                dataOut = [{dataOut},{labelOut}];

            else % Do not apply mixing

                dataOut = [{x},{infoIn.Label}];
                infoOut = infoIn;

            end
        end

    end
end