how to implement training for UNET implementation on Wolfram Neural Net Repository?

Question

My question is with reference to the U-NET implementation present on the Wolfram Neural Net Repository

The construction notebook present on the page (link: http://www.wolframcloud.com/files/1737200a-b043-413c-ad37-477e208472ad?contentDisposition=attachment) contains all the necessary functions for constructing the net. However, it does not contain the procedure for training the neural net.

I am trying to implement a simple training procedure so that I can firstly train the net myself on the same dataset which the net was initially trained on (https://www.dropbox.com/sh/8dcqxlj94fyyop0/AADib7XPcVkJ1PHddD2Nm9Moa?dl=0). Thereafter, I would like to use a different dataset for training.

Please download the construction notebook and the training dataset before proceeding. The only code that I have added to the construction notebook is mentioned below:

(*loading the images, resizing them and augmenting them to produce the training dataset;
background labelled as 1 and cells in the foreground as 2.*)

fnamesimages = Import["C:\\Users\\aliha\\Downloads\\dataset\\images\\"];
ordering = Ordering@Flatten@StringCases[fnamesimages, (p : DigitCharacter ..) ~~ ".tif" :> FromDigits@p];
fnamesimages = fnamesimages[[ordering]];

images = Import["C:\\Users\\aliha\\Downloads\\dataset\\images\\" <> #] &/@fnamesimages;
images = ImageResize[#, {388, 388}] & /@ images;
masks = Import["C:\\Users\\aliha\\Downloads\\dataset\\segmentation\\" <> #] &/@fnamesimages;
allmasks = Flatten@Table[ImageRotate[j, i], {j, masks}, {i, {0, Pi/2, Pi, 3/2 Pi}}];
allmasks = Join[allmasks, ImageReflect /@ allmasks];
maskres = ImageResize[#, {388, 388}] & /@ allmasks;

m = ArrayComponents[ImageData@#, 2, {0. -> 1, n_ /; n != 0. -> 2}] &/@maskres;
allimages = Flatten@Table[ImageRotate[j, i], {j, images}, {i, {0, Pi/2, Pi, 3/2 Pi}}];
allimages = Join[allimages, ImageReflect /@ allimages];

(* using a small subset of images and segmented images because of GPU memory crash*)
trained = NetTrain[unet, allimages[[1 ;; 50]] -> m[[1 ;; 50]], All, BatchSize -> 5, MaxTrainingRounds -> 1, TargetDevice -> "GPU"];
trainedNet = trained["TrainedNet"];

In addition I am using the code in the example notebook (present on the same page) to then evaluate the trained net on a test image.

Clear@netevaluate;
netevaluate[img_, device_ : "CPU"] :=
Block[{net = trainedNet, dims = ImageDimensions[img], pads, mask},
pads = Map[{Floor[#], Ceiling[#]} &, Mod[4 - dims, 16]/2];
mask = NetReplacePart[net,
  {"Input" -> 
    NetEncoder[{"Image", Ceiling[dims - 4, 16] + 188, 
      ColorSpace -> "Grayscale"}],
   "Output" -> 
    NetDecoder[{"Class", Range[2], "InputDepth" -> 3}]}][
 ImagePad[ColorConvert[img, "Grayscale"], pads + 92, 
  Padding -> "Reversed"],
 TargetDevice -> device
 ];
Take[mask, {1, -1} Reverse[pads[[2]] + 1], {1, -1} (pads[[1]] + 1)]
];

we can now load the test image and apply the net.

testimg = Import["C:\\Users\\aliha\\Downloads\\dataset\\test image\\t099.tif];
netevaluate[testimg]//Colorize

Unfortunately I do not get any segmentations back. I just get the background. Could someone kindly let me know where I may be having the issue? Thanks !

Answer 1

Your masks are all black, aren't they? So, it is not possible that U-net can learn anything.

I think the first step, is to use correct masks

Secondly:

increase the MaxTrainingRounds to 1000 or so... or as long as the algorithm is learning.

I worked with Mathematica and U-net 1 year ago with cell data and it works properly.

Answer 2

  imageinputsize = {128, 128};

poolingLayer := PoolingLayer[{2, 2}, 2];

convolution[n_, p_: 0] := NetChain[{If[p == 1, poolingLayer, Nothing],
    ConvolutionLayer[n, 3, "PaddingSize" -> {1, 1}], Ramp, 
    BatchNormalizationLayer[],
    ConvolutionLayer[n, 3, "PaddingSize" -> {1, 1}], Ramp, 
    BatchNormalizationLayer[]
    }];

deconvolution[n_] := NetGraph[{
    "deconvolution" -> 
     DeconvolutionLayer[n, {2, 2}, "Stride" -> {2, 2}],
    "catenate" -> CatenateLayer[],
    "convolution" -> convolution[n]},
   {NetPort["Input1"] -> "catenate", 
    NetPort["Input2"] -> 
     "deconvolution" -> "catenate" -> "convolution"}
   ];

UNET2D[NChannel_: 1, Nclass_: 1] := NetGraph[<|
    "enc_1" -> convolution[64], "enc_2" -> convolution[128, 1], 
    "enc_3" -> convolution[256, 1], "enc_4" -> convolution[512, 1], 
    "enc_5" -> convolution[1024, 1],
    "dec_1" -> deconvolution[512], "dec_2" -> deconvolution[256], 
    "dec_3" -> deconvolution[128], "dec_4" -> deconvolution[64],
    "map" -> {ConvolutionLayer[Nclass, {1, 1}], LogisticSigmoid, 
      If[Nclass > 1, TransposeLayer[{1 <-> 3, 1 <-> 2}], Nothing], 
      If[Nclass > 1, SoftmaxLayer[], Nothing], 
      If[Nclass > 1, Nothing, FlattenLayer[1]]}
    |>,
   {NetPort["Input"] -> 
     "enc_1" -> "enc_2" -> "enc_3" -> "enc_4" -> "enc_5",
    {"enc_4", "enc_5"} -> "dec_1", {"enc_3", "dec_1"} -> "dec_2",
    {"enc_2", "dec_2"} -> "dec_3", {"enc_1", "dec_3"} -> "dec_4", 
    "dec_4" -> "map"},
   "Input" -> {NChannel, imageinputsize[[1]], imageinputsize[[2]]}, 
   "Output" -> 
    If[Nclass > 1, 
     NetDecoder[{"Class", "Labels" -> Range[1, Nclass], 
       "InputDepth" -> 3}], Automatic]
   ];

dir = "C:\\machinelearning\\image segmentation\\Human blast data"

SetDirectory[dir];

images = Partition[
   ImageData[#] & /@ (ColorConvert[#, 
        "Grayscale"] & /@ (ImageResize[Import[dir <> "\\" <> #], 
          imageinputsize] & /@ FileNames[])), 1];

shuffledimages = RandomSample@Thread[Range@Length@images -> images];

keysshuffle = Keys@shuffledimages;

dir = "C:\machinelearning\image segmentation\Human blast masks"

SetDirectory[dir];

masks = ImageData[#] & /@ (ImageResize[Import[dir <> "\\" <> #], 
       imageinputsize] & /@ FileNames[]);

shuffledmasks = 
  Lookup[<|Thread[Range@Length@masks -> masks]|>, keysshuffle];

{dataset, validationset} = 
  TakeList[Values@shuffledimages, {Floor[Length[images] 0.8], 
    Length[images] - Floor[Length[images] 0.8]}];
{labeldataset, labelvalidationset} = 
  TakeList[shuffledmasks, {Floor[Length[images] 0.8], 
    Length[images] - Floor[Length[images] 0.8]}];

trainedNN = 
 NetTrain[UNET2D[1, 1], Thread[dataset -> labeldataset], 
  ValidationSet -> Thread[validationset -> labelvalidationset], 
  MaxTrainingRounds -> 100]

res = Table[Image[trainedNN[dataset[[i]]]], {i, 1, 20}];

m = MorphologicalPerimeter[#, 0.5] & /@ res
m1 = MorphologicalPerimeter[#, 0.99] & /@ res

Grid[Table[{Image[dataset[[i, 1]]], Image[trainedNN[dataset[[i]]]], 
   Image[shuffledmasks[[i]]]}, {i, 10, 15}]]

Table[HighlightImage[
  HighlightImage[
   Image[dataset[[i, 1]]], {Darker[Red], m[[i]]}], {Darker[Blue, 0.7],
    m1[[i]]}], {i, 10, 15}]

and here are the results:

If we train the above U-net model and apply it, it works very well...

Here are the original data sets...

https://lmb.informatik.uni-freiburg.de/resources/datasets/tem.en.html

see human blast data...

I hope this helps...