hthierno/pytorch-stardist

PyTorch implementation of 2D and 3D StarDist - Object Detection with Star-convex Shapes

PyTorch StarDist

This repository contains PyTorch implementations for both StarDist 2D and StarDist 3D as described in:

• Uwe Schmidt, Martin Weigert, Coleman Broaddus, and Gene Myers.
  Cell Detection with Star-convex Polygons.
  International Conference on Medical Image Computing and Computer-Assisted Intervention (MICCAI), Granada, Spain, September 2018.

• Martin Weigert, Uwe Schmidt, Robert Haase, Ko Sugawara, and Gene Myers.
  Star-convex Polyhedra for 3D Object Detection and Segmentation in Microscopy.
  The IEEE Winter Conference on Applications of Computer Vision (WACV), Snowmass Village, Colorado, March 2020

Installation

You should have a C++ compiler installed as this code relies on C/C++ extensions that need to be compiled. This code has been tested with Build Tools for Visual Studio on Windows and GCC on linux.

Follow this step to install pytorch stardist:

1. Download the repo
2. Create a conda environment using the file environment.yml: conda env create --file environment.yml
3. Activate the environment and install the package stardist_tools:

   conda activate pytorch-stardist
   cd stardist_tools_
   python setup.py install
   

Training and Inference

Notebooks

The notebooks at examples/3D and examples/2D show in details how to perform training and inference.

Command line scipts

You can also use command line script.

Let's download some data for the demonstration:

python download_demo_data.py

Training

You need a YAML file containing the training configurations to run the train.py script . Check confs\dsb2018_2d.yml and confs\demo_3d.yml for examples of configuration files.

Run the following command to train the model with the configurations in confs\dsb2018.yml:

python train.py --yaml_conf .\confs\dsb2018.yml

It will train a starDist2D model on a subsample of the Data Science Bowl 2018 dataset.

Inference

We can perform predictions with the trained model using:

python predict.py --name dsb2018 --checkpoints_dir .\checkpoints --image_dir .\datasets\dsb2018\test\images --result_dir .\datasets\dsb2018\preds --use_gpu --use_amp

--name dsb2018 indicates the name of the experiment given in the YAML configuration file used for the training.

Notes

• The code is inspired by pytorch-CycleGAN-and-pix2pix and stardist
• The code in the folder stardist_tools_ is from the StarDist repo.