milesial/pytorch-unet

This repository is a Python implementation of a neural network called U-Net, used for image segmentation. Image segmentation means automatically drawing precise outlines around objects in a photo, for example identifying exactly which pixels belong to a car versus the background. The original U-Net design came from medical imaging research, but it works on many kinds of photos. The code was built and tested for a Kaggle competition called the Carvana Image Masking Challenge, where the task was to draw accurate outlines around cars in high-resolution photographs. The model was trained on roughly 5,000 images and scored very well on over 100,000 test images. Beyond cars, the same approach applies to portrait photos, medical scans, and other segmentation tasks. You can run this code either by installing the required Python libraries directly on your machine, or through a pre-packaged Docker container that includes everything already set up. Training is done from the command line with options to control how many passes through the data to run, how large each batch is, the learning rate, and whether to use mixed precision (a technique that speeds up training on modern graphics cards while using less memory). Once a model is trained and saved, a separate prediction script takes one or more image files and outputs the segmentation masks. A pretrained model trained on the Carvana dataset is also available for direct download or loading with one line of code, so you can get results without training from scratch. Training progress, loss curves, and predicted mask previews can be monitored in real time through Weights and Biases, an external experiment-tracking service.