coolzhao / pytorch-unet

PyTorch-2D-3D-UNet-Tutorial

A beginner-friendly tutorial to start a 2D or 3D image segmentation deep learning project with PyTorch & the U-Net architecture. Based on the blog series "Creating and training a U-Net model with PyTorch for 2D & 3D semantic segmentation - A guide to semantic segmentation with PyTorch and the U-Net".

Installation

1. Set up a new environment with an environment manager (recommended):
   1. conda:
      1. conda create --name unet-tutorial -y
      2. conda activate unet-tutorial
      3. conda install python=3.8 -y
   2. venv:
      1. python3 -m venv unet-tutorial
      2. source unet-tutorial/bin/activate
2. Install the libraries: pip install -r requirements.txt
3. Start a jupyter server: jupyter-notebook OR jupyter-lab

Note: This will install the CPU-version of torch. If you want to use a GPU or TPU, please refer to the instructions on the PyTorch website

Summary

• Part I: Building a dataset in PyTorch & visualizing it with napari
• Part II: Creating the U-Net model in PyTorch & information about model input and output
• Part III: Training a 2D U-Net model on a sample of the Carvana dataset with improving datasets (caching, multiprocessing)
• Part IV: Running inference on test data
• Part V: Building a 3D dataset
• Part VI: Running an experiment with the experiment tracker neptune.ai and the high-level PyTorch library PyTorch Ligthning

Note: Due to updates in the neptune API, part VI will probably not work and needs to be migrated to the new neptune API first

Dataset samples

This repository contains a sample of a 2D and a 3D dataset for semantic segmentation.

• 2D: Carvana challenge
• 3D: Electron tomograms showing Microtubules from trypanosomes

U-Net model

If you are unsure what arguments to pass in the Unet class, please take a look at the enums in unet.py and view examples in test_unet.py.

Note: Weights are initialized randomly (xavier initialization) and training can take some time. To faster train a segmentation model, it is recommended to use a pretrained backbone such as ResNet for 2D or even 3D tasks.