function createSpaceNetLabelImages(spaceNetVegasDataDir,buildingLabelDir,distanceLabelDir)
% Write bigimage masks and label images for images in the directory
% vegasRGBImageDir. The building polygons are specified in the CSV files in
% the directory 'summaryData' within spaceNetVegasDataDir. The masks are
% written to disk at the specified location, maskDir. The building label
% images are written to disk at the specified location, buildingLabelDir
% and building label images are written to disk at the specified location,
% distanceLabelDir
%
% The building masks are obtained using L*a*b color segmentation. They 
% 
% The label images are obtained by using the pixel-coordinates of building polygons in CSV files. The polygons are used to
% create label images of the same spatial dimensions as the RGB images. The
% label images are written to disk as explained earlier.

% Copyright 2020 The MathWorks, Inc.

buildingCSVFile = fullfile(spaceNetVegasDataDir,'summaryData','AOI_2_Vegas_Train_Building_Solutions.csv');
csvds = datastore(buildingCSVFile);
buildingPolygons = readall(csvds);

vegasRGBImageDir = fullfile(spaceNetVegasDataDir, 'RGB-PanSharpen');

extension = '.tif';
imds = imageDatastore(vegasRGBImageDir,'FileExtensions',{extension});
imgIndices = 1:length(imds.Files);

for idx = 1:numel(imgIndices)
    [~,fileName] = fileparts(imds.Files{imgIndices(idx)});
    
    distanceLabelDirName = strcat(fileName,'_distance');
    buildingLabelDirName = strcat(fileName,'_label');
    if  ~exist(fullfile(distanceLabelDir,[distanceLabelDirName  extension]),'file') || ...
            ~exist(fullfile(buildingLabelDir,[buildingLabelDirName extension]),'file')
        
        I = imread(imds.Files{imgIndices(idx)});
        
        fileNameInJSON = fileName(strfind(fileName,'AOI'):end);
        bIdx = ismember(buildingPolygons.ImageId,fileNameInJSON);
        buildingPtsStr = buildingPolygons.PolygonWKT_Pix(bIdx);
        if isempty(buildingPtsStr) || strcmpi(buildingPtsStr{1},'POLYGON EMPTY')
            continue;
        end
        
        buildingPtsStr = cellfun(@(x)replace(x, {'0,'}, ';'), buildingPtsStr, 'UniformOutput', false);
        buildingPtsStr = cellfun(@(x)replace(x, {'POLYGON ('}, ''), buildingPtsStr, 'UniformOutput', false);
        buildingPtsStr = cellfun(@(x)replace(x, {'0)'}, ']'), buildingPtsStr, 'UniformOutput', false);
        buildingPtsStr = cellfun(@(x)replace(x, ',(', ';['), buildingPtsStr, 'UniformOutput', false);
        buildingPtsStr = cellfun(@(x)replace(x, '(', '['), buildingPtsStr, 'UniformOutput', false);
        buildingPtsStr = cellfun(@(x)replace(x, ')', ''), buildingPtsStr, 'UniformOutput', false);
        
        buildingPts = cellfun(@(x)round(str2num(x)),buildingPtsStr, 'UniformOutput', false);
        buildingPts = cellfun(@(x)fixPolygonCoordinatesValues(x, size(I)),buildingPts, 'UniformOutput', false);
        
        %% Building Label Images
        sizeI = size(I);
        % Meshgrid in world units for all pixels in this block.
        [xgrid, ygrid] = meshgrid(1:sizeI(1),1:sizeI(2));
        buildingLabel = zeros(size(I));
        totalRegion = zeros(size(I));
        onRegionTotal = zeros(size(I));
        for ptsIdx = 1:numel(buildingPts)
            % Continue when a polygon is empty
            if isempty(buildingPts{ptsIdx})
                continue;
            end
            
            buildingRegion = inpolygon(xgrid,ygrid,buildingPts{ptsIdx,1}(:,1),buildingPts{ptsIdx,1}(:,2));
            % Eg  'RGB-PanSharpen_AOI_2_Vegas_img953' has region cut-outs
            % within the building polygon.  Look at buildingPtsStr{end}
            % there are 2 cut-out regions within the main that need to be
            % handled properly.
            buildingRegion = double(buildingRegion);
            se = strel('disk',3);
            totalRegion = totalRegion + buildingRegion;
            inRegion = imerode(buildingRegion,se);
            onRegion = buildingRegion-inRegion;
            onRegionTotal = onRegionTotal + onRegion;
            buildingLabel = buildingLabel + imlincomb(0.5,inRegion,1,onRegion);
        end
        
        buildingLabel = imquantize(buildingLabel,[0 0.5],1:3);
        if  ~exist(fullfile(buildingLabelDir,[buildingLabelDirName extension]),'file')
            % Make a bigimage out of the image mask and set its spatial reference
            buildingLabelDirPathName = fullfile(buildingLabelDir,buildingLabelDirName);
            imwrite(uint8(buildingLabel),[buildingLabelDirPathName extension],'TIFF');
        end
        
        %% Distance Label Images
        D = bwdist(onRegionTotal);
        Dsigned = round(-1 .* (D .* ~totalRegion) + (D .* totalRegion));
        % Need to use negative values to represent outside building pixels.
        Dsignedq = imquantize(Dsigned,-64:63,1:129);
                
        % Make a bigimage out of the image mask and set its spatial reference
        if ~exist(fullfile(distanceLabelDir,[distanceLabelDirName  extension]),'file')
            distanceLabelDirPathName = fullfile(distanceLabelDir,distanceLabelDirName);
            imwrite(uint8(Dsignedq),[distanceLabelDirPathName extension],'TIFF');
        end
                
    end
    
end
end

function x = fixPolygonCoordinatesValues(x, imageSize)

% Early return when the polygon is empty
if isempty(x)
    return;
end

if any(any(x == 0))
    x(x==0) = 1;
end

if any(any(x(:,1) > imageSize(1)))
    x(x>imageSize(1)) = imageSize(1);
end

if any(any(x(:,2) > imageSize(2)))
    x(x>imageSize(2)) = imageSize(2);
end

end