Pytorch UNet
New checkpoints
고화질 이미지에서 Kaggle의 Carvana 이미지 마스킹 챌린지를 위해 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 데이터 세트에는 사전 제한 모델을 사용할 수 있습니다. 토치에서로드 할 수도 있습니다. 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 폴더에는 탐욕스러운 데이터 로더로 인해 하위 폴더 또는 기타 파일이 포함되어 있지 않아야합니다). 카바나의 경우 이미지는 RGB이고 마스크는 흑백입니다.
utils/data_loading.py 에 올바르게로드되어 있는지 확인하는 한 자신의 데이터 세트를 사용할 수 있습니다.
Olaf Ronneberger, Philipp Fischer, Thomas Brox의 오리지널 논문 :
U-Net : 생의학 이미지 세분화를위한 컨볼 루션 네트워크
