UNET: neural network for 2D & 3D image segmentation w/ medical examples

[mfroeling (Martijn Froeling)]

UNET

---

A package to generate and train a UNET deep convolutional network for 2D and 3D image segmentation.

Some code was based on work by @Ali Hashmi, which was also dicussed in this post
The full version of the toolbox can be found on my github page.

• Information
• Install toolbox
• Using the toolbox
• Functionality
• Visualization
• Example

Information

UNET is developed for Mathematica.
It contains the following toolboxes:

• UnetCore
• UnetSupport

Documentation of all functions and their options is fully integrated in the Mathematica documentation.
The toolbox always works within the latest version of Mathematica and does not support any backward compatibility.

All code and documentation is maintained and uploaded to github using Workbench.

Install toolbox

Install the toolbox in the Mathematica UserBaseDirectory > Applications.

FileNameJoin[{$UserBaseDirectory, "Applications"}]

Using the toolbox

The toolbox can be loaded by using <<UNET`

The notbook UNET.nb shows examples of how to use the toolbox on artificially generated 2D data.
There are also examples how to visualize the layer of your trained network and how to visualize the training itself.

Functionality

The network supports multi channel inputs and multi class segmentation.

• UNET generates a UNET convolutional network.
  • 2D UNET

* 3D UNET  

• Loss Layers: Training the data is done using three loss layers: a SoftDiceLossLayer, BrierLossLayer and a CrossEntropyLossLayer.

• Convolution Blocks: The toobox contains five different convolution blocks that build up the network: UNET, UResNet, RestNet, UDenseNet, DensNet.

• SplitTrainData splits the data and labels into training, validation and test data.

• TrainUNET trains the network.

Visualization

• Visualize the network and results.
  • Visualize the features of the layers.

* Visualize the results.  

* Animate the training process.  

Example

• Example: 3D segmentation of lower legg muscles using MRI data.

UNET.nb (1.22 MB)

[sghuisman (Sander Huisman)]

Really neat! I wish I understood how this works! I still don’t fully understand how you can make a neural network that has a different amount of outputs. Is that easily explained or can i read about this somewhere?

Groetjes!

[mfroeling (Martijn Froeling)]

Thanks! What do you mean with different amounts of outputs?

This is my understanding, but i might be mistaken.
In my functions i layout the network architecture. Once the function is called you have to fill in the number of input classes, data dimensions, number of features and the number of output classes. From that moment everything is fixed. The UNET only has one output, an array with the binary segmentation.

The only place where there are multiple output ports is in the loss function, which is well explained in the help of mathematica LossFunction. But these are only used in training.

My approach was to create the neural net with only an Input, the data, and an Output, the segmentation, defined.
Only when I start training the net I add the appropriate loss functions which have as input the Output of the net and the ground truth Target for comparison. The loss function it self just outputs a number. You can add as many loss functions as needed and you can even provide a different Target to each. As i understand the training tries to minimize the sum of all the loss functions.

Hope this helps.

[wolframcommunity (EDITORIAL BOARD)]

• Congratulations! This post is now a Staff Pick as distinguished by a badge on your profile! Thank you, keep it coming, and consider contributing your work to the The Notebook Archive!

[szhorvat (Szabolcs Horvát)]

Be sure to put it on PackageData:

http://packagedata.net/

[szhorvat (Szabolcs Horvát)]

It looks like a nice package, and I am very happy to see that more and more people are publishing packages, despite of WRI’s rather weak support for package developers and package development.

I was very disappointed to read this in the recent blog post:

Minimal dependencies (no collections of competing libraries from different sources with independent and shifting compatibility).

It sounds as if the goal were to explicitly avoid structuring Mathematica into packages in a logical way (everything goes into System) and want Mathematica to have everything built-in and Mathematica users to not rely on third-party packages. That’s impossible, WRI doesn’t have the resources to produce even a tiny fraction of what we users need.

But, as disappointing as this public statement is (there are several other things in the blog post suggesting the same attitude), that’s a discussion for another time.

---

Unfortunately, it does seem to me that many users are distrusting of any functionality that doesn’t come from Wolfram directly. “I know you have this package but I’d rather use a built-in …” This situation was created by Wolfram in part by not encouraging package development and not providing functionality/resources (and implying that “it should be built-in”), but also by not educating the community about good package development practices. The unfortunate truth is that many of packages out there follow at least one questionable practice, which will eventually affect those that try to use it by causing unintended side effects, breakage, all sorts of problems. No wonder people don’t trust packages developed by random people.

What can the community do about this?

I would like to propose setting up an informal network of WL package developers who help and support each other by looking at each other’s work, doing informal code reviews, educating each other and learning from each other.

I am going to start here and I will point out a few small things that could be improved in this package (UNET). I hope you will find it helpful.

[mfroeling (Martijn Froeling)]

will do

[mfroeling (Martijn Froeling)]

All input is welcome.

I agree, there is little information about how to write a robust package/paclets. Most I learned along the way from forums.
I like using workbench in eclipse for debugging and paclet development and generating documentation.

Regrettably it can sometimes be a bit buggy and for the HTML documentation build you need to fix some things but if you do it is easily used online.

[szhorvat (Szabolcs Horvát)]

Please keep in mind that doing this review is an experiment, and I am hoping to have a positive outcome. I do not mean to intrude, or hijack Martijn’s package, and I do respect his choices, whatever those may be.

I am making these comments in the hope to benefit both the developer and the user community.

I also hope that others will spontaneously comment on my own work in a similar manner!

---

Due to lack of time, the below observations will be a bit superficial. I might try to dive deeper later.

init.m

Kernel/init.m is run when $Context is still Global` . Any non-System symbols you use in this file will pollute the Global namespace. Try to avoid this. When the package gets loaded, it should not create any symbols in Global.

