{"id":14958151,"url":"https://github.com/milesial/pytorch-unet","last_synced_at":"2025-05-13T21:09:57.031Z","repository":{"id":37470736,"uuid":"100485912","full_name":"milesial/Pytorch-UNet","owner":"milesial","description":"PyTorch implementation of the U-Net for image semantic segmentation with high quality images","archived":false,"fork":false,"pushed_at":"2024-08-11T03:00:15.000Z","size":48537,"stargazers_count":10069,"open_issues_count":117,"forks_count":2605,"subscribers_count":70,"default_branch":"master","last_synced_at":"2025-05-08T10:05:03.647Z","etag":null,"topics":["convolutional-networks","convolutional-neural-networks","deep-learning","kaggle","pytorch","pytorch-unet","semantic-segmentation","tensorboard","unet","wandb","weights-and-biases"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"gpl-3.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/milesial.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,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2017-08-16T12:17:08.000Z","updated_at":"2025-05-08T09:16:29.000Z","dependencies_parsed_at":"2024-09-22T06:02:06.793Z","dependency_job_id":null,"html_url":"https://github.com/milesial/Pytorch-UNet","commit_stats":{"total_commits":132,"total_committers":25,"mean_commits":5.28,"dds":"0.24242424242424243","last_synced_commit":"21d7850f2af30a9695bbeea75f3136aa538cfc4a"},"previous_names":[],"tags_count":4,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/milesial%2FPytorch-UNet","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/milesial%2FPytorch-UNet/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/milesial%2FPytorch-UNet/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/milesial%2FPytorch-UNet/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/milesial","download_url":"https://codeload.github.com/milesial/Pytorch-UNet/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":254029002,"owners_count":22002283,"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":["convolutional-networks","convolutional-neural-networks","deep-learning","kaggle","pytorch","pytorch-unet","semantic-segmentation","tensorboard","unet","wandb","weights-and-biases"],"created_at":"2024-09-24T13:16:21.864Z","updated_at":"2025-05-13T21:09:51.955Z","avatar_url":"https://github.com/milesial.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# U-Net: Semantic segmentation with PyTorch\n\u003ca href=\"#\"\u003e\u003cimg src=\"https://img.shields.io/github/actions/workflow/status/milesial/PyTorch-UNet/main.yml?logo=github\u0026style=for-the-badge\" /\u003e\u003c/a\u003e\n\u003ca href=\"https://hub.docker.com/r/milesial/unet\"\u003e\u003cimg src=\"https://img.shields.io/badge/docker%20image-available-blue?logo=Docker\u0026style=for-the-badge\" /\u003e\u003c/a\u003e\n\u003ca href=\"https://pytorch.org/\"\u003e\u003cimg src=\"https://img.shields.io/badge/PyTorch-v1.13+-red.svg?logo=PyTorch\u0026style=for-the-badge\" /\u003e\u003c/a\u003e\n\u003ca href=\"#\"\u003e\u003cimg src=\"https://img.shields.io/badge/python-v3.6+-blue.svg?logo=python\u0026style=for-the-badge\" /\u003e\u003c/a\u003e\n\n![input and output for a random image in the test dataset](https://i.imgur.com/GD8FcB7.png)\n\n\nCustomized implementation of the [U-Net](https://arxiv.org/abs/1505.04597) in PyTorch for Kaggle's [Carvana Image Masking Challenge](https://www.kaggle.com/c/carvana-image-masking-challenge) from high definition images.\n\n- [Quick start](#quick-start)\n  - [Without Docker](#without-docker)\n  - [With Docker](#with-docker)\n- [Description](#description)\n- [Usage](#usage)\n  - [Docker](#docker)\n  - [Training](#training)\n  - [Prediction](#prediction)\n- [Weights \u0026 Biases](#weights--biases)\n- [Pretrained model](#pretrained-model)\n- [Data](#data)\n\n## Quick start\n\n### Without Docker\n\n1. [Install CUDA](https://developer.nvidia.com/cuda-downloads)\n\n2. [Install PyTorch 1.13 or later](https://pytorch.org/get-started/locally/)\n\n3. Install dependencies\n```bash\npip install -r requirements.txt\n```\n\n4. Download the data and run training:\n```bash\nbash scripts/download_data.sh\npython train.py --amp\n```\n\n### With Docker\n\n1. [Install Docker 19.03 or later:](https://docs.docker.com/get-docker/)\n```bash\ncurl https://get.docker.com | sh \u0026\u0026 sudo systemctl --now enable docker\n```\n2. [Install the NVIDIA container toolkit:](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html)\n```bash\ndistribution=$(. /etc/os-release;echo $ID$VERSION_ID) \\\n   \u0026\u0026 curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - \\\n   \u0026\u0026 curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list\nsudo apt-get update\nsudo apt-get install -y nvidia-docker2\nsudo systemctl restart docker\n```\n3. [Download and run the image:](https://hub.docker.com/repository/docker/milesial/unet)\n```bash\nsudo docker run --rm --shm-size=8g --ulimit memlock=-1 --gpus all -it milesial/unet\n```\n\n4. Download the data and run training:\n```bash\nbash scripts/download_data.sh\npython train.py --amp\n```\n\n## Description\nThis model was trained from scratch with 5k images and scored a [Dice coefficient](https://en.wikipedia.org/wiki/S%C3%B8rensen%E2%80%93Dice_coefficient) of 0.988423 on over 100k test images.\n\nIt can be easily used for multiclass segmentation, portrait segmentation, medical segmentation, ...\n\n\n## Usage\n**Note : Use Python 3.6 or newer**\n\n### Docker\n\nA docker image containing the code and the dependencies is available on [DockerHub](https://hub.docker.com/repository/docker/milesial/unet).\nYou can download and jump in the container with ([docker \u003e=19.03](https://docs.docker.com/get-docker/)):\n\n```console\ndocker run -it --rm --shm-size=8g --ulimit memlock=-1 --gpus all milesial/unet\n```\n\n\n### Training\n\n```console\n\u003e python train.py -h\nusage: train.py [-h] [--epochs E] [--batch-size B] [--learning-rate LR]\n                [--load LOAD] [--scale SCALE] [--validation VAL] [--amp]\n\nTrain the UNet on images and target masks\n\noptional arguments:\n  -h, --help            show this help message and exit\n  --epochs E, -e E      Number of epochs\n  --batch-size B, -b B  Batch size\n  --learning-rate LR, -l LR\n                        Learning rate\n  --load LOAD, -f LOAD  Load model from a .pth file\n  --scale SCALE, -s SCALE\n                        Downscaling factor of the images\n  --validation VAL, -v VAL\n                        Percent of the data that is used as validation (0-100)\n  --amp                 Use mixed precision\n```\n\nBy default, the `scale` is 0.5, so if you wish to obtain better results (but use more memory), set it to 1.\n\nAutomatic mixed precision is also available with the `--amp` flag. [Mixed precision](https://arxiv.org/abs/1710.03740) allows the model to use less memory and to be faster on recent GPUs by using FP16 arithmetic. Enabling AMP is recommended.\n\n\n### Prediction\n\nAfter training your model and saving it to `MODEL.pth`, you can easily test the output masks on your images via the CLI.\n\nTo predict a single image and save it:\n\n`python predict.py -i image.jpg -o output.jpg`\n\nTo predict a multiple images and show them without saving them:\n\n`python predict.py -i image1.jpg image2.jpg --viz --no-save`\n\n```console\n\u003e python predict.py -h\nusage: predict.py [-h] [--model FILE] --input INPUT [INPUT ...] \n                  [--output INPUT [INPUT ...]] [--viz] [--no-save]\n                  [--mask-threshold MASK_THRESHOLD] [--scale SCALE]\n\nPredict masks from input images\n\noptional arguments:\n  -h, --help            show this help message and exit\n  --model FILE, -m FILE\n                        Specify the file in which the model is stored\n  --input INPUT [INPUT ...], -i INPUT [INPUT ...]\n                        Filenames of input images\n  --output INPUT [INPUT ...], -o INPUT [INPUT ...]\n                        Filenames of output images\n  --viz, -v             Visualize the images as they are processed\n  --no-save, -n         Do not save the output masks\n  --mask-threshold MASK_THRESHOLD, -t MASK_THRESHOLD\n                        Minimum probability value to consider a mask pixel white\n  --scale SCALE, -s SCALE\n                        Scale factor for the input images\n```\nYou can specify which model file to use with `--model MODEL.pth`.\n\n## Weights \u0026 Biases\n\nThe training progress can be visualized in real-time using [Weights \u0026 Biases](https://wandb.ai/).  Loss curves, validation curves, weights and gradient histograms, as well as predicted masks are logged to the platform.\n\nWhen launching a training, a link will be printed in the console. Click on it to go to your dashboard. If you have an existing W\u0026B account, you can link it\n by setting the `WANDB_API_KEY` environment variable. If not, it will create an anonymous run which is automatically deleted after 7 days.\n\n\n## Pretrained model\nA [pretrained model](https://github.com/milesial/Pytorch-UNet/releases/tag/v3.0) is available for the Carvana dataset. It can also be loaded from torch.hub:\n\n```python\nnet = torch.hub.load('milesial/Pytorch-UNet', 'unet_carvana', pretrained=True, scale=0.5)\n```\nAvailable scales are 0.5 and 1.0.\n\n## Data\nThe Carvana data is available on the [Kaggle website](https://www.kaggle.com/c/carvana-image-masking-challenge/data).\n\nYou can also download it using the helper script:\n\n```\nbash scripts/download_data.sh\n```\n\nThe input images and target masks should be in the `data/imgs` and `data/masks` folders respectively (note that the `imgs` and `masks` folder should not contain any sub-folder or any other files, due to the greedy data-loader). For Carvana, images are RGB and masks are black and white.\n\nYou can use your own dataset as long as you make sure it is loaded properly in `utils/data_loading.py`.\n\n\n---\n\nOriginal paper by Olaf Ronneberger, Philipp Fischer, Thomas Brox:\n\n[U-Net: Convolutional Networks for Biomedical Image Segmentation](https://arxiv.org/abs/1505.04597)\n\n![network architecture](https://i.imgur.com/jeDVpqF.png)\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmilesial%2Fpytorch-unet","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fmilesial%2Fpytorch-unet","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmilesial%2Fpytorch-unet/lists"}