-
Notifications
You must be signed in to change notification settings - Fork 5
/
detectMaskRCNN.m
145 lines (109 loc) · 4.49 KB
/
detectMaskRCNN.m
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
function [bboxes, scores, labels, finalMasks] = detectMaskRCNN(dlnet, maskSubnet, image, params, executionEnvironment)
% detectMaskRCNN runs prediction on a trained maskrcnn network
%
% Inputs:
% dlnet - Pretrained MaskRCNN dlnetwork
% maskSubnet - Pretrained mask branch of maskRCNN dlnetwork
% image - RGB image to run prediction on. (H x W x 3)
% params - MaskRCNN network configuration object Created using
% helper.createNetworkConfiguration
%
% Outputs:
% bboxes - Final bounding box detections ([x y w h]) formatted as
% NumDetections x 4.
% scores - NumDetections x 1 classification scores.
% labels - NumDetections x 1 categorical class labels.
% finalMasks - Binary object Masks detections formatted as
% H x W x NumDetections
% Copyright 2020 The MathWorks, Inc.
% Prepare input image for prediction
if(executionEnvironment == "gpu")
image = gpuArray(image);
end
% Cast the image to a dlarray
X = dlarray(single(image),'SSCB');
imageSize = size(image);
%%% Detector prediction
% Outputs for object detection
featureMapNode = 'res4b22_relu';
outputNodes = {'rpl', 'rcnnSoftmax', 'fcBoxDeltas'};
outputNodes = [outputNodes featureMapNode];
% Run prediction on the inputs
[bboxes, YRCNNClass, YRCNNReg, featureMap] = predict(...
dlnet, X, 'Outputs', outputNodes);
% Extract data from the output dlarrays
bboxes = extractdata(bboxes)';
YRCNNClass = extractdata(YRCNNClass);
YRCNNReg = extractdata(YRCNNReg);
% Network outputs proposals in [x1 y1 x2 y2] format in the image
% space. Convert to [x y w h] for further processing.
bboxes = vision.internal.cnn.boxUtils.x1y1x2y2ToXYWH(bboxes);
% Classification data processing
% Calculate scores
allScores = squeeze(YRCNNClass);
% remove scores associated with background
bgIndex = strcmp(string(params.ClassNames),params.BackgroundClass);
allScores(bgIndex,:) = [];
scores = reshape(allScores',[],1);
numClasses = numel(params.ClassNames)-1; % exclude background
% replicate each proposal box for each class.
% The boxes are replicated as [b1 b1 b1 b1 b2 b2 b2 b2 ....]
bboxes = repmat(bboxes(:, 1:4), numClasses, 1);
numObservations = size(YRCNNReg,2);
% replicate labels
% Group the labels as [label1 label1 ... numObservation times, label2
% label2.... numObservation times ...]'
classNames = categorical(params.ClassNames, params.ClassNames);
% Remove background
classNames(classNames==params.BackgroundClass) = [];
classNames = removecats(classNames, params.BackgroundClass);
ind = (1:numel(classNames))';
labels = classNames(repelem(ind,numObservations,1));
labels = labels';
% Regression data processing
% YRCNNreg is [numClasses*4 numObs]. reshape
% to 4-by-(numClasses * numObs), where each column is
% the target for a specific classes. Data is arranged
% as follows:
%
% [c1 c2 c1 c2 c1 c2], where rows stride by
% numClasses goes to next observation.
reg = reshape(YRCNNReg, 4, numClasses, numObservations);
reg = permute(reg, [1 3 2]);
reg = reshape(reg, 4, [])';
% Apply regression
bboxes = helper.applyRegression(bboxes, reg, params.MinSize, params.MaxSize);
% filter invalid predictions
[bboxes, scores, labels] = helper.filterBoxesAfterRegression(bboxes,scores,labels, imageSize);
% Filter boxes clip boxes on the edge
bboxes = vision.internal.detector.clipBBox(bboxes,imageSize);
% Remove boxes that are too small after clipping.
tooSmall = any(bboxes(:,3:4) < params.MinSize,2);
bboxes(tooSmall,:) = [];
scores(tooSmall,:) = [];
labels(tooSmall,:) = [];
% Filter scores lesser than the Threshold
keep = scores >= params.Threshold;
bboxes = bboxes(keep,1:4);
scores = scores(keep,:);
labels = labels(keep,:);
% Perform NMS
if params.SelectStrongest
[bboxes, scores, labels] = selectStrongestBboxMulticlass(bboxes,scores,labels,...
'RatioType', 'Min', 'OverlapThreshold', 0.7);
end
%%% Mask segmentation
% Prepare final bboc detections for mask prediction
bboxesX1Y1 = vision.internal.cnn.boxUtils.xywhToX1Y1X2Y2(bboxes);
roiIn = dlarray([bboxesX1Y1 ones(size(bboxesX1Y1,1),1)]', "SSCB");
mask = predict(maskSubnet, roiIn, featureMap);
bboxes = gather(bboxes);
mask = gather(squeeze(extractdata(mask)));
finalMasks = false([imageSize(1) imageSize(2) size(bboxes,1)]);
% Resize and insert masks
for i = 1:size(bboxes,1)
m = imresize(mask(:,:,i), [bboxes(i,4) bboxes(i,3)],'cubic') > 0.5 ;
finalMasks(bboxes(i,2):bboxes(i,2)+bboxes(i,4)-1, ...
bboxes(i,1):bboxes(i,1)+bboxes(i,3)-1, i) = m;
end
end