Skip to content

wiktorlazarski/head-segmentation

Repository files navigation


πŸ‘¦ Human Head Semantic Segmentation

πŸ§‘β€πŸŽ“ Wiktor πŸ§‘β€πŸŽ“ Kuba


ci-testing Open In Collab Code style: black

πŸ’Ž 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>.

Model mobilenetv2_256 mobilenetv2_512 resnet34_512
head IoU 0.967606 0.967337 0.968457
background IoU 0.942936 0.942160 0.944469
mIoU 0.955271 0.954749 0.956463

🧐 Qualitative results

alt text alt text alt text alt text

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.

save figure just save final result
cpu around 2.1s around 0.8s
gpu around 1.4s around 0.15s

πŸ€— Enjoy the model!