Based on code in Divam Gupta's image-segmentation-keras repository.
In [ ]:
    
import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    
In [ ]:
    
! pip install opencv-contrib-python
    
In [ ]:
    
! pip install tqdm
    
In [ ]:
    
! pip install keras_segmentation
    
In [2]:
    
import keras_segmentation
    
    
In [3]:
    
model = keras_segmentation.models.unet.vgg_unet(n_classes=51, 
                                                      input_height=416, 
                                                      input_width=608)
    
In [4]:
    
model.train(
    train_images = "image-seg-data/images_prepped_train/",
    train_annotations = "image-seg-data/annotations_prepped_train/",
    checkpoints_path = "model_output/image-seg/",
    epochs = 10,
    validate=True,
    val_images = "image-seg-data/images_prepped_test/",
    val_annotations = "image-seg-data/annotations_prepped_test/"
)
    
    
    
    
    
    
    
In [9]:
    
out = model.predict_segmentation(
    inp="image-seg-data/images_prepped_test/0016E5_07965.png",
    out_fname="output-10epochs.png"
)
    
In [16]:
    
model.load_weights('model_output/image-seg.8')
    
In [17]:
    
out = model.predict_segmentation(
    inp="image-seg-data/images_prepped_test/0016E5_07965.png",
    out_fname="output-9epochs.png"
)
    
In [ ]: