Pytorch UNet
New checkpoints
Implémentation personnalisée du U-Net dans Pytorch pour le défi de masquage de l'image Carvana de Kaggle à partir d'images haute définition.
Installer Cuda
Installer Pytorch 1.13 ou version ultérieure
Installer des dépendances
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 --ampCe modèle a été entraîné à partir de zéro avec des images 5K et a marqué un coefficient de dés de 0,988423 sur plus de 100 000 images de test.
Il peut être facilement utilisé pour la segmentation multiclasse, la segmentation des portraits, la segmentation médicale, ...
Remarque: utilisez Python 3.6 ou plus récent
Une image Docker contenant le code et les dépendances est disponible sur DockerHub. Vous pouvez télécharger et sauter dans le conteneur avec (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 Par défaut, l' scale est de 0,5, donc si vous souhaitez obtenir de meilleurs résultats (mais utilisez plus de mémoire), définissez-le sur 1.
Une précision mélangée automatique est également disponible avec le drapeau --amp . La précision mixte permet au modèle d'utiliser moins de mémoire et d'être plus rapide sur les GPU récents en utilisant l'arithmétique FP16. L'activation de l'ampli est recommandée.
Après avoir entraîné votre modèle et l'enregistrer sur MODEL.pth , vous pouvez facilement tester les masques de sortie sur vos images via la CLI.
Pour prédire une seule image et enregistrez-la:
python predict.py -i image.jpg -o output.jpg
Pour prédire plusieurs images et les montrer sans les enregistrer:
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 Vous pouvez spécifier le fichier de modèle à utiliser avec --model MODEL.pth .
Les progrès de l'entraînement peuvent être visualisés en temps réel en utilisant des poids et des biais. Les courbes de perte, les courbes de validation, les poids et les histogrammes de gradient, ainsi que les masques prévus sont enregistrés sur la plate-forme.
Lors du lancement d'une formation, un lien sera imprimé dans la console. Cliquez dessus pour accéder à votre tableau de bord. Si vous avez un compte W&B existant, vous pouvez le lier en définissant la variable d'environnement WANDB_API_KEY . Sinon, il créera une exécution anonyme qui est automatiquement supprimée après 7 jours.
Un modèle pré-entraîné est disponible pour l'ensemble de données Carvana. Il peut également être chargé à partir de torche.hub:
net = torch . hub . load ( 'milesial/Pytorch-UNet' , 'unet_carvana' , pretrained = True , scale = 0.5 )Les échelles disponibles sont de 0,5 et 1,0.
Les données de Carvana sont disponibles sur le site Web de Kaggle.
Vous pouvez également le télécharger à l'aide du script d'assistance:
bash scripts/download_data.sh
Les images d'entrée et les masques cibles doivent être respectivement dans les dossiers data/imgs et data/masks (notez que le dossier imgs et masks ne doit contenir aucun sous-dossier ou aucun autre fichier, en raison du chargeur de données gourmand). Pour Carvana, les images sont RVB et les masques sont en noir et blanc.
Vous pouvez utiliser votre propre ensemble de données tant que vous vous assurez qu'il est chargé correctement dans utils/data_loading.py .
Papier original par Olaf Ronneberger, Philipp Fischer, Thomas Brox:
U-net: réseaux convolutionnels pour la segmentation de l'image biomédicale
