End-to-end Ultrametric Learning for Hierarchical Segmentation

Python implementation of End-to-end Ultrametric Learning for Hierarchical Segmentation.

Abstract

Hierarchical image segmentation aims to capture the structure of objects of different sizes at different scales and helps to understand the scene. With the success of neural networks for image segmentation and the recent emergence of object and part segmentation datasets, the task of supervised learning of segmentation hierarchies naturally arises to improve performance over semi-supervised and unsupervised approaches. In a previous work, we proposed a differentiable ultrametric layer that transforms any dissimilarity measure into an ultrametric distance equivalent to a hierarchical segmentation.
In this paper, we study several loss functions for optimising ultrametrics in order to obtain a neural network model for end-to-end learning of hierarchical segmentations. In particular, we propose a generalisation of the Rand index to hierarchical segmentations and propose exact and approximate algorithms to compute it. The suitability of the proposed pipeline is demonstrated with several possible loss function combinations on a simulated hierarchical dataset. To evaluate the effectiveness of these optimizations, we introduce two complementary metrics that provide precise and well-ordered hierarchical segmentation assessments. A repository as well as the dataset will be made accessible shortly.

Usage

Download this repository.

Download the Humanoid Dataset and put it in data/humanoid.

(Optional) Download the lightning checkpoints of the experiments of the paper (~8Go compressed, ~8.3Go uncompressed) and put them in lightning_logs. If you do so, you will be able to load them in lightning_logs/<exp_name>/checkpoints/, check their configurations in lightning_logs/<exp_name>/config.yaml and visualize the training process in TensorBoard by installing it (pip install tensorboard), then opening Tensorboard with the following command: tensorboard --logdir lightning_logs/.

(Optional) Download the predictions of the experiments of the paper (~220Mo compressed, ~2.7Go uncompressed) and put them in saved_data/predictions/humanoid/. If you do so, you will be able to visualize the predictions of the experiments of the paper in the corresponding notebook.

(Optional) Download the metrics of the experiments of the paper (~200Mo compressed, ~11Go uncompressed) and put them in saved_data/metrics/humanoid/. If you do so, you will be able to visualize the metrics of the paper in the corresponding notebook.

To train a new neural network, modify config.yaml as you want, then:

• if you are in a cluster using SLURM, modify run-training.sh's SBATCH arguments as well as the MODULE_ENV according to your cluster, then run it with the command sbatch run-training.sh. It will train a neural network and queue the metrics computation, which will queue the metrics combination after it's done.
• if you are not in a cluster, well i gotta add instructions here [TODO].

To replicate the experiments in the paper, modify run-training.sh's CONFIG_FILE argument with the experiment you want to replicate.

To visualize the metrics, open the corresponding notebook, modify the exp_names and friendly_exp_names with your own experiments if you want, and then execute it.

Requirements

• Classical Python scientific packages (numpy, scipy, scikit-image, matplotlib, glob, itertools, pickle)
• Pytorch: optimization library
• TorchVision: vision models library
• PyTorch Lightning: optimization framework for structured and flexible code
• Segmentation Models Pytorch: pretrained segmentation models in PyTorch
• Albumentations: data augmentation library
• Higra: graph library
• ImageIO: image IO library
• ImageSize for reading image size without loading them
• PyYAML: YAML library for configuration files
• IPython for managing MatPlotLib displays

How to cite

bibtex

@InProceedings{10.1007/978-3-031-57793-2_22,
author="Lapertot, Raphael
and Chierchia, Giovanni
and Perret, Benjamin",
editor="Brunetti, Sara
and Frosini, Andrea
and Rinaldi, Simone",
title="End-to-End Ultrametric Learning for Hierarchical Segmentation",
booktitle="Discrete Geometry and Mathematical Morphology",
year="2024",
publisher="Springer Nature Switzerland",
address="Cham",
pages="286--297",
abstract="Hierarchical image segmentation aims to capture the structure of objects of different sizes at different scales and helps to understand the scene. With the success of neural networks for image segmentation and the recent emergence of object and part segmentation datasets, the task of supervised learning of segmentation hierarchies naturally arises. In a previous work, we proposed a differentiable ultrametric layer that transforms any dissimilarity measure into an ultrametric distance equivalent to a hierarchical segmentation. In this paper, we study several loss functions for end-to-end learning of a neural network model predicting hierarchical segmentations. In particular, we propose a generalization of the Rand index for hierarchical segmentation and propose exact and approximate algorithms to compute it. We introduce new metrics to compare hierarchical segmentations, and we demonstrate the suitability of the proposed pipeline with several possible loss function combinations on a simulated hierarchical dataset.",
isbn="978-3-031-57793-2"
}