I have seen the following update of the context path multiple times:

$ContextPath = Union[$ContextPath, System`$UNETContextPaths];

Note that this reorders the context path!! That includes not having System and Global at the very end, where they should be.

Please don’t do this. It will cause unexpected and hard to track down problems for your users.

I noticed that all symbol names are Removed at the start. Maybe I am missing something, but such drastic measures are not usually necessary. ClearAll will do fine.

Note that constructs like ClearAll["foo`*"] or Protect["foo`*"] do work. SetAttributes["foo`*", ReadProtected] does not work, but SetAttributes[Evaluate@Names["foo`*"], ReadProtected] does. Note that none of these constructs introduce any new symbols, thus they are convenient to use in init.m.

Consider not printing anything during package loading without good reason. Note that people might load the package as part of another package. They might even load it on subkernels.

(I do believe that sometimes there is a practical reason. I do it with my IGraph/M package, though I suppress printing when I can detect that IGraph/M is loaded as part of another package. I keep the printed message as small as possible, yet some users still complain about it. It can be suppressed with Block[{Print}, Needs["..."]])

Usage messages

I wanted to note that it’s great that you took care to make them work well with auto-completion.

Miscellaneous

size is not localized in VisualizeUNET2D.

AddLossLyaer is misspelt in the doc guide page.

UNET symbol link doesn’t work in the docs. Was this renamed to MakeUNET?

Cosmetic

Make the version shown in PacletInfo match the actual version. It’s show in the doc browser.

Compatibility

Indicate in the README which Mathematica version is required and try to test with that version. The IDEA WL plugin is capable of highlighting symbols not present in the WL version you set. It’s useful for detecting certain compatibility problems early. http://wlplugin.halirutan.de

[sghuisman (Sander Huisman)]

I mean, how can the network ‘decide’ that there is one object detected, or two, or three, or four… (i know the output will be some matrix with indices). But I don’t get how the network ‘numbers’ the output…

Thanks for the explanation!

[wolfganghitzl]

Congratulations and many thanks for all your efforts, suggestions. Especially, I do really like the idea to illustrate and animate the training process.

For some applications, it might be interesting to use the segmented regions as boundaries for a partial differential equation system (e.g. to descrive a diffusion process or a wave or so)…

Frankly speaking, I have never seen these interesting loss layers, like SoftDiceLossLayer, BrierLossLayer.

Time to get closer to these network architectures!

[mfroeling (Martijn Froeling)]

Hi Szabolcs,

Thank for the valuable comments.
The Print was there for debugging and should have been removed and agree with you completely on that, all print statements during loading are triggered by verbose.
I will fix the documentation.

I had not realized that init pollutes the global namespace.
I was thinking of moving the stuff from the init file to a UnetLoader.wl file such that init only contains Get["UnetLoader`"] which could be a nicer solution but somehow i have issues with initiating the sub-packages from that location.
Also I would like to hear your thought a bit more on the clearing of the namespace. I use the Remove since that is the only way i could clear shadowed definitions once they have been created.
If you define a function in Global and then initiate a toolbox with the same function even when using ClearAll the symbol remains shadowed. Only by adding remove this is solved.

Here you can see that without Remove the functions are highlighted in red indicating they are shadowed while the only Context they have is Main.

[mfroeling (Martijn Froeling)]

Hope this helps

We have one input channel Nchan = 1

Step 1: split the lables into classes, for this exmaple Nclass = 4 (background, sphere, rectangle and overlap), so the network looks for 4 output lables.

Step 2: The first encoding layer split our input channel into 32 channels 1->32

Step 3: by the time we arrive at the deepest level our channels are downsampled but we have a lot of them

Step 4: After decoding we go back to 32 channels again. Hopefully by now each channel contains unique information about the image.

Step 5: We are looking for 4 labels so the 32 channels are mapped to 4 channels

Step 6: The 4 channels are converted to probability maps.

Step 7: The probability map become our 4 labels. In this case the network fail to correctly segment our example.

Our background and overlap are segmented whit a very high probability because they have a very distinct contrast which is easy to detect. However our circle and rectangle have the same contrast and have to be labeled based on their shape and not contrast. This is much more difficult and as such mistakes are made.

[joseafalm (José Antonio Fernández)]

Is there any reason, other than memory limitations, to limit the maximum input sizes to 256 for each dimension? Shouldn’t it be enough with being multiple of 16?

I have overriden the definition of TrainUNET so as to allow that circunventing the size cap and am training on 400x304 images, and definitely training- time blows up, but hope the rest is OK. Maybe some hyper parameters would need to be adapted for the larger images?

[anon21378091]

No there is no reason why the max size cannot be greater that 256, as long as the inputs size can be divided by 2 four times.

So any dimension that can be generated by (dim)2^4 should be fine. But indeed you will need more memory on your graphics card and training will take longer.

In TrainUNET the line that check dimensions I actually already increased for my own working version.

In this line the value max can be whatever you want, it was originally set to 20 but can be larger. 30 would do the trick for you.

(*check if data dimensions are valid for net*)
	If[(!AllTrue[MemberQ[Range[max]2^4,#]&/@datDim,TrueQ])&&(MemberQ[{2,3},netDim]),

The TrainUNET you don’t have to use, it just does some extra stuff I found useful back when I was using it.

I don’t know the optimal settings for the hyper parameters, the values that are in the UNET are ones that worked well for my problems so they might need tweaking.

[drargento (Claudio Argento)]

Martjin,

I am very confused on how to load the paclets to run UNET. What is the file structure and where should I put it? When I use “File-Install-Package”, what paclets should I use? I am not really familiar with the paclet structure in Mathematica.

Your help would be very much appreciated.

Claudio

[mfroeling (Martijn Froeling)]

Hi,

i wasnt sure what the current latest release contained. I just uploaded a new release 1.3 which contains a *.paclet file. download this file and use PacletInstall["fileDir\\UNET-1.3.0.paclet"] which will install the paclet. You can then load the toolbox with <<UNET`.

Some examples can be found in the UNET.nb file in the zip files also in the release.

[drargento (Claudio Argento)]

Thank you!