{"id":24378219,"url":"https://github.com/wiktorlazarski/head-segmentation","last_synced_at":"2026-03-14T17:33:19.201Z","repository":{"id":39118181,"uuid":"455985143","full_name":"wiktorlazarski/head-segmentation","owner":"wiktorlazarski","description":"👦 Human head semantic segmentation","archived":false,"fork":false,"pushed_at":"2024-05-31T09:55:31.000Z","size":7112,"stargazers_count":81,"open_issues_count":3,"forks_count":18,"subscribers_count":3,"default_branch":"main","last_synced_at":"2025-03-24T17:50:10.151Z","etag":null,"topics":["celeba-dataset","computer-vision","deep-learning","human-head","hydra","pytorch","pytorch-lightning","semantic-segmentation","semantic-segmentation-pytorch","unet","wandb"],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"other","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/wiktorlazarski.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null}},"created_at":"2022-02-05T21:03:58.000Z","updated_at":"2025-03-20T14:28:45.000Z","dependencies_parsed_at":"2023-09-23T03:18:39.897Z","dependency_job_id":null,"html_url":"https://github.com/wiktorlazarski/head-segmentation","commit_stats":null,"previous_names":[],"tags_count":4,"template":false,"template_full_name":"wiktorlazarski/ai-awesome-project-template","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/wiktorlazarski%2Fhead-segmentation","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/wiktorlazarski%2Fhead-segmentation/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/wiktorlazarski%2Fhead-segmentation/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/wiktorlazarski%2Fhead-segmentation/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/wiktorlazarski","download_url":"https://codeload.github.com/wiktorlazarski/head-segmentation/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":248283811,"owners_count":21077922,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["celeba-dataset","computer-vision","deep-learning","human-head","hydra","pytorch","pytorch-lightning","semantic-segmentation","semantic-segmentation-pytorch","unet","wandb"],"created_at":"2025-01-19T06:18:51.940Z","updated_at":"2026-03-14T17:33:14.150Z","avatar_url":"https://github.com/wiktorlazarski.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"______________________________________________________________________\n\u003cdiv align=\"center\"\u003e\n\n# 👦 Human Head Semantic Segmentation\n\n\u003cp align=\"center\"\u003e\n  \u003ca href=\"https://github.com/wiktorlazarski\"\u003e🧑‍🎓 Wiktor\u003c/a\u003e\n  \u003ca href=\"https://github.com/Szuumii\"\u003e🧑‍🎓 Kuba\u003c/a\u003e\n\u003c/p\u003e\n\n______________________________________________________________________\n\n[![ci-testing](https://github.com/wiktorlazarski/head-segmentation/actions/workflows/ci-testing.yml/badge.svg?branch=main\u0026event=push)](https://github.com/wiktorlazarski/head-segmentation/actions/workflows/ci-testing.yml)\n[![Open In Collab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1QScgxBRXWbGbQ3DmYIJ2Ja3sD_mJ_Efx?usp=sharing)\n[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)\n\n\u003c/div\u003e\n\n## 💎 Installation with `pip`\n\nInstallation is as simple as running:\n\n```bash\npip install git+https://github.com/wiktorlazarski/head-segmentation.git\n```\n\n## 🔨 How to use\n\n### 🤔 Inference\n```python\nimport head_segmentation.segmentation_pipeline as seg_pipeline\n\nsegmentation_pipeline = seg_pipeline.HumanHeadSegmentationPipeline()\n\nsegmentation_map = segmentation_pipeline.predict(image)\n```\n\n### 🎨 Visualizing\n\n```python\nimport matplotlib.pyplot as plt\n\nimport head_segmentation.visualization as vis\n\nvisualizer = vis.VisualizationModule()\n\nfigure, _ = visualizer.visualize_prediction(image, segmentation_map)\nplt.show()\n```\n\n## ⚙️ Setup for development\n\n```bash\n# Clone repo\ngit clone https://github.com/wiktorlazarski/head-segmentation.git\n\n# Go to repo directory\ncd head-segmentation\n\n# (Optional) Create virtual environment\npython -m venv venv\nsource ./venv/bin/activate\n\n# Install project in editable mode\npip install -e .[dev]\n\n# (Optional but recommended) Install pre-commit hooks to preserve code format consistency\npre-commit install\n```\n\n## 🐍 Setup for development with Anaconda or Miniconda\n\n```bash\n# Clone repo\ngit clone https://github.com/wiktorlazarski/head-segmentation.git\n\n# Go to repo directory\ncd head-segmentation\n\n# Create and activate conda environment\nconda env create -f ./conda_env.yml\nconda activate head_segmentation\n\n# (Optional but recommended) Install pre-commit hooks to preserve code format consistency\npre-commit install\n```\n\n## 🔬 Quantitative results\n\n**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.\n\nThe table below presents results, computed on the full scale test set images, of three best models we trained. Model naming convention is as followed: `\u003cbackbone\u003e_\u003cnn_input_image_resultion\u003e`.\n\n\n|      Model     | mobilenetv2_256 | mobilenetv2_512 | resnet34_512 |\n|:--------------:|:---------------:|:---------------:|:------------:|\n|    head IoU    |     0.967606    |     0.967337    | **0.968457** |\n| background IoU |     0.942936    |     0.942160    | **0.944469** |\n|      mIoU      |     0.955271    |     0.954749    | **0.956463** |\n\n\n## 🧐 Qualitative results\n\n![alt text](https://github.com/wiktorlazarski/head-segmentation/blob/main/doc/images/wiktor.png)\n![alt text](https://github.com/wiktorlazarski/head-segmentation/blob/main/doc/images/kuba.png)\n![alt text](https://github.com/wiktorlazarski/head-segmentation/blob/main/doc/images/wiktor_with_glasses.png)\n![alt text](https://github.com/wiktorlazarski/head-segmentation/blob/main/doc/images/kuba_with_helmet.png)\n\nIf you want to check predictions on some of your images, please feel free to use our Streamlit application.\n\n```bash\ncd head-segmentation\n\nstreamlit run ./scripts/apps/web_checking.py\n```\n\n## ⏰ Inference time\n\nIf 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.\n\n```python\nimport torch\nfrom PIL import Image\nimport head_segmentation.segmentation_pipeline as seg_pipeline\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nsegmentation_pipeline = seg_pipeline.HumanHeadSegmentationPipeline(device=device)\n\nsegmentation_map = segmentation_pipeline.predict(image)\n\nsegmented_region = image * cv2.cvtColor(segmentation_map, cv2.COLOR_GRAY2RGB)\n\npil_image = Image.fromarray(segmented_region)\npil_image.save(save_path)\n```\n\nThe table below presents inference time which is tested on Tesla T4 (just for reference). The first image will take more time.\n\n|                |       save figure     | just save final result|\n|:--------------:|:---------------------:|:---------------------:|\n|       cpu      |       around 2.1s     |       around 0.8s     |\n|       gpu      |       around 1.4s     |       around 0.15s    |\n\n\u003cdiv align=\"center\"\u003e\n\n### 🤗 Enjoy the model!\n\n\u003c/div\u003e\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fwiktorlazarski%2Fhead-segmentation","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fwiktorlazarski%2Fhead-segmentation","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fwiktorlazarski%2Fhead-segmentation/lists"}