% This code takes a structure of images and their corresponding masks and
% splits them in individual images corresponding to objects in the original
% image. This is done applying the watershed algorithm.

% This code uses functions found in Image Processing Toolbox and the 
% HyperTools toolbox (https://www.hypertools.org/).

% image = structure. Size = number of images.
% images(i).data = cube (height x width x spectrum) containing the data of 
% a hyperspectral image. No axes or labels.
% image(i).mask = logical array (height x width in pixels) stating if each
% pixel is included in the analysis. 1=Included 0=Not included

% Author: Jokin Ezenarro Garate. 17/02/2025
% Citation and illustrative example: https://doi.org/10.1016/j.saa.2025.125906

images= ; %Select your dataset

MinItemSize = 50; % Size (in pixels) of the smallest thing considered as object.
for i=1:size(images)  %Loop through images
    close all
    imageCut=split(images(i).data,images(i).mask,MinItemSize);
    ask=1;
    while ask==1
        close all
        x = input("1=Continue 0=Repeat \n");
        if x==0
            imageCut=split(images(i).data,images(i).mask,30);
        else
            ask=0;
        end
    end
    imCut{i}=imageCut;
end
disp("You're done!")

%%
function imageCut=split(data,mask,MinItemSize)
    %bw = mask ~= 0; The predefined mask could be used but it's not so
    %sensitive to gradients in the spectra
    A = htfakergb(data,mask);    
    A1=1-((squeeze(double(A(:,:,2)))-30)./256); %Using one of the channels 
    % is more sensitive to empty spaces and divides objects better. The
    % number 30 is an adjustment to this sensitivity, can be changed.
    bw=round(A1,0);

    bw = imerode(bw,strel("disk",1)); %Remove small objects
    bw = bwareaopen(~bw,20);
    bw = bwareaopen(~bw,200);
    bw = imerode(bw,strel("disk",1));
    
    D = bwdist(~bw);
    D = -D;
    %%
    S = watershed(D,26); %Do the object division
    S(~bw) = 0;
    %imagesc(S);
    
    Mitems=max(max(S));

    %% Divide objects in masks
    clear a b m
    items=0;
    for i=1:Mitems
        a=find(S==i);
        if size(a,1)>MinItemSize
            items=items+1;
            b=zeros(size(mask,1),size(mask,2));
            b(a)=1;
            m(items,:,:)=b;      
        end
    end
    %% Plot objects masks on original image
    A = htfakergb(data,mask);
    figure,imagesc(A);
    hold on
    for i=1:items
        n=squeeze(m(i,:,:));
        f=imshow(n);
        set(f, 'AlphaData', n*0.3);
        text(mean(find(sum(n,1)>0)),mean(find(sum(n,2)>0)),string(i))
    end
    %% Plot each of the new objects
    f=figure;
    for i=1:items
        n=squeeze(m(i,:,:));
        ny=sum(n,2);
        ly=find(ny>0);
        nx=sum(n,1);
        lx=find(nx>0);
        imageCut(i).data=data(ly(1):ly(end),lx(1):lx(end),:);
        imageCut(i).mask=logical(n(ly(1):ly(end),lx(1):lx(end),:));
    
        nexttile
        v = htfakergb(imageCut(i).data,imageCut(i).mask);
        imagesc(v);
        title(i)
    end
    clear n ny nx ly lx v

    %% Check if the objects are alright or should be modified
    ask=1;
    while ask==1 
        clear x
        x = input("Remove or combine any objects? \n 0=No \n i=Remove object i \n [i1,i2]=Combine those objects \n");
        if size(x,2)<2
            if x==0
                ask=0;
            else
                imageCut(x)=[];
                m(x,:,:)=[];
                %Plot again
                close(f)
                f=figure;
                for i=1:size(imageCut,2)
                    nexttile
                    v = htfakergb(imageCut(i).data,imageCut(i).mask);
                    imagesc(v);
                    title(i)
                end
            end
        else
            over=min(x);
            del=max(x);
            n1=squeeze(m(x(1),:,:));
            n2=squeeze(m(x(2),:,:));
            n=n1+n2;
            ny=sum(n,2);
            ly=find(ny>0);
            nx=sum(n,1);
            lx=find(nx>0);
            imageCut(over).data=data(ly(1):ly(end),lx(1):lx(end),:);
            imageCut(over).mask=logical(n(ly(1):ly(end),lx(1):lx(end),:));
            m(over,:,:)=n;
            imageCut(del)=[];
            m(del,:,:)=[];
            %Plot again
            close(f)
            f=figure;
            for i=1:size(imageCut,2)
                nexttile
                v = htfakergb(imageCut(i).data,imageCut(i).mask);
                imagesc(v);
                set(gca,'XTick',[], 'YTick', [])
                title(i)
            end
        
        end
    end
end