Image Segmentation with VGGNet U-Net

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


Using TensorFlow backend.

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/"
)


  3%|▎         | 12/367 [00:00<00:02, 119.63it/s]
Verifying train dataset
100%|██████████| 367/367 [00:01<00:00, 203.14it/s]
 21%|██        | 21/101 [00:00<00:00, 202.12it/s]
Dataset verified! 
Verifying val dataset
100%|██████████| 101/101 [00:00<00:00, 199.39it/s]
Dataset verified! 
Starting Epoch  0
Epoch 1/1
512/512 [==============================] - 139s - loss: 1.2426 - acc: 0.6802 - val_loss: 1.2152 - val_acc: 0.6010
saved  model_output/image-seg.model.0
Finished Epoch 0
Starting Epoch  1
Epoch 1/1
512/512 [==============================] - 128s - loss: 0.7111 - acc: 0.7843 - val_loss: 0.8228 - val_acc: 0.7302
saved  model_output/image-seg.model.1
Finished Epoch 1
Starting Epoch  2
Epoch 1/1
512/512 [==============================] - 128s - loss: 0.5134 - acc: 0.8424 - val_loss: 0.4538 - val_acc: 0.8648
saved  model_output/image-seg.model.2
Finished Epoch 2
Starting Epoch  3
Epoch 1/1
512/512 [==============================] - 128s - loss: 0.4178 - acc: 0.8713 - val_loss: 0.4496 - val_acc: 0.8576
saved  model_output/image-seg.model.3
Finished Epoch 3
Starting Epoch  4
Epoch 1/1
512/512 [==============================] - 127s - loss: 0.3410 - acc: 0.8936 - val_loss: 0.3900 - val_acc: 0.8749
saved  model_output/image-seg.model.4
Finished Epoch 4
Starting Epoch  5
Epoch 1/1
512/512 [==============================] - 128s - loss: 0.2911 - acc: 0.9083 - val_loss: 0.3558 - val_acc: 0.8881
saved  model_output/image-seg.model.5
Finished Epoch 5
Starting Epoch  6
Epoch 1/1
512/512 [==============================] - 128s - loss: 0.2539 - acc: 0.9193 - val_loss: 0.3884 - val_acc: 0.8793
saved  model_output/image-seg.model.6
Finished Epoch 6
Starting Epoch  7
Epoch 1/1
512/512 [==============================] - 128s - loss: 0.2179 - acc: 0.9304 - val_loss: 0.3390 - val_acc: 0.8997
saved  model_output/image-seg.model.7
Finished Epoch 7
Starting Epoch  8
Epoch 1/1
512/512 [==============================] - 128s - loss: 0.1837 - acc: 0.9406 - val_loss: 0.2975 - val_acc: 0.9111
saved  model_output/image-seg.model.8
Finished Epoch 8
Starting Epoch  9
Epoch 1/1
512/512 [==============================] - 128s - loss: 0.1614 - acc: 0.9472 - val_loss: 0.4474 - val_acc: 0.8777
saved  model_output/image-seg.model.9
Finished Epoch 9

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 [ ]: