Find a file
milesial 8b614c3e31 Modified to take any image size (with even width, height > width/2)
Former-commit-id: 2751e6a3df45c1527376a4697d3804d683095d83
2017-11-30 07:19:52 +01:00
unet Added CLI for predict, cleaned up code, updated README 2017-11-30 06:45:19 +01:00
.gitignore Added CLI for predict, cleaned up code, updated README 2017-11-30 06:45:19 +01:00
crf.py Added simple predict + submit script 2017-08-21 18:00:07 +02:00
data_vis.py Created a basic train loop + changed a bit loss and utils 2017-08-17 21:16:19 +02:00
eval.py Added simple predict + submit script 2017-08-21 18:00:07 +02:00
load.py Added simple predict + submit script 2017-08-21 18:00:07 +02:00
main.py Added CLI for predict, cleaned up code, updated README 2017-11-30 06:45:19 +01:00
MODEL.pth.REMOVED.git-id Added simple trained model 2017-08-23 17:38:55 +02:00
myloss.py Created a basic train loop + changed a bit loss and utils 2017-08-17 21:16:19 +02:00
predict.py Modified to take any image size (with even width, height > width/2) 2017-11-30 07:19:52 +01:00
README.md Modified to take any image size (with even width, height > width/2) 2017-11-30 07:19:52 +01:00
submit.py Modified to take any image size (with even width, height > width/2) 2017-11-30 07:19:52 +01:00
train.py Added CLI for predict, cleaned up code, updated README 2017-11-30 06:45:19 +01:00
utils.py Added CLI for predict, cleaned up code, updated README 2017-11-30 06:45:19 +01:00

Pytorch-UNet

Customized implementation of the U-Net in Pytorch for Kaggle's Carvana Image Masking Challenge.

This model scored a dice coefficient of 0.988423 (511 out of 735), which is bad but could be improved with more training, data augmentation, fine tuning, and playing with CRF post-processing.

The model used for the last submission is stored in the MODEL.pth file, if you wish to play with it. The data is available on the Kaggle website.

Usage

### Prediction

You can easily test the output masks on your images via the CLI. To see all options: python predict.py -h

To predict a single image and save it:

python predict.py -i image.jpg -o ouput.jpg

To predict a multiple images and show them without saving them:

python predict.py -i image1.jpg image2.jpg --viz --no-save

You can use the cpu-only version with --cpu.

You can specify which model file to use with --model MODEL.pth.

Warning

In order to process the image, it is splitted into two squares (a left on and a right one), and each square is passed into the net. The two square masks are then merged again to produce the final image. As a consequence, the height of the image must be strictly superior than half the width. Make sure the width is even too.