A comprehensive tutorial on how to implement and train variational UNet based ousing PyTorch
-
Go to the cell tracking chanllenge website to download the HeLa cells on a flat glass training and test dataset. The dataset can be downloaded and unzipped manually or use the PythonDownloadAndUnzip notebook to download programmably.
-
Run the Preprocess notebook to perform erosion and spatial weight calculation preprocessing on the dataset.
-
Run the TrainSimpleUNetWithWeight notebook or TrainSimpleUNetWithoutWeight notebook to train UNet with or without spatial weighted loss.
-
Run the DirectInference notebook or OverlapTileInference notebook to segment new (larger) image using the trained UNet model through direct inference or overlap tile strategy.
-
Run the Evaluation notebook to calculate intersection over union (IoU) of the trained UNet model.
Preprocessing result:
UNet model trained without weighted loss function segmentation result (valdiation set IoU = 85.36%):
UNet model trained with weighted loss function segmentation result (valdiation set IoU = 85.61%):
Overlap tile strategy implementation result:
This repo has been implemented and tested on the following dependencies:
- Python 3.10.13
- matplotlib 3.8.2
- numpy 1.26.2
- torch 2.1.1+cu118
- torchvision 0.16.1+cu118
- notebook 7.0.6
- opencv-python 4.10.0.84
This repo has been tested on a laptop computer with the following specs:
- CPU: Intel(R) Core(TM) i7-9750H CPU
- Memory: 32GB
- GPU: NVIDIA GeForce RTX 2060
[1] Ronneberger, O., Fischer, P. & Brox, T. U-NET: Convolutional Networks for Biomedical Image Segmentation. in Lecture notes in computer science 234–241 (2015). doi:10.1007/978-3-319-24574-4_28.
[2] Maska, M., (...), de Solorzano, C.O.: A benchmark for comparison of cell tracking algorithms. Bioinformatics 30, 1609-1617 (2014)
[1] Cell tracking challenge website URL: https://celltrackingchallenge.net/