Skip to content

johschmidt42/PyTorch-2D-3D-UNet-Tutorial

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

49 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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".

image1 image2

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.

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.

About

No description, website, or topics provided.

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published