Sliding FullyConvolutional net over larger images

Question

Take for instance the following net based on LeNet but with $1 \times 1$ convolutions instead of the MLP (the last two, fully connected layers) at the end of the classical LeNet.

lenetFullyConvolutional = NetInitialize@NetChain[{
ConvolutionLayer[20, 5], Ramp, PoolingLayer[2, 2],
ConvolutionLayer[50, 5], Ramp, PoolingLayer[2, 2],
(* MLP equivalent *)
ConvolutionLayer[500, {4, 4}], Ramp, 
ConvolutionLayer[10, {1, 1}] , FlattenLayer[] , SoftmaxLayer[]},
"Output" -> NetDecoder[{"Class", Range[0, 9]}],
"Input" -> NetEncoder[{"Image", {28, 28}, "Grayscale"}]
]

It can be trained exactly like the LeNet example from the documentation with

resource = ResourceObject["MNIST"]
trainingData = ResourceData[resource, "TrainingData"];
testData = ResourceData[resource, "TestData"];

(* train the net *)
trained = 
NetTrain[lenetFullyConvolutional, trainingData, 
ValidationSet -> Scaled[0.1], MaxTrainingRounds -> 3] 

Since the main part of the net only consists of convolutions it should be possible to slide it over images of arbitrary size such as a $1\times28\times56$ tensor (two appended MNIST images) that would result not in a single class vector ($1\times1\times10$) but a output volume such as ($1\times2\times10$) when using stride size 28.

See here for reference.

Does anyone know how to archive this using the current neural networks capabilities in Mathematica?

I tried taking the ConvolutionLayers with Take[lenetFullyConvolutional, {1, -1}] but I am not sure how to change the input dimensions accordingly so to be able to apply the net to arbitrarly sized input.

Answer 1

Introduction

The fully convolutional neural network (FCN) is very useful for object detection and segmentation. One of the first demonstration of FCN for semantic segmentation is in this paper. FCN has also been used in the recent Mask-RCNN architecture as a critical part for its high accuracy for semantic segmentation.

This answer tries to provide an example of using an FCN to generate heat maps that show the location of objects in the image.

We will work on this image (taken from here)

We crop it to have a dimension of 2000 by 1600 to avoid crashing Mathematica

i = ImageCrop[img,{2000, 1600}];

And we will convert a VGG-16 network to an FCN and slides on the image to generate the heat map. The results look like this

Conver to FCN

In order to convert a neural network to an FCN, we have to convert all the fully connected layers into convolutional layers. The conversion should be in such a way that it preserves the output of the original network on images with the size of the original input (224 by 224 for VGG-16).

We use the trained version of the VGG-16 network

vgg = NetModel["VGG-16 Trained on ImageNet Competition Data"];

It has three fully connected layers and we need to convert each of them into convolutional layers

fc6Weights = NetExtract[vgg, {"fc6", "Weights"}];
fc6Biases = NetExtract[vgg, {"fc6", "Biases"}];

conv6Weights = ArrayReshape[fc6Weights, {4096, 512, 7, 7}];
conv6Biases = fc6Biases;

conv6 = ConvolutionLayer["Weights" -> conv6Weights, 
   "Biases" -> conv6Biases];

fc7Weights = NetExtract[vgg, {"fc7", "Weights"}];
fc7Biases = NetExtract[vgg, {"fc7", "Biases"}];

conv7Weights = ArrayReshape[fc7Weights, {4096, 4096, 1, 1}];
conv7Biases = fc7Biases;

conv7 = ConvolutionLayer["Weights" -> conv7Weights, 
   "Biases" -> conv7Biases];

fc8Weights = NetExtract[vgg, {"fc8", "Weights"}];
fc8Biases = NetExtract[vgg, {"fc8", "Biases"}];

conv8Weights = ArrayReshape[fc8Weights, {1000, 4096, 1, 1}];
conv8Biases = fc8Biases;

conv8 = ConvolutionLayer["Weights" -> conv8Weights, 
   "Biases" -> conv8Biases];

Now we can construct the FCN together with these converted layers

net = NetChain[
  Flatten@{Values@Normal@Take[vgg, {1, "pool5"}], conv6, Ramp, 
    DropoutLayer[0.5], conv7, Ramp, DropoutLayer[0.5], conv8}, 
  "Input" -> 
   NetEncoder[{"Image", {224, 224}, 
     "MeanImage" -> {0.4850196078431373, 0.457956862745098, 
       0.4076039215686274}}]]

To make it truly an FCN, we have to remove the information of the input size on the layers, we can use the removeInputInformation in Sascha answer to do that:

fcn = {Values@Normal@Take[vgg, {1, "pool5"}] // Take[#, {1, -1}] & // 
      Normal // Map[removeInputInformation], conv6, Ramp, 
    DropoutLayer[0.5], conv7, Ramp, DropoutLayer[0.5], conv8} // 
   Flatten // NetChain

Removing the input size information allows us to apply the network to images with arbitrary sizes.

Convert to heat map

Before we apply our FCN to the image, we need to encode the image and subtract the mean image of the VGG network

vggmeanimage = {0.4850196078431373, 0.457956862745098, 
   0.4076039215686274};
