Application of the U-Net convolutional neural network in the segmentation and classification of regions in white blood cells using the Raabin WBC Data dataset.
Using the PyTorch language, the architecture of the original U-Net was developed based on the article "U-Net: Convolutional Networks for Biomedical Image Segmentation" by Olaf Ronneberger et al (2015), according to Figure 1 presented in this article.
Composition of the dataset used:
- Total of 1145 images;
- 218 images of basophil cells;
- 201 images of eosinophil cells;
- 242 images of lymphocyte cells;
- 242 images of monocyte cells;
- 242 images of neutrophil cells.
You can download the dataset from Raabin WBC Data .
Dataset split: 85% for training and 15% for validation (there was no test).
The training environment was on Google Colab, using the free resources available and the possibility of allocation in the GPU.
Training parameters:
- 30 epochs;
- Batch size: 5;
- Learning Rate: 0.0001;
- Optimizer: optm.Adam, in PyTorch;
- Loss function: nn.BCEWithLogitsLoss, in PyTorch;
- Data Augmentation: Rotation, Horizontal Flip, Vertical Flip, and Perspective.
The result of the training with 30 epochs can be observed in the following graph:
As metrics to evaluate the performance of the network, Accuracy and Dice Score were used, reaching 99.34% of Accuracy and 97.16% of Dice Score in the end.
All the resulting images from the training can be accessed in the "results" folder, but here are the results for eosinophil cells:
Hope you enjoy it!
Diego Oliveira