Semantic segmentation

The goal of this workflow is assign a class to each pixel of the input image.

• Input:

  • Image.

  • Class mask where each pixel is labeled with an integer representing a class.

• Output:

  • Image with the probability of being part of each class.

In the figure below an example of this workflow’s input is depicted. There, only two labels are present in the mask: black pixels, with value 0, represent the background and white ones the mitochondria, labeled with 1. The number of classes is defined by MODEL.N_CLASSES variable.

Input image.

Input class mask (ground truth).

The output in case that only two classes are present, as in this example, will be an image where each pixel will have the probability of being of class 1.

If there are 3 or more classes, the output will be a multi-channel image, with the same number of channels as classes, and the same pixel in each channel will be the probability (in [0-1] range) of being of the class that represents that channel number. For instance, with 3 classes, e.g. background, mitochondria and contours, the fist channel will represent background, the second mitochondria and the last contour class.

Data preparation

To ensure the proper operation of the library the data directory tree should be something like this:

dataset/
├── train
│   ├── x
│   │   ├── training-0001.tif
│   │   ├── training-0002.tif
│   │   ├── . . .
│   │   ├── training-9999.tif
│   └── y
│       ├── training_groundtruth-0001.tif
│       ├── training_groundtruth-0002.tif
│       ├── . . .
│       ├── training_groundtruth-9999.tif
└── test
    ├── x
    │   ├── testing-0001.tif
    │   ├── testing-0002.tif
    │   ├── . . .
    │   ├── testing-9999.tif
    └── y
        ├── testing_groundtruth-0001.tif
        ├── testing_groundtruth-0002.tif
        ├── . . .
        ├── testing_groundtruth-9999.tif

Warning

Ensure that images and their corresponding masks are sorted in the same way. A common approach is to fill with zeros the image number added to the filenames (as in the example).

Configuration

Find in templates/semantic_segmentation folder of BiaPy a few YAML configuration templates for this workflow.

Special workflow configuration

Here some special configuration options that can be selected in this workflow are described:

• Data loading: if you want to select DATA.EXTRACT_RANDOM_PATCH you can also set DATA.PROBABILITY_MAP to create a probability map so the patches extracted will have a high probability of having an object in the middle of it. Useful to avoid extracting patches which no foreground class information. That map will be saved in PATHS.PROB_MAP_DIR.

Furthermore, when this is enabled, in PATHS.DA_SAMPLES path, i.e. aug folder by default (see Results), two more images will be created so you can check how this probability map is working. These images will have painted a blue square and a red point in its middle, which correspond to the patch area extracted and the central point selected respectively. One image will be named as mark_x and the other one as mark_y, which correspond to the input image and ground truth respectively.

• Metrics: during the inference phase the performance of the test data is measured using different metrics if test masks were provided (i.e. ground truth) and, consequently, DATA.TEST.LOAD_GT is enabled. In the case of semantic segmentation the Intersection over Union (IoU) metrics is calculated. This metric, also referred as the Jaccard index, is essentially a method to quantify the percent of overlap between the target mask and the prediction output. Depending on the configuration different values are calculated (as explained in Test phase). This values can vary a lot as stated in [FBMunozBAC21].

  • per patch values are calculated if TEST.STATS.PER_PATCH is enabled. IoU is calculated for each patch separately and then averaged.

  • merge patches values are calculated if TEST.STATS.MERGE_PATCHES is enabled. Notice that depending on the amount of overlap/padding selected the merged image can be different than just concatenating each patch.

  • full image values are calculated if TEST.STATS.FULL_IMG is enabled. This can be done if the model selected if fully convolutional. The results may be slightly different from merge patches as you may notice and probably no border effect will be seen.

• Post-processing: When PROBLEM.NDIM is 2D the post-processing will be enabled only if TEST.STATS.FULL_IMG is enabled. In that case the post-processing will process all 2D predicted images as a unique 3D stack. On the other hand, when PROBLEM.NDIM is 3D the post-processing will be applied when TEST.STATS.PER_PATCH and TEST.STATS.MERGE_PATCHES is selected. In this case, each 3D predicted image will be processed individually.

  • Z-filtering: to apply a median filtering in z axis. Useful to maintain class coherence across 3D volumes. Enable it with TEST.POST_PROCESSING.Z_FILTERING and use TEST.POST_PROCESSING.Z_FILTERING_SIZE for the size of the median filter.

  • YZ-filtering: to apply a median filtering in y and z axes. Useful to maintain class coherence across 3D volumes that can work slightly better than Z-filtering. Enable it with TEST.POST_PROCESSING.YZ_FILTERING and use TEST.POST_PROCESSING.YZ_FILTERING_SIZE for the size of the median filter.