enc[i_] := 
 Transpose[Map[# - vggmeanimage &, ImageData[i], {-2}], {2, 3, 1}]

Now we can apply our FCN to the image

res = fcn[enc[i]]; // AbsoluteTiming
{150.764, Null}

Since our image has a size larger than 224x244, the FCN slides on the image, takes patches of size 224x224 and generate a classification result for each patch. We can see that the result we get from the FCN has dimension 1000 by 44 by 56, corresponding to the 1000 classes and 44 by 56 sliding patches.

res // Dimensions
(* {1000, 44, 56} *)

For example, res[[All,1,1]] tells us the probabilities of each class on the first patch.

In order to make a heat map of the probability of the dog, we use only the results from the four highest dog classes

dogindex = 
 Ordering[SoftmaxLayer[][
   Take[vgg, {1, -2}][ImageCrop[i, {1000, 1000}, Top]]], -5]
(* {233, 228, 249, 265, 264} *)

vgg16Classes = 
  NetExtract[
    NetModel["VGG-16 Trained on ImageNet Competition Data", 
     "UninitializedEvaluationNet"], "Output"][["Labels"]];

vgg16Classes[[dogindex]]
(* {Entity["Concept", "BorderCollie::463w2"], 
 Entity["Concept", "Kelpie::k6795"], 
 Entity["Concept", "EskimoDog::2vm97"], 
 Entity["Concept", "Cardigan::b724j"], 
 Entity["Concept", "Pembroke::95g54"]} *)

And we can plot the heat map and overlay it on the original image

heat = ListDensityPlot[#/Max[#] &@Reverse[Plus @@ res[[dogindex]]], 
   ColorFunction -> (Blend[{Transparent, Yellow, Red}, #] &), 
   PlotRange -> {All, All, {0.2, 0.8}}, ClippingStyle -> Automatic, 
   PlotRangePadding -> None, ImagePadding -> None, Frame -> None, 
   ImageSize -> ImageDimensions[i], 
   AspectRatio -> (ImageDimensions[i][[2]]/ImageDimensions[i][[1]])];

ImageCompose[i, {heat, 0.7}]

Reference

1. Fully Convolutional Networks for Semantic Segmentation, arXiv:1605.06211
2. Mask R-CNN, arXiv:1703.06870
3. convnets-keras, https://github.com/heuritech/convnets-keras

Answer 2

This is probably only a partial answer to my question but I have found a way to convert a fully convolutional net (FCN) trained on images of a specific size to an net that takes arbitrary sized inputs.

To this end we first define and train a FCN for use with MNIST images ($28\times28$ pixel) as has been done in the question text already.

Next we need to remove the input size information from the trained net

I had no luck doing this with any combination of Take or NetExtract on the whole net so I wrote the following function which takes a layer, extracts its relevant properties and creates a copy of it but with it's input attributes set to Automatic for the spacial dimensions (which happens automatically by not specifying any input dimensions).

removeInputInformation[layer_ConvolutionLayer] := 
With[{
k = NetExtract[layer, "OutputChannels"], 
kernelSize = NetExtract[layer, "KernelSize"] , 
weights = NetExtract[layer, "Weights"], 
biases = NetExtract[layer, "Biases"], 
padding = NetExtract[layer, "PaddingSize"], 
stride = NetExtract[layer, "Stride"],
dilation = NetExtract[layer, "Dilation"]
},
ConvolutionLayer[k, kernelSize, "Weights" -> weights, 
"Biases" -> biases, "PaddingSize" -> padding, "Stride" -> stride, 
"Dilation" -> dilation]
]

removeInputInformation[layer_PoolingLayer] := 
With[{
f = NetExtract[layer, "Function"], 
kernelSize = NetExtract[layer, "KernelSize"] , 
padding = NetExtract[layer, "PaddingSize"], 
stride = NetExtract[layer, "Stride"]
},
PoolingLayer[kernelSize, stride, "PaddingSize" -> padding, 
"Function" -> f ]
]

removeInputInformation[layer_ElementwiseLayer] := 
With[{f = NetExtract[layer, "Function"]},
ElementwiseLayer[f]
]

It is then possible to take the relevant part of the net with Take, disassemble it with Normal, transform each layer with removeInputInformation and reassemble the net again to get a net that will accept an input of any dimensions.

newFCN =  trained //Take[#, {1, -3}] & //Normal  
//Map[removeInputInformation] //Append[{TransposeLayer[1 -> 2]}] 
// Flatten  //NetChain

Note that I appended a TransposeLayer to the net for convenience later on.

---

Testing the net

To create some test images for the net I wrote the following function that joins multiples images and adds some adjustable padding in-between them.

appendImages[images_, pad_: 1] := 
With[{dim = ImageDimensions@First@images}, 
images // Riffle[#, Image@ConstantArray[0.5, dim]] & // 
Map[Transpose@*ImageData] // Apply[Join] // Transpose // Image]

Generating some test data:

sample = RandomSample[testData, 10] // Part[#, All, 1] &;
strip = appendImages[sample, 28]

and applying the net to the data

newFCN[{ImageData@strip}] // Flatten[#, 1] & 
//ArrayPlot[#,FrameTicks -> {{Transpose@{Range[10], Range[0, 9]}, None}, {None, None}}] &

Adding a PoolingLayer(max pooling) to the net helps to improve the number detection/recognition.

newFCN = trained //Take[#, {1, -3}] & //Normal    //Map[removeInputInformation] 
//Append[{TransposeLayer[1 -> 2],PoolingLayer[{1, 4}, "Stride" -> {1, 4}, 
   "PaddingSize" -> {0, 4}, "Function" -> Max]}] 
//Flatten //NetChain