ROMI CNN-based segmentation module

Installation

Clone the sources

git clone https://github.com/romi/romiseg

Install PyTorch & dependencies

python -m pip install torch==1.12.1+cu102 torchvision==0.13.1+cu102 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu102

We install it for CUDA 10.2, as it is reasonably old to work with your CUDA toolkit.

Install romiseg

cd romiseg
python -m pip install -e .

Segmentation2D module

This module requires Colmap task to be completed and preferably uses Undistorted images. For each image in the scan dataset, it generates 6 images corresponding to the prediction mask of the classes: Background, Stem, Flower, Fruit, Leaf and Peduncle.

These images are generated by a segmentation neural network. This network has been trained on virtual images of arabidopsis generated with ROMI's blender virtual scanner.

It is possible to annotate manually real images taken with the scanner to improve the segmentation predictions and 3D reconstruction with the Annotation and Fine-Tuning tool.

Annotation and fine-tuning tool

Introduction

This script allows you to fine-tune a pixel-per-pixel segmentation network on your own dataset of images.

The first part of the script allows you to generate this dataset by manually annotating images of your choice using LabelMe.

The second part consists in training a pre-trained network of your choice on your new dataset.

Example of use: ROMI pipeline segmentation. The segmentation networks integrated in the ROMI pipeline were trained on arabidopsis. If you want to reconstruct another plant species, for example from a 3D scan of a tomato, it will hardly work: The present module allows you to annotate manually images from the scan of tomatoes, and re-train the network. Then, launch the pipeline again with the updated network. The results are good enough to allow a 3D reconstruction.

This tool relies on the file pipeline.toml used by the virtual scanner.

Preliminary set-up

First update the .toml file of the luigi task Segmentation2D according to the one in this directory.

The code will mount db.romi-project.eu locally to allow saving the annotations directly on the shared database, and to collect tall the annotated images from the database to do the training. The weights will be saved in the cache, the mounted database too, as well as the tensorboard log.

[Segmentation2D]
upstream_task = "Scan"
query = "{\"channel\":\"rgb\"}"
labels = "background,flower,peduncle,stem,bud,leaf,fruit"
model_name = "Resnet101"
model_segmentation_name = "Resnet101_896_896_epoch51.pt" #model to finetune
Sx = 896
Sy = 896
learning_rate = 0.0001

[Finetune]
finetune_epochs = 10
batch = 1

If the pre-trained network is not already present in the cache, it will be fetched from the database db.romi-project.eu.

You can also change the number of epochs for the training.

Running the tool

finetune.py --config [/path/to/config_file]

Set-by-step

First you will have to select images you want to annotate. 3 images should be enough.

Then, LabelMe will pop up with the first selected image. You can annotate the image with the "Create Polygon" function. The classes are flower, stem, leaf, fruit, peduncle and are already set-up in LabelMe.

Save the .json label file in the suggested folder (it corresponds to directory_images where the images you just selected have also been copied).

When you close LabelMe the next image you have selected will automatically pop up.

When all the images are labelled a pop-up image will show a random sample of the training dataset you have generated. Close this window and press enter.

The training will automatically start for the number of epochs indicated in pipeline.toml.

Once the training is over, the weights are automatically saved in the weights folder and the pipeline.toml file is updated with the name of the fine_tuned segmentation network.

You can launch the plant-3d-vision pipeline on the scan and the result should be better than before the fine-tuning. The name of the new network is automatically updated.