Wolfram Neural Net Repository

Immediate Computable Access to Neural Net Models

DoubleU-Net Trained on Medical Image Segmentation Datasets

Segment medical images performing a task-dependent binary classification

The authors propose a novel architecture called DoubleU-Net, combining two U-Net blocks stacked on top of each other. The first U-Net works as an encoder, using a VGG-19 pre-tained on the ImageNet dataset, while the second U-Net is added to capture more semantic information. An atrous spatial pyramid pooling (ASPP) block is added to capture contextual information within the network. The architecture was trained over different medical image segmentation datasets covering areas such as colonoscopy, dermoscopy and microscopy. The result outperforms U-Net for all the tasks.

Training Set Information

Model Information

Examples

Resource retrieval

Get the pre-trained net:

In[1]:=
NetModel["DoubleU-Net Trained on Medical Image Segmentation Datasets"]
Out[1]=

NetModel parameters

This model consists of a family of individual nets, each identified by a specific parameter combination. Inspect the available parameters:

In[2]:=
NetModel["DoubleU-Net Trained on Medical Image Segmentation Datasets", "ParametersInformation"]
Out[2]=

Pick a non-default net by specifying the parameters:

In[3]:=
NetModel["DoubleU-Net Trained on Medical Image Segmentation Datasets", "Task" -> "CellNuclei"]
Out[3]=

Pick a non-default uninitialized net:

In[4]:=
NetModel[{"DoubleU-Net Trained on Medical Image Segmentation Datasets", "Task" -> "ColonPolyp"}, "UninitializedEvaluationNet"]
Out[4]=

Evaluation function

Define an evaluation function to select the evaluation mode and postprocess the net output:

In[5]:=
netevaluate[net_, img_, fast_ : False, device_ : "CPU"] := Module[{probs}, probs = If[fast, Flatten[net[img, NetPort[{"FastMask"}], TargetDevice -> device], 1], net[img, TargetDevice -> device] ]; probs = ArrayResample[probs, Reverse@ImageDimensions@img]; Round[probs] ];

Basic usage

Obtain the segmentation mask of a skin lesion:

In[6]:=
(* Evaluate this cell to get the example input *) CloudGet["https://www.wolframcloud.com/obj/dee58408-0861-4a1a-93c5-d504c1eadcbe"]
In[7]:=
mask = netevaluate[ NetModel[ "DoubleU-Net Trained on Medical Image Segmentation Datasets"], img];

Visualize the mask:

In[8]:=
ArrayPlot[mask]
Out[8]=

The mask is a matrix of 0 and 1 whose size matches the dimensions of the input image:

In[9]:=
DeleteDuplicates@Flatten[mask]
Out[9]=
In[10]:=
Dimensions[mask] ImageDimensions[img]
Out[10]=

Overlay the mask on the input image:

In[11]:=
HighlightImage[img, mask]
Out[11]=

Performance tradeoff

The nets compute a lower-quality segmentation mask as an intermediate step. It is possible to stop the computation at such intermediate step to obtain a segmentation mask faster. This can be enabled by setting the third argument of the evaluation function to True:

In[12]:=
(* Evaluate this cell to get the example input *) CloudGet["https://www.wolframcloud.com/obj/768d29f7-1603-42fd-8f6b-07e764075381"]
In[13]:=
fastMask = netevaluate[ NetModel[ "DoubleU-Net Trained on Medical Image Segmentation Datasets"], img, True];

Visualize the mask:

In[14]:=
ArrayPlot[fastMask]
Out[14]=

Overlay the mask on the input image:

In[15]:=
HighlightImage[img, fastMask]
Out[15]=

Compare the evaluation timings of the fast and slow masks:

In[16]:=
{slowTime, slowMask} = RepeatedTiming@ netevaluate[ NetModel[ "DoubleU-Net Trained on Medical Image Segmentation Datasets"], img]; {fastTime, fastMask} = RepeatedTiming@ netevaluate[ NetModel[ "DoubleU-Net Trained on Medical Image Segmentation Datasets"], img, True];
In[17]:=
{slowTime, fastTime}
Out[17]=

Compare the quality of the two segmentation masks. The fast mask is in blue and the slow mask is in red:

In[18]:=
HighlightImage[img, {"Boundary", Blue, Image[fastMask], Red, Image[slowMask]}]
Out[18]=

Net information

Obtain the total number of parameters:

In[19]:=
Information[ NetModel[ "DoubleU-Net Trained on Medical Image Segmentation Datasets"], "ArraysTotalElementCount"]
Out[19]=

Obtain the layer type counts:

In[20]:=
Information[ NetModel[ "DoubleU-Net Trained on Medical Image Segmentation Datasets"], "LayerTypeCounts"]
Out[20]=

Display the summary graphic:

In[21]:=
Information[ NetModel[ "DoubleU-Net Trained on Medical Image Segmentation Datasets"], "SummaryGraphic"]
Out[21]=

Resource History

  • Date Created: 26 December 2022

Reference