/leaf-segmentation-unet

a unet model trained for the semantic segmentation of leaf images 🍃

Primary LanguagePythonMIT LicenseMIT

leaf-segmentation-unet

a unet model trained for the semantic segmentation of leaf images

folder structure

.
├── dataset
│   └── note.md
├── model
│   ├── pretrained
│   │   └── download.md
│   ├── __init__.py
│   └── model.py
├── preprocess
│   ├── generate_dataset.py
│   ├── __init__.py
│   └── README.md
├── test
│   └── get_testset.sh
├── utilities
│   ├── __init__.py
│   └── utility.py
├── get_dataset.sh
├── get_pretrained.sh
├── LICENSE
├── predict.py
├── README.md
├── requirements.txt
└── train.py

output samples

1 2 3

input images are from the dataset of Plant Pathology 2021 Challenge

instructions

train from scratch

  • prepare dataset

    linux users can run get_dataset.sh instead of first three steps

    1. cd into ./preprocess

    2. download DenseLeaves.zip from here

    3. unzip the downloaded file as ./preprocess/DenseLeaves/

    4. run python generate_dataset.py

    The newly processed dataset is now saved at ./dataset

  • cd back to project root and run python train.py to train the model

When training, the model saves the weights in the ./model/pretrained model.

  • latest_weights.pth - weights saved at the end of the last epoch
  • best_val_weights.pth - weights saved when the model obtained minimum validation loss

use pretrained (prediction)

  • download pretrained weights

linux users can run get_pretrained.sh (make sure gdown is installed -- pip install gdown)

others can download the weights from links provided in this file

  • specify test image location

edit the TEST_DIR variable in predict.py to specify custom images or download a sample dataset by running get_testset.sh in ./test folder.

tip: if the segmentation results are not satisfactory, modify the mask threshold values in predict.py file

acknowledgement

  • DenseLeaves dataset - Michigan State University visit
  • Plant Pathology 2021 dataset - Kaggle

update log

2021-06-24: first code upload, most of the code is really bad (I wrote them a while ago). I shall refactor them soon.