Run

Command lineDockerGoogle Colab

Open a terminal as described in Installation. For instance, using 2d_semantic_segmentation.yaml template file, the code can be run as follows:

# Configuration file
job_cfg_file=/home/user/2d_semantic_segmentation.yaml
# Where the experiment output directory should be created
result_dir=/home/user/exp_results
# Just a name for the job
job_name=my_2d_semantic_segmentation
# Number that should be increased when one need to run the same job multiple times (reproducibility)
job_counter=1
# Number of the GPU to run the job in (according to 'nvidia-smi' command)
gpu_number=0

# Move where BiaPy installation resides
cd BiaPy

# Load the environment
conda activate BiaPy_env
source $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh

python -u main.py \
    --config $job_cfg_file \
    --result_dir $result_dir  \
    --name $job_name    \
    --run_id $job_counter  \
    --gpu $gpu_number

For multi-GPU training you can call BiaPy as follows:

gpu_number="0, 1, 2"
python -u -m torch.distributed.run \
    --nproc_per_node=3 \
    main.py \
    --config $job_cfg_file \
    --result_dir $result_dir  \
    --name $job_name    \
    --run_id $job_counter  \
    --gpu $gpu_number

nproc_per_node need to be equal to the number of GPUs you are using (e.g. gpu_number length).

Results

The results are placed in results folder under --result_dir directory with the --name given. An example of this workflow is depicted below:

Example of semantic segmentation model predictions. From left to right: input image, its mask and the overlap between the mask and the model’s output binarized.

Following the example, you should see that the directory /home/user/exp_results/my_2d_semantic_segmentation has been created. If the same experiment is run 5 times, varying --run_id argument only, you should find the following directory tree:

my_2d_semantic_segmentation/
├── config_files/
│   └── my_2d_semantic_segmentation_1.yaml
├── checkpoints
│   └── my_2d_semantic_segmentation_1-checkpoint-best.pth
└── results
    ├── my_2d_semantic_segmentation_1
    ├── . . .
    └── my_2d_semantic_segmentation_5
        ├── aug
        │   └── .tif files
        ├── charts
        │   ├── my_2d_semantic_segmentation_1_*.png
        │   ├── my_2d_semantic_segmentation_1_loss.png
        │   └── model_plot_my_2d_semantic_segmentation_1.png
        ├── full_image
        │   └── .tif files
        ├── full_image_binarized
        │   └── .tif files
        ├── full_post_processing
        │   └── .tif files
        ├── per_image
        │   └── .tif files
        ├── per_image_binarized
        │   └── .tif files
        ├── tensorboard
        └── train_logs

• config_files: directory where the .yaml filed used in the experiment is stored.

  • my_2d_semantic_segmentation.yaml: YAML configuration file used (it will be overwrited every time the code is run)

• checkpoints: directory where model’s weights are stored.

  • my_2d_semantic_segmentation_1-checkpoint-best.pth: checkpoint file (best in validation) where the model’s weights are stored among other information.

• results: directory where all the generated checks and results will be stored. There, one folder per each run are going to be placed.

  • my_2d_semantic_segmentation_1: run 1 experiment folder.

    • aug: image augmentation samples.

    • charts:

      • my_2d_semantic_segmentation_1_*.png: Plot of each metric used during training.

      • my_2d_semantic_segmentation_1_loss.png: Loss over epochs plot (when training is done).

      • model_plot_my_2d_semantic_segmentation_1.png: plot of the model.

    • full_image:

      • .tif files: output of the model when feeding entire images (without patching).

    • full_image_binarized:

      • .tif files: Same as full_image but with the image binarized.

    • full_post_processing (optional if any post-processing was selected):

      • .tif files: output of the model when feeding entire images (without patching) and applying post-processing, which in this case only y and z axes filtering was selected.

    • per_image:

      • .tif files: reconstructed images from patches.

    • per_image_binarized:

      • .tif files: Same as per_image but with the images binarized.

    • tensorboard: Tensorboard logs.

    • train_logs: each row represents a summary of each epoch stats. Only avaialable if training was done.

Note

Here, for visualization purposes, only my_2d_semantic_segmentation_1 has been described but my_2d_semantic_segmentation_2, my_2d_semantic_segmentation_3, my_2d_semantic_segmentation_4 and my_2d_semantic_segmentation_5 will follow the same structure.