Pytorch UNet

U-Net: семантическая сегментация с Pytorch

Индивидуальная реализация U-Net в Pytorch для Kaggle's Carvana Masking Challenge из изображений высокой четкости.

• Быстрый старт
  • Без докера
  • С докером
• Описание
• Использование
  • Докер
  • Обучение
  • Прогноз
• Вес и предвзятость
• Предварительная модель
• Данные

Быстрый старт

Без докера

1. Установите CUDA

2. Установите Pytorch 1.13 или позже

3. Установить зависимости

pip install -r requirements.txt

4. Загрузите данные и запустите обучение:

bash scripts/download_data.sh
python train.py --amp

С докером

1. Установите Docker 19.03 или позже:

curl https://get.docker.com | sh && sudo systemctl --now enable docker

2. Установите Toolkit nvidia Container:

distribution= $( . /etc/os-release ; echo $ID$VERSION_ID ) 
   && curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - 
   && curl -s -L https://nvidia.github.io/nvidia-docker/ $distribution /nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update
sudo apt-get install -y nvidia-docker2
sudo systemctl restart docker

3. Загрузите и запустите изображение:

sudo docker run --rm --shm-size=8g --ulimit memlock=-1 --gpus all -it milesial/unet

4. Загрузите данные и запустите обучение:

bash scripts/download_data.sh
python train.py --amp

Описание

Эта модель была обучена с нуля с изображениями 5K и набрала коэффициент костей 0,988423 на более 100 тысячах тестовых изображений.

Его можно легко использовать для сегментации мультикласс, портретной сегментации, медицинской сегментации, ...

Использование

Примечание: используйте Python 3.6 или более новый

Докер

Изображение Docker, содержащее код и зависимости, доступно на Dockerhub. Вы можете скачать и прыгнуть в контейнер с (Docker> = 19.03):

 docker run -it --rm --shm-size=8g --ulimit memlock=-1 --gpus all milesial/unet

Обучение

> python train.py -h
usage: train.py [-h] [--epochs E] [--batch-size B] [--learning-rate LR]
                [--load LOAD] [--scale SCALE] [--validation VAL] [--amp]

Train the UNet on images and target masks

optional arguments:
  -h, --help            show this help message and exit
  --epochs E, -e E      Number of epochs
  --batch-size B, -b B  Batch size
  --learning-rate LR, -l LR
                        Learning rate
  --load LOAD, -f LOAD  Load model from a .pth file
  --scale SCALE, -s SCALE
                        Downscaling factor of the images
  --validation VAL, -v VAL
                        Percent of the data that is used as validation (0-100)
  --amp                 Use mixed precision

По умолчанию scale составляет 0,5, поэтому, если вы хотите получить лучшие результаты (но используйте больше памяти), установите ее на 1.

Автоматическая смешанная точность также доступна с флагом --amp . Смешанная точность позволяет модели использовать меньше памяти и быть быстрее на недавних графических процессорах, используя арифметику FP16. Рекомендуется включение усилителя.

Прогноз

После обучения своей модели и сохранения ее на MODEL.pth .

Чтобы предсказать одно изображение и сохранить его:

python predict.py -i image.jpg -o output.jpg

Чтобы предсказать несколько изображений и показать их без сохранения:

python predict.py -i image1.jpg image2.jpg --viz --no-save

> python predict.py -h
usage: predict.py [-h] [--model FILE] --input INPUT [INPUT ...] 
                  [--output INPUT [INPUT ...]] [--viz] [--no-save]
                  [--mask-threshold MASK_THRESHOLD] [--scale SCALE]

Predict masks from input images

optional arguments:
  -h, --help            show this help message and exit
  --model FILE, -m FILE
                        Specify the file in which the model is stored
  --input INPUT [INPUT ...], -i INPUT [INPUT ...]
                        Filenames of input images
  --output INPUT [INPUT ...], -o INPUT [INPUT ...]
                        Filenames of output images
  --viz, -v             Visualize the images as they are processed
  --no-save, -n         Do not save the output masks
  --mask-threshold MASK_THRESHOLD, -t MASK_THRESHOLD
                        Minimum probability value to consider a mask pixel white
  --scale SCALE, -s SCALE
                        Scale factor for the input images

Вы можете указать, какой файл модели использовать с помощью --model MODEL.pth .

Вес и предвзятость

Прогресс обучения может быть визуализирован в режиме реального времени с использованием весов и уклонов. Кривые потерь, кривые валидации, веса и градиентные гистограммы, а также предсказанные маски регистрируются на платформе.

При запуске обучения в консоли будет напечатана ссылка. Нажмите на него, чтобы перейти на приборную панель. Если у вас есть существующая учетная запись W & B, вы можете связать ее, установив переменную среды WANDB_API_KEY . Если нет, это создаст анонимный прогон, который автоматически удаляется через 7 дней.

Предварительная модель

Предварительная модель доступна для набора данных Carvana. Это также может быть загружен из Torch.hub:

 net = torch . hub . load ( 'milesial/Pytorch-UNet' , 'unet_carvana' , pretrained = True , scale = 0.5 )

Доступные масштабы 0,5 и 1,0.

Данные

Данные Carvana доступны на веб -сайте Kaggle.

Вы также можете скачать его с помощью сценария Helper:

 bash scripts/download_data.sh

Входные изображения и целевые маски должны быть в папках data/imgs и data/masks соответственно (обратите внимание, что папка imgs и masks не должна содержать какого-либо подразделения или любых других файлов из-за жадного загрузчика данных). Для Карваны изображения - RGB, а маски черно -белые.

Вы можете использовать свой собственный набор данных, если убедиться, что он правильно загружен в utils/data_loading.py .

---

Оригинальная бумага Олафа Роннебергера, Филиппа Фишера, Томаса Брокс:

U-Net: сверточные сети для сегментации биомедицинской изображения