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