pytorch-unet-segmentation
Members : PyeongEun Kim, JuHyung Lee, MiJeong Lee
Supervisors : Utku Ozbulak, Wesley De Neve
Description
This project aims to implement biomedical image segmentation with the use of U-Net model. The below image briefly explains the output we want:
The dataset we used is Transmission Electron Microscopy (ssTEM) data set of the Drosophila first instar larva ventral nerve cord (VNC), which is dowloaded from [ISBI Challenge: Segmentation of of neural structures in EM stacks](http://brainiac2.mit.edu/isbi_challenge/home)
The dataset contains 30 images (.png) of size 512x512 for each train, train-labels and test.
## Table of Content
* [Dataset](#dataset)
* [Preprocessing](#preprocessing)
* [Model](#model)
* [Loss function](#lossfunction)
* [Post-processing](#postprocessing)
* [Results](#results)
* [Dependency](#dependency)
* [References](#references)
## Dataset
```ruby
class SEMDataTrain(Dataset):
def __init__(self, image_path, mask_path, in_size=572, out_size=388):
"""
Args:
image_path (str): the path where the image is located
mask_path (str): the path where the mask is located
option (str): decide which dataset to import
"""
# All file names
# Lists of image path and list of labels
# Calculate len
# Calculate mean and stdev
def __getitem__(self, index):
"""Get specific data corresponding to the index
Args:
index (int): index of the data
Returns:
Tensor: specific data on index which is converted to Tensor
"""
"""
# GET IMAGE
"""
#Augmentation on image
# Flip
# Gaussian_noise
# Uniform_noise
# Brightness
# Elastic distort {0: distort, 1:no distort}
# Crop the image
# Pad the image
# Sanity Check for Cropped image
# Normalize the image
# Add additional dimension
# Convert numpy array to tensor
#Augmentation on mask
# Flip same way with image
# Elastic distort same way with image
# Crop the same part that was cropped on image
# Sanity Check
# Normalize the mask to 0 and 1
# Add additional dimension
# Convert numpy array to tensor
return (img_as_tensor, msk_as_tensor)
def __len__(self):
"""
Returns:
length (int): length of the data
"""
```
## Preprocessing
We preprocessed the images for data augmentation. Following preprocessing are :
* Flip
* Gaussian noise
* Uniform noise
* Brightness
* Elastic deformation
* Crop
* Pad
#### Image Augmentation
Original Image
| Image |
| Flip |
Vertical |
Horizontal |
Both |
| Gaussian noise |
Standard Deviation: 10 |
Standard Deviation: 50 |
Standard Deviation: 100 |
| Uniform noise |
Intensity: 10 |
Intensity: 50 |
Intensity: 100 |
| Brightness |
Intensity: 10 |
Intensity: 20 |
Intensity: 30 |
| Elastic deformation |
Random Deformation: 1 |
Random Deformation: 2 |
Random Deformation: 3 |
#### Crop and Pad
| Crop |
Left Bottom |
Left Top |
Right Bottom |
Right Top |
Padding process is compulsory after the cropping process as the image has to fit the input size of the U-Net model.
In terms of the padding method, **symmetric padding** was done in which the pad is the reflection of the vector mirrored along the edge of the array. We selected the symmetric padding over several other padding options because it reduces the loss the most.
To help with observation, a  'yellow border' is added around the original image: outside the border indicates symmetric padding whereas inside indicates the original image.
| Pad |
Left Bottom |
Left Top |
Right bottom |
Right Top |
## Model
#### Architecture
We have same structure as U-Net Model architecture but we made a small modification to make the model smaller.
## Loss function
We used a loss function where pixel-wise softmax is combined with cross entropy.
#### Softmax
.png)
#### Cross entropy
.png)
## Post-processing
In attempt of reducing the loss, we did a post-processing on the prediction results. We applied the concept of watershed segmentation in order to point out the certain foreground regions and remove regions in the prediction image which seem to be noises.
The numbered images in the figure above indicates the stpes we took in the post-processing. To name those steps in slightly more detail:
* 1. Convertion into grayscale
* 2. Conversion into binary image
* 3. Morphological transformation: Closing
* 4. Determination of the certain background
* 5. Calculation of the distance
* 6. Determination of the certain foreground
* 7. Determination of the unknown region
* 8. Application of watershed
* 9. Determination of the final result
### Conversion into grayscale
The first step is there just in case the input image has more than 1 color channel (e.g. RGB image has 3 channels)
### Conversion into binary image
Convert the gray-scale image into binary image by processing the image with a threshold value: pixels equal to or lower than 127 will be pushed down to 0 and greater will be pushed up to 255. Such process is compulsory as later transformation processes takes in binary images.
### Morphological transformation: Closing.
We used **morphologyEX()** function in cv2 module which removes black noises (background) within white regions (foreground).
### Determination of the certain background
We used **dilate()** function in cv2 module which emphasizes/increases the white region (foreground). By doing so, we connect detached white regions together - for example, connecting detached cell membranes together - to make ensure the background region.
### Caculation of the distance
This step labels the foreground with a color code:  red color indicates farthest from the background while  blue color indicates closest to the background.
### Determination of the foreground
Now that we have an idea of how far the foreground is from the background, we apply a threshold value to decide which part could surely be the foreground.
The threshold value is the maximum distance (calculated from the previous step) multiplied by a hyper-parameter that we have to manually tune. The greater the hyper-parameter value, the greater the threshold value, and therefore we will get less area of certain foreground.
### Determination of the unknown region
From previous steps, we determined sure foreground and background regions. The rest will be classified as *'unknown'* regions.
### Label the foreground: markers
We applied **connectedComponents()** function from the cv2 module on the foreground to label the foreground regions with color to distinguish different foreground objects. We named it as a 'marker'.
### Application of watershed and Determination of the final result
After applying **watershed()** function from cv2 module on the marker, we obtained an array of -1, 1, and many others.
* -1 = Border region that distinguishes foreground and background
* 1 = Background region
To see the result, we created a clean white page of the same size with the input image. then we copied all the values from the watershed result to the white page except 1, which means that we excluded the background.
## Results
| Optimizer |
Learning Rate |
Lowest Loss |
Epoch |
Highest Accuracy |
Epoch |
| SGD |
0.001 |
0.196972 |
1445 |
0.921032 |
1855 |
| 0.005 |
0.205802 |
1815 |
0.918425 |
1795 |
| 0.01 |
0.193328 |
450 |
0.922908 |
450 |
| RMS_prop |
0.0001 |
0.203431 |
185 |
0.924543 |
230 |
| 0.0002 |
0.193456 |
270 |
0.926245 |
500 |
| 0.001 |
0.268246 |
1655 |
0.882229 |
1915 |
| Adam |
0.0001 |
0.194180 |
140 |
0.924470 |
300 |
| 0.0005 |
0.185212 |
135 |
0.925519 |
135 |
| 0.001 |
0.222277 |
165 |
0.912364 |
180 |
We chose the best learning rate that fits the optimizer based on **how fast the model converges to the lowest error**. In other word, the learning rate should make model to reach optimal solution in shortest epoch repeated. However, the intersting fact was that the epochs of lowest loss and highest accuracy were not corresponding. This might be due to the nature of loss function (Loss function is log scale, thus an extreme deviation might occur). For example, if the softmax probability of one pixel is 0.001, then the -log(0.001) would be 1000 which is a huge value that contributes to loss.
For consistency, we chose to focus on accuracy as our criterion of correctness of model.
| Accuracy and Loss Graph |
 |
 |
 |
SGD
(lr=0.01,momentum=0.99) |
RMS prop
(lr=0.0002) |
Adam
(lr=0.0005) |
We used two different optimizers (SGD, RMS PROP, and Adam). In case of SGD the momentum is manually set (0.99) whereas in case of other optimizers (RMS Prop and Adam) it is calculated automatically.
### Model Downloads
Model trained with SGD can be downloaded via **dropbox**:
https://www.dropbox.com/s/ge9654nhgv1namr/model_epoch_2290.pwf?dl=0
Model trained with RMS prop can be downloaded via **dropbox**:
https://www.dropbox.com/s/cdwltzhbs3tiiwb/model_epoch_440.pwf?dl=0
Model trained with Adam can be downloaded via **dropbox**:
https://www.dropbox.com/s/tpch6u41jrdgswk/model_epoch_100.pwf?dl=0
### Example
Input Image
| Results comparsion |
 |
 |
 |
 |
| original image mask |
RMS prop optimizer
(Accuracy 92.48 %) |
SGD optimizer
(Accuracy 91.52 %) |
Adam optimizer
(Accuracy 92.55 %) |
## Dependency
Following modules are used in the project:
* python >= 3.6
* numpy >= 1.14.5
* torch >= 0.4.0
* PIL >= 5.2.0
* scipy >= 1.1.0
* matplotlib >= 2.2.2
## References
:
[1] O. Ronneberger, P. Fischer, and T. Brox. U-Net: Convolutional Networks for Biomedical Image Segmentation, http://arxiv.org/pdf/1505.04597.pdf
[2] P.Y. Simard, D. Steinkraus, J.C. Platt. Best Practices for Convolutional Neural Networks Applied to Visual Document Analysis, http://cognitivemedium.com/assets/rmnist/Simard.pdf
