Create Instance Segmentation Training Data From Ground Truth

Open Live Script

This example shows how to create instance segmentation training data from a groundTruth object. To train a Mask R-CNN and SOLO v2 network using the trainMaskRCNN and trainSOLOv2 functions, respectively, format your input data as a 1-by-4 cell array containing the RGB training image, bounding boxes, instance labels, and instance masks. Convert your polygon ground data into a set of binary instance masks and axis-aligned rectangles for training.

First, load the ground truth, the images visionteam.jpg and visionteam1.jpg, and create a groundTruth object.

gt = load("gTruth.mat");
src = groundTruthDataSource(["visionteam.jpg","visionteam1.jpg"]);
gTruth = groundTruth(src,gt.labelDef,gt.labelData)

gTruth = 
  groundTruth with properties:

          DataSource: [1×1 groundTruthDataSource]
    LabelDefinitions: [1×5 table]
           LabelData: [2×1 table]

Load the training images using the ImageDatastore object.

imds = imageDatastore(gTruth.DataSource.Source);

Gather the polygon data from the groundTruth object and group it by label type so that all the polygon data for all classes is combined.

labelData = gatherLabelData(gTruth,labelType.Polygon,GroupLabelData="LabelType");

Identify which images have label data and only keep those rows of data.

hasData = rowfun(@(x)~isempty(x{1}),labelData{1});
imds = subset(imds,hasData{:,1});
labels = labelData{1}(hasData{:,1},:);

Create an arrayDatastore to read the label data.

labelDS = arrayDatastore(labels{:,1},IterationDimension=1,OutputType="same");

Combine the image and label data. Apply a data transform using the helperPolygonToMaskTransform helper function to convert each polygon into a bounding box and a stack of binary masks. Then, preview the results.

ds = combine(imds,labelDS);
tds = transform(ds,@(data,info)helperPolygonToMaskTransform(data,info),IncludeInfo=true);
data = preview(tds)

data=1×4 cell array
    413×800×3 uint8    3×4 double    3×1 cell    413×800×3 logical

To speed up training, write the binary instance masks, bounding boxes, and labels to disk into a .mat file, instead of computing them during training, using the helperWriteMaskBoxesAndLabels helper function. The name of the MAT file matches the name of the corresponding image.

labelDataFolder = fullfile(pwd,"masksBoxesAndLabels");
if ~isfolder(labelDataFolder)
    mkdir(labelDataFolder);
    classes = string(gTruth.LabelDefinitions.Name);
mkdir(labelDataFolder);
    writeall(tds,labelDataFolder,...
        WriteFcn=@(data,info,fmt)helperWriteMaskBoxesAndLabels(data,info,fmt,classes));
end

Create a fileDatastore object to read MAT files during training using the helperReadMaskBoxesAndLabels helper function. Ensure that you read only valid MAT files using the helperValidMATFiles helper function.

fs = helperValidMATFiles(labelDataFolder);
labelDS = fileDatastore(fs,...
    IncludeSubfolders=true,...
    ReadFcn=@helperReadMaskBoxesAndLabels,...
    FileExtensions=".mat",...
    UniformRead=true);

Order the image and MAT files so that they are in the same order.

[~,ordImgs] = sortrows(imds.Files);
[~,ordLabels] = sortrows(labelDS.Files);
imds = subset(imds,ordImgs);
labelDS = subset(labelDS,ordLabels);

Combine the image and label datastores into a datastore you can use for training the Mask R-CNN and SOLOv2 instance segmentation networks. Preview the training datastore.

ds = combine(imds,labelDS);
preview(ds)

ans=1×4 cell array
    413×800×3 uint8    3×4 double    3×1 categorical    413×800×3 logical

Supporting Functions

`helperPolygonToMaskTransform`

function [out,info] = helperPolygonToMaskTransform(data,info)
% Transform polygons into binary instance masks and bounding boxes.
out = cell(size(data,1),4);
for i = 1:size(data,1)
    I = data{i,1};
    annotationStruct = data{i,2};
    if isempty(annotationStruct)
        maskset = false(0,0,0);
        boxes = zeros(0,4);
        labels = createArray(0,1,Like="");
    else
        s = data{i,2};
        % Convert polygon into a set of binary masks and a set of bounding
        % boxes.
        [maskset,boxes,labels] = helperPolygonsToMasksAndBoxes(s,size(I));
    end

    % Package output in the required for training.
    out{i,1} = I;
    out{i,2} = boxes;
    out{i,3} = labels;
    out{i,4} = maskset;
end
end

`helperPolygonsToMasksAndBoxes`

function [maskset,boxes,labels] = helperPolygonsToMasksAndBoxes(s,sizeI)
% Transform polygons into binary instance masks and bounding boxes.
% The input s is an M-by-N cell. Each row contains ROI data for a particular class.
numInstances = size(s,1);
maskset = false([sizeI(1:2) numInstances]);

% Generate labels for each object instance.
labels = s(:,2);

% Convert polygons into binary instance masks and bounding boxes.
boxes = zeros(numInstances,4);
for i = 1:numInstances
    poly = s{i,1};
    minxy = min(poly);
    maxxy = max(poly);
    boxes(i,:) = reshape([minxy maxxy-minxy],[],4);
    maskset(:,:,i) = poly2mask(poly(:,1),poly(:,2),sizeI(1),sizeI(2));
end
end

`helperWriteMaskBoxesAndLabels`

function helperWriteMaskBoxesAndLabels(data,writeInfo,~,classes)
% Write binary instance masks, boxes, and labels into a MAT file.
[fPath,fName,~] = fileparts(writeInfo.SuggestedOutputName);
labelData.Boxes = data{2};
labelData.Labels = categorical(string(data{3}),classes);
labelData.Masks = data{4};
save(fullfile(fPath,fName+".mat"),"labelData");
end

`helperReadMaskBoxesAndLabels`

function data = helperReadMaskBoxesAndLabels(filename)
% Read instance masks, boxes, and labels from a MAT file.
loaded = load(filename);
data = cell(1,3);
data{1} = loaded.labelData.Boxes;
data{2} = loaded.labelData.Labels;
data{3} = loaded.labelData.Masks;
end

`helperValidMATFiles`

function fs = helperValidMATFiles(labelDataFolder)
fs = matlab.io.datastore.FileSet(labelDataFolder,"FileExtensions","*.mat",IncludeSubfolders=true);
if ismac
    % Remove files beginning with "._" that mac OS adds when writing to
    % certain file systems. 
    files = fs.FileInfo.Filename;
    fs = fs.subset(~files.contains("._"));
end
end