Pytorch UNet

U-NET：pytorch的语义分割

从高清图像中的Kaggle的Carvana Image掩盖挑战中，在Pytorch中定制了U-NET实现。

• 快速开始
  • 没有Docker
  • 与Docker
• 描述
• 用法
  • Docker
  • 训练
  • 预言
• 重量和偏见
• 预验证的模型
• 数据

快速开始

没有Docker

1. 安装CUDA

2. 安装Pytorch 1.13或更高版本

3. 安装依赖项

pip install -r requirements.txt

4. 下载数据并运行培训：

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

与Docker

1. 安装Docker 19.03或更高版本：

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

2. 安装NVIDIA容器工具包：

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图像，并在超过100K的测试图像上得分为0.988423。

它可以轻松用于多类分割，肖像细分，医学细分，...

用法

注意：使用Python 3.6或更新

Docker

包含代码和依赖项的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算术使用更少的内存，并在最近的GPU上更快。建议启用放大器。

预言

训练模型并将其保存到MODEL.pth之后，您可以通过CLI轻松测试图像上的输出掩码。

预测单个图像并保存它：

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网站上找到。

您也可以使用助手脚本下载它：

 bash scripts/download_data.sh

输入图像和目标掩码应分别在data/imgs和data/masks文件夹中（请注意，由于贪婪的数据加载程序， imgs和masks文件夹不应包含任何子文件夹或任何其他文件）。对于carvana，图像是RGB，面具是黑色和白色。

只要确保在utils/data_loading.py中正确加载它，就可以使用自己的数据集。

---

Olaf Ronneberger，Philipp Fischer，Thomas Brox的原始论文：

U-NET：生物医学图像分割的卷积网络