Pytorch UNet
New checkpoints
التنفيذ المخصص لـ U-NET في Pytorch for Kaggle's Carvana Image Lask Challenge من صور عالية الدقة.
تثبيت 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 وسجل معامل النرد من 0.988423 على أكثر من 100 ألف صورة اختبار.
يمكن استخدامه بسهولة للتجزئة متعددة الطوائف ، تجزئة الصور ، التجزئة الطبية ، ...
ملاحظة: استخدم 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. يوصى بتمكين 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. يمكن أيضًا تحميله من 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 يجب ألا يحتوي على أي مجلد فرعي أو أي ملفات أخرى ، بسبب محمل البيانات الجشع). بالنسبة لكارفانا ، الصور RGB والأقنعة أسود وأبيض.
يمكنك استخدام مجموعة البيانات الخاصة بك طالما أنك تتأكد من تحميلها بشكل صحيح في utils/data_loading.py .
الورقة الأصلية من تأليف أولاف رونبرغر ، فيليب فيشر ، توماس بروكس:
U-Net: شبكات تلافيفية لتجزئة الصور الطبية الحيوية
