💎 Installation with pip

Installation is as simple as running:

pip install git+https://github.com/wiktorlazarski/head-segmentation.git

🔨 How to use

🤔 Inference

import head_segmentation.segmentation_pipeline as seg_pipeline

segmentation_pipeline = seg_pipeline.HumanHeadSegmentationPipeline()

segmentation_map = segmentation_pipeline.predict(image)

🎨 Visualizing

import matplotlib.pyplot as plt

import head_segmentation.visualization as vis

visualizer = vis.VisualizationModule()

figure, _ = visualizer.visualize_prediction(image, segmentation_map)
plt.show()

⚙️ Setup for development

# Clone repo
git clone https://github.com/wiktorlazarski/head-segmentation.git

# Go to repo directory
cd head-segmentation

# (Optional) Create virtual environment
python -m venv venv
source ./venv/bin/activate

# Install project in editable mode
pip install -e .[dev]

# (Optional but recommended) Install pre-commit hooks to preserve code format consistency
pre-commit install

🐍 Setup for development with Anaconda or Miniconda

# Clone repo
git clone https://github.com/wiktorlazarski/head-segmentation.git

# Go to repo directory
cd head-segmentation

# Create and activate conda environment
conda env create -f ./conda_env.yml
conda activate head_segmentation

# (Optional but recommended) Install pre-commit hooks to preserve code format consistency
pre-commit install

🔬 Quantitative results

Keep in mind that we trained our model with CelebA dataset, which means that our model may not necessarily perform well on your data, since they may come from a different distribution than CelebA.

The table below presents results, computed on the full scale test set images, of three best models we trained. Model naming convention is as followed: <backbone>_<nn_input_image_resultion>.

🧐 Qualitative results

If you want to check predictions on some of your images, please feel free to use our Streamlit application.

cd head-segmentation

streamlit run ./scripts/apps/web_checking.py

⏰ Inference time

If you are strict with time, you can use gpu to acclerate inference. Visualization also consume some time, you can just save the final result as below.

import torch
from PIL import Image
import head_segmentation.segmentation_pipeline as seg_pipeline

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

segmentation_pipeline = seg_pipeline.HumanHeadSegmentationPipeline(device=device)

segmentation_map = segmentation_pipeline.predict(image)

segmented_region = image * cv2.cvtColor(segmentation_map, cv2.COLOR_GRAY2RGB)

pil_image = Image.fromarray(segmented_region)
pil_image.save(save_path)

The table below presents inference time which is tested on Tesla T4 (just for reference). The first image will take more time.