Pytorch UNet
New checkpoints
高解像度画像からのKaggleのCarvana Image Masking ChallengeのPytorchでのU-Netのカスタマイズされた実装。
cudaをインストールします
Pytorch 1.13以降をインストールします
依存関係をインストールします
pip install -r requirements.txtbash scripts/download_data.sh
python train.py --ampcurl https://get.docker.com | sh && sudo systemctl --now enable dockerdistribution= $( . /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 dockersudo docker run --rm --shm-size=8g --ulimit memlock=-1 --gpus all -it milesial/unetbash scripts/download_data.sh
python train.py --ampこのモデルは、5K画像でゼロからトレーニングされ、100Kを超えるテスト画像で0.988423のサイコロ係数を獲得しました。
マルチクラスセグメンテーション、ポートレートセグメンテーション、医療セグメンテーション、...
注：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算術を使用して最近のGPUでより速くすることができます。 AMPを有効にすることをお勧めします。
モデルをトレーニングして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 Datasetでは、事前に処理されたモデルを利用できます。 Torch.hubからロードすることもできます。
net = torch . hub . load ( 'milesial/Pytorch-UNet' , 'unet_carvana' , pretrained = True , scale = 0.5 )利用可能なスケールは0.5および1.0です。
Carvanaデータは、Kaggle Webサイトで入手できます。
ヘルパースクリプトを使用してダウンロードすることもできます。
bash scripts/download_data.sh
入力画像とターゲットマスクは、それぞれdata/imgsとdata/masksフォルダーにある必要があります（ imgsおよびmasksフォルダーには、貪欲なデータローダーのため、サブフォルダーまたはその他のファイルを含めるべきではないことに注意してください）。カルバナの場合、画像はRGBで、マスクは白黒です。
utils/data_loading.pyで適切にロードされていることを確認する限り、独自のデータセットを使用できます。
トーマス・ブロックス、フィリップ・フィッシャー、オラフ・ロンネバーガーによるオリジナルペーパー：
u-net：生物医学画像セグメンテーションのための畳み込みネットワーク
