# How can image segmentation from UNET be improved?

EDIT 2: code can be found @ https://github.com/alihashmiii/UNet-Segmentation-Wolfram

EDIT 1: applying Binarize to the masks after resizing helps to ensure the data for mask is binary

I have made an attempt to implement the neural network (written in R) @ https://github.com/rstudio/keras/blob/master/vignettes/examples/unet.R in Mathematica. The purpose of the network is to take in a grayscale image of an aggregate of cells and output a binarized mask (shown below).

The dataset (images and corresponding labels) for training can be downloaded from:

for images

https://www.dropbox.com/sh/bi0sky2hejjhqit/AAAT3YSzVwvDO0HW3POLJs0Ba?dl=0

for masks

https://www.dropbox.com/sh/yb0itaoyuz2y422/AAB25g6i9A2Mdlz1TwJgfzhka?dl=0

Note: although i am not sure how well i did with translating the network. It seems to be working. However, I am looking for performance boosts

Creating The Net

convolutionModule[kernelsize_, padsize_, stride_] := NetChain[
BatchNormalizationLayer[],
ElementwiseLayer[Ramp]}
];

{DeconvolutionLayer[1, {2, 2}, "PaddingSize" -> 0, "Stride" -> {2, 2}],
BatchNormalizationLayer[],
ElementwiseLayer[Ramp],
];

PoolingLayer[{2, 2}, "Function" -> Max, "Stride" -> {2, 2}]
];

cropLayer[{dim1_, dim2_}] := NetChain[
{PartLayer[{1, 1 ;; dim1, 1 ;; dim2}],
ReshapeLayer[{1, dim1, dim2}]}
];

UNET := Block[{kernelsize = {3, 3}, padsize = {1, 1}, stride = {1, 1}},
NetGraph[
<|"enc_1" -> encodeModule[kernelsize, 0, stride],
"dropout_1" -> DropoutLayer[],
"dec_2" -> decoderModule[kernelsize, padsize + 1, stride],
"crop_1" -> cropLayer[{20, 20}],
"concat_1" -> CatenateLayer[],
"dropout_2" -> DropoutLayer[],
"crop_2" -> cropLayer[{40, 40}],
"concat_2" -> CatenateLayer[],
"dropout_3" -> DropoutLayer[],
"crop_3" -> cropLayer[{80, 80}],
"concat_3" -> CatenateLayer[],
"dropout_4" -> DropoutLayer[],
"map" -> {ConvolutionLayer[1, kernelsize, "Stride" -> stride,
{{NetPort["Input"] ->
"enc_1" ->
"enc_2" ->
"enc_3" ->
"enc_4" -> "dropout_1" -> "enc_5" -> "dec_1" -> "dec_2"},
"dec_2" -> "crop_1",
{"enc_3", "crop_1"} ->
"concat_1" -> "dropout_2" -> "conv_1" -> "dec_3",
"enc_2" -> "crop_2",
{"crop_2", "dec_3"} ->
"concat_2" -> "dropout_3" -> "conv_2" -> "dec_4",
"enc_1" -> "crop_3",
{"crop_3", "dec_4"} ->
"concat_3" -> "dropout_4" -> "conv_3" -> "dec_5" -> "map"
},
"Input" ->
NetEncoder[{"Image", {168, 168}, ColorSpace -> "Grayscale"}],
"Output" ->
NetDecoder[{"Image", Automatic, ColorSpace -> "Grayscale"}]
]
] // NetInitialize;


Training The Net

dir = "C:\\Users\\aliha\\Downloads\\fabrice-ali\\deeplearning\\data\\train\\train_images_8bit";

SetDirectory[dir];

images = ImageResize[Import[dir <> "\\" <> #], {168, 168}] & /@ FileNames[];
(* we shuffle the images *)
keysshuffle = Keys@shuffledimages;

SetDirectory[dir];

masks = Binarize[ImageResize[Import[dir <> "\\" <> #],{160, 160}]]&/@FileNames[];
(* resize masks to 160,160 same size as the output of the net *)

(* shuffle the masks -> same order as images *)

{dataset, validationset, unseen} = TakeList[Values@shuffledimages, {290, 80, 20}];

{labeldataset, labelvalidationset, groundTruth} = TakeList[shuffledmasks, {290, 80, 20}];

trainedNN =
ValidationSet -> Thread[validationset -> labelvalidationset], BatchSize -> 8,
MaxTrainingRounds -> 100, TargetDevice -> "GPU" ];


As can be seen from the image above, the net seems to be working after a 100 rounds of training. For some cases the network works reasonably but for some not so well. However, I think the performance can be improved and I am happy to put anyone in credits who can help me improve the performance.

• I am facing a few problems with my classifier as it is not improving. How do you augment the images and it's mask? Do you use a generator from keras or do you define your own? – sriram vijendran Sep 4 '18 at 8:47
• Hi Sriram, this should be added as a new question if you're working in Wolfram Mathematica, and moved elsewhere if not. :) – Carl Lange Sep 4 '18 at 9:13
• @sriramvijendran did you try using the links I have mentioned in the question, namely in Edit 2 – Ali Hashmi Sep 8 '18 at 14:16
• What is this dataset that you are using, I want to use this dataset for my work. Can I know of the exact link and citation details – Srijita Roy Dec 12 '18 at 10:01

I built a U-net version in MMA following the paper exactly then modified it for my application (which used 8-band images and 10 output classes). It worked pretty well for me. Your implementation (at least from the code you've put in the post above) differs from the original U-net paper in several ways.

• I noticed that the number of convolution channels (or trainable convolution kernels) isn't what's described in the paper - these increase in number in the contraction stages and decrease in the expansion stages. This is important for the way the net works.
• You seem to have only one convolution layer rather than two in each contraction block, this will probably affect performance as well.
• I stuck with the original image dimensions described in the paper (768x768 input and 388x388 output) and used a reflected image padding to increase the input image size.
• I used training set augmentation quite heavily - rotations, scaling and mirroring - and this improved classification accuracy,
• I used an intersection-over-union measure to measure classification accuracy.

One of the things I haven't yet implemented in my version of U-net is introducing weighted training data to deal with class imbalance in the training data; there is also a more tricky weighting of data for class boundary pixels. This is described in the paper but isn't trivial to implement. Class imbalance didn't have as large an impact on performance as I expected but that may be just the training data and classifications I've been doing.

One other thing I suppose I should note is that I had buy a new, bigger GPU to train the network (a GTX 1080) as there wasn't enough memory on the M1000 I started with (although it did train painfully slowly on CPU).

I hope this all helps

• thanks for pointing out the differences. Actually if you see the link posted above for the code, Alexey and I have modified it to keep the spirit of the original version. It is doing a good job now ! – Ali Hashmi May 7 '18 at 9:51
• OK, glad to hear it's working better. I did try the github link but couldn't get the notebook or source file to open correctly in MMA - presumably I need some package to access .nb or .m MMA files from github? – EstabanW May 7 '18 at 13:08
• @EstabanW Would you be able to publish your UNET implementation somewhere? – Carl Lange Aug 25 '18 at 19:28
• This will be a good reference github.com/zhixuhao/unet/blob/master/model.py – Wenuka Dec 12 '19 at 13:55