Pytorch UNet
New checkpoints
Customisierte Implementierung der U-NET in Pytorch für Kaggle's Carvana Image-Maskierungsherausforderung aus High Definition-Bildern.
CUDA installieren
Installieren Sie Pytorch 1.13 oder höher
Abhängigkeiten installieren
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 --ampDieses Modell wurde von Grund auf mit 5K -Bildern trainiert und einen Würfelkoeffizienten von 0,988423 auf über 100K -Testbildern erzielt.
Es kann leicht für die Segmentierung von Multiclas, Porträtsegmentierung, medizinische Segmentierung, ...
Hinweis: Verwenden Sie Python 3.6 oder neuer
Ein Docker -Bild, das den Code und die Abhängigkeiten enthält, finden Sie in DockerHub. Sie können den Container herunterladen und springen mit (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 Standardmäßig beträgt die scale 0,5. Wenn Sie also bessere Ergebnisse erzielen möchten (aber mehr Speicher verwenden), setzen Sie sie auf 1.
Die automatische gemischte Präzision ist auch mit der Flagge --amp erhältlich. Durch gemischte Präzision kann das Modell weniger Speicher verwenden und bei der kürzlich durchgeführten GPUs durch die Verwendung von FP16 -Arithmetik schneller sein. Das Aktivieren von AMP wird empfohlen.
Nachdem Sie Ihr Modell trainiert und auf MODEL.pth sparen.
Um ein einzelnes Bild vorherzusagen und es zu speichern:
python predict.py -i image.jpg -o output.jpg
Multiple Bilder vorherzusagen und sie zu zeigen, ohne sie zu retten:
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 Sie können angeben, welche Modelldatei mit --model MODEL.pth verwendet werden soll.
Der Trainingsfortschritt kann in Echtzeit mit Gewichten und Vorurteilen visualisiert werden. Verlustkurven, Validierungskurven, Gewichte und Gradientenhistogramme sowie vorhergesagte Masken werden auf der Plattform protokolliert.
Beim Start eines Trainings wird ein Link in der Konsole gedruckt. Klicken Sie darauf, um zu Ihrem Dashboard zu gehen. Wenn Sie über ein vorhandenes W & B -Konto verfügen, können Sie es verknüpfen, indem Sie die Umgebungsvariable WANDB_API_KEY einstellen. Wenn nicht, wird ein anonymer Lauf erstellt, der nach 7 Tagen automatisch gelöscht wird.
Für den Carvana -Datensatz steht ein vorbereitendes Modell zur Verfügung. Es kann auch von Torch geladen werden.
net = torch . hub . load ( 'milesial/Pytorch-UNet' , 'unet_carvana' , pretrained = True , scale = 0.5 )Die verfügbaren Skalen betragen 0,5 und 1,0.
Die Carvana -Daten sind auf der Kaggle -Website verfügbar.
Sie können es auch mit dem Helfer -Skript herunterladen:
bash scripts/download_data.sh
Die Eingabebilder und Zielmasken sollten sich in den Ordnern data/imgs bzw. data/masks befinden (beachten Sie, dass der Ordner imgs und masks aufgrund des gierigen Datenladers keinen Unterordner oder andere Dateien enthalten sollte). Für Carvana sind Bilder RGB und Masken sind schwarz und weiß.
Sie können Ihren eigenen Datensatz verwenden, solange Sie sicherstellen, dass er in utils/data_loading.py ordnungsgemäß geladen ist.
Originalpapier von Olaf Ronneberger, Philipp Fischer, Thomas Brox:
U-NET: Faltungsnetzwerke für die biomedizinische Bildsegmentierung
