Installation is as simple as running:
pip install git+https://github.com/wiktorlazarski/head-segmentation.gitimport head_segmentation.segmentation_pipeline as seg_pipeline
segmentation_pipeline = seg_pipeline.HumanHeadSegmentationPipeline()
segmentation_map = segmentation_pipeline.predict(image)import matplotlib.pyplot as plt
import head_segmentation.visualization as vis
visualizer = vis.VisualizationModule()
figure, _ = visualizer.visualize_prediction(image, segmentation_map)
plt.show()# 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# 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 installKeep 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 |
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.pyIf 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 |