SelfPretrain-SegFinetune

Pre-text Pre-training -> Segmentation Fine-tuning pipeline

Developed with the software and tools below.

Table of Contents

• Overview
• Repository Structure
• Modules
• Getting Started
  • Installation
  • Data and Data Download
  • Experiments Reproduction
    • 1, Data Similarity
    • 2, Fine-tuning Data Size
    • 3, Baseline Model
• Acknowledgments

---

Overview

The SelfPretrain-SegFinetune project is designed to optimize image segmentation models through self-supervised learning techniques, including extensive data augmentation and contrastive learning techniques. The project includes autoencoder with ViT as encoder and MLP as decoder and leverages contrastive learning durning pre-training. Core experiments within the project assess performance variations using diverse training datasets and different piplines.

Example Result

---

Repository Structure

└── SelfPretrain-SegFinetune/
    ├── data_augmentation.py
    ├── data_utils.py
    ├── instruction.pdf
    ├── losses.py
    ├── metrics.py
    ├── models.py
    ├── requirements.txt
    ├── utils.py
    ├── experiments
    │   ├── compare_data_similarity_for_segmentation.ipynb
    │   ├── compare_data_similarity_for_segmentation.py
    │   ├── compare_pretrained_model_finetuning_sizes.ipynb
    │   ├── compare_pretrained_model_finetuning_sizes.py
    │   ├── compare_pretrained_with_baseline_transfer.ipynb
    │   ├── compare_pretrained_with_baseline.ipynb
    │   └── compare_pretrained_with_baseline.py
    ├── Images
    │   ├── compare_data_similarity_for_segmentation
    │   ├── compare_pretrained_model_finetuning_sizes
    │   └── compare_pretrained_with_baseline
    ├── datasets
    │   ├── data
    │   └── 102flowers

---

Modules

.

File	Summary
data_augmentation.py	Enables the enhancement of pre-training image datasets by applying various transformations and augmentations, supporting improved model generalization within the project, crucial for experiments involving pre-trained model evaluations and segmentation tasks.
data_utils.py	Establishes tools for handling image datasets crucial for contrastive learning, including classes for loading and transforming both augmented and original images. It provisions functionality for ensuring image correct format in both training and testing phases.
losses.py	Introduces specialized loss functions, including ContrastiveLoss for pre-training, and DiceLoss for fine-tuning.
metrics.py	Defines functions for model evaluation, specifically computing pixel-wise accuracy and Intersection over Union (IoU) score.
models.py	Defines models essential for pre-text and segmentation tasks, including an Autoencoder and a Decoder.
requirements.txt	Specifies the necessary libraries for this project.
utils.py	Facilitates visualization of segmentation results within the SelfSup-SegFinetune project, enhancing comparison between ground-truth and predicted segmentation masks.

experiments

This folder contains the experiments in this project, in both python scripts and notebooks. All the results are shown in these notebooks.

File	Summary
compare_data_similarity_for_segmentation.ipynb	This notebook contains the code to explore how similar the pre-training and fine-tuning data need to be for better segmentation performance, as well as the results.
compare_data_similarity_for_segmentation.py	This python script contains the code to explore how similar the pre-training and fine-tuning data need to be for better segmentation performance.
compare_pretrained_model_finetuning_sizes.ipynb	This notebook contains the code to explore the effects of the fine-tuning data sized on the segmentation performance, as well as the results.
compare_pretrained_model_finetuning_sizes.py	This python script contains the code to explore the effects of the fine-tuning data sizes on the segmentation performance.
compare_pretrained_with_baseline_transfer.ipynb	This notebook contains code to compare the pre-text pre-training -> segmentation fine-tuning model with the model using fully supervised methods. Additionally, the pre-trained weights of vit_b_16 downloaded from pytorch is directly used. In other words, pre-training is not down in this experiment.
compare_pretrained_with_baseline.ipynb	This notebook contains code to compare the pre-text pre-training -> segmentation fine-tuning model with the model using fully supervised methods, as well as the results.
compare_pretrained_with_baseline.py	This python script contains code to compare the pre-text pre-training -> segmentation fine-tuning model with the model using fully supervised methods.

---

Getting Started

System Requirements:

• Python: version 3.10.8

Installation

From source

1. Clone the repository:

$ git clone https://github.com/kangchengX/SelfPretrain-SegFinetune.git

2. Change to the project directory:

$ cd SelfPretrain-SegFinetune

3. Install the dependencies:

$ pip install -r requirements.txt

Data and Data Download

Fine-tuning-data: This Oxford-IIIT Pet Dataset data for fine-tuning across different experiments. This doesn't need to be manually downloaded, since this is handled in the code.

Pre-training-data-relevant: This dataset contains pets-related data. This needs to be mannally downloaded to datasets/data. We use this data during pre-training. The augmented-data, i.e., data after augmentation, is also available.

Pre-training-data-irrelevant: This dataset contains flower images. This needs to be mannally downloaded to datasets/102flowers.

Experiments Reproduction

Three experiments were conducted, which can be reproduced by either python script or notebook.

1, Data Similarity

The purpose of this experiment is to explore how similar the pre-training and fine-tuning data need to be for better segmentation performance. All of the three datasets - Fine-tuning-data, Pre-training-data-relevant and Pre-training-data-irrelevant are used.

From source

Use python script

Run the command below:

$ python experiments/compare_data_similarity_for_segmentation.py

The visualization results will be stored in Images/compare_data_similarity_for_segmentation

Use Jupyter Notebook

Run the cells in experiments/compare_data_similarity_for_segmentation.ipynb

All the results will be shown in the notebook.

2, Fine-tuning Data Size

The purpose of this experiment is to explore the effects of the fine-tuning data sizes on the segmentation performance. Datasets Fine-tuning-data and Pre-training-data-relevant are used.

From source

Use python script

Run the command below:

$ python experiments/compare_pretrained_model_finetuning_sizes.py

The visualization results will be stored in Images/compare_pretrained_model_finetuning_sizes

Use Jupyter Notebook

Run the cells in experiments/compare_pretrained_model_finetuning_sizes.ipynb

All the results will be shown in the notebook.

3, Baseline Model

The purpose of this experiment is to compare the pre-text pre-training -> segmentation fine-tuning model with the model using fully supervised methods. Datasets Fine-tuning-data and Pre-training-data-relevant are used.

From source

Use python script

Run the command below:

$ python experiments/compare_pretrained_with_baseline.py

The visualization results will be stored in Images/compare_pretrained_with_baseline

Use Jupyter Notebook

Run the cells in experiments/compare_pretrained_with_baseline.ipynb

All the results will be shown in the notebook.

---

Acknowledgments

Thanks all the 8 members of our team - Gao, Yuefeng; Shi, Yunfan; Wu, Yifan; Xu, Kangcheng; Xue, Hongjia; Zhang, Ze; Zhuang, Yan; Zuo, Hancheng.

---