Compile Unet based on pytorch for cell structure segmentation

Paper address: U-Net: Convolutional Networks for Biomedical Image Segmentation
Code address: https://github.com/laisimiao/Unet
Unet must have heard its name, a very famous paper in the direction of medical image segmentation, I Also with curiosity, motivated by my desire to have a preliminary understanding of segmentation tasks, I realized the challenge mentioned in the paper (less training data and short training time): ISBI ​​Challenge: Segmentation of neuronal structures in EM stacks , about this data set: this is what the paper says:

The training data is a set of 30 images (512x512 pixels) from serial section transmission electron microscopy of the Drosophila first instar larva ventral nerve cord (VNC). The general idea
is: there are 30 training data with a resolution of 512x512. These pictures are fruit Electron micrograph of the fly.

Look at the following animation, you can understand what task to do: give a cell structure diagram, we have to separate it from each other.
Insert picture description here
In fact, the segmentation task I understand is to classify each pixel, so its label is also the same as the resolution of the input image. For the fruit fly cell map here, it is just fine to divide it into two categories : the cell membrane is black. The cell is white, so the feature map shape of our prediction result can be 2x512x512 (the same for VOC and COCO multi-classification, the feature map of CxHxW shape can be predicted, C is the number of target categories, H and W are respectively the input image Height and width), where the first channel is the mask at the cell membrane, and the second channel is the mask inside the cell. Later, I found that for this two-classification problem, only one channel is actually output, that is, 1x512x512, because at the end, except for black It is white, and the original image is very close to the grayscale image, there is no big difference after being converted to a single channel, so the input and output are all single channel, and the calculation of pixel-wise logistic loss is also very convenient for code implementation.
Putting a Unet network structure diagram, it is relatively easy to implement with pytorch , but the resolution of the contracting path and the expanding path implemented here are the same, instead of the contracting path crop and the expanding path as in the paper. Concat:
Unet network structure diagram
write a little understanding and not too clear:

  • After the network is trained, it is like a texture detector. After inputting a test picture, its activation value at the cell boundary is large, and after the sigmoid function, it tends to 1, which is black; the activation value in other places It is low, and after the sigmoid function, it tends to 0, which is white.
  • For images in .tif format, it seems that only Image and TIFF can be used to read in. But once I read in this cell image, the mode of Image is L , which is grayscale, so I don’t know. It turned out to be 3-channel or 1-channel, but this does not affect the later

For the specific implementation, you can see the code . Two versions of keras 2.2.4 and pytorch 1.0 are provided. During actual training, it was found that when the batch size of training in the keras version was 2, the GPU memory occupied by about 10GB, but only in pytorch It took about 4GB, because the decode part of the pytorch version is achieved by upsampling rather than transposed convolution, here are two renderings:
This is the result image of the keras version
This is the result picture of the pytorch version
attach a link to the evaluation index dice coefficient: Dice Loss of medical image segmentation

Guess you like

Origin blog.csdn.net/laizi_laizi/article/details/103863756