Image Segmentation via Keras

Implementation of Segnet, FCN, UNet, PSPNet and other models in Keras. This repo is cloned and modify based on https://github.com/divamgupta/image-segmentation-keras. Click here to see original README

Important Dependencies

• python 3.5+

• Keras 2.0

• python_opencv

• Theano / Tensorflow / CNTK

• numpy, matplotlib

Use following to install necessary dependencies

python setup.py install

Prepare data for segmentation

I use CUB-200-2011 and VOC-2012 for image segmentation. You should first download these dataset.

After downloading the dataset, you should put corresponding dataset into Datasets, it should be like

-- Datasets
---- CUB_200_2011
------ images
------ segmentations
---- VOC
------ JEPGImages
------ SegmentationClass
------ SegmentationObject

To prepare data for segmentation task, use DataPreprocess.py in utils, for example

convert_segmentations_CUB()
make_dataset_CUB()

convert_segmentations_VOC()
make_dataset_VOC

After that, there are train and test directory in above mentioned datasets

Run the code

To run the code, you should see the example in run.py.

import keras_segmentation
from keras.models import load_model
import tensorflow as tf
import os

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" 
os.environ["CUDA_VISIBLE_DEVICES"]="2, 3" 
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8)

DATA_NAME = "VOC"
EPOCH = [5, 10]
CHECKOUTPOINT_PATH = "./tmp/voc_5_10"

if DATA_NAME == "VOC":
    train_images_path = "./Datasets/VOC/train/imgs/"
    train_segs_path = "./Datasets/VOC/train/segs/"
    test_images_path = "./Datasets/VOC/test/imgs/"
    test_segs_path = "./Datasets/VOC/test/segs"
    class_num = 21
if DATA_NAME == "CUB":
    train_images_path = "./Datasets/CUB_200_2011/train/imgs/"
    train_segs_path = "./Datasets/CUB_200_2011/train/segs/"
    test_images_path = "./Datasets/CUB_200_2011/test/imgs/"
    test_segs_path = "./Datasets/CUB_200_2011/test/segs"
    class_num = 2
'''
Change model name
'''
model = keras_segmentation.models.unet.resnet50_unet(n_classes=class_num,  input_height=416, input_width=608)
# model.load_weights("./tmp/cub_psspnet_vgg_pspnet.9")

for i in range(EPOCH[0]):

    '''
    Train
    '''
    model.train( 
        train_images =  train_images_path,
        train_annotations = train_segs_path,
        checkpoints_path = CHECKOUTPOINT_PATH , epochs=EPOCH[1], verify_dataset = False
    )

    '''
    Output
    '''
    # for (i, image_dir) in enumerate(os.listdir(test_images_path)):
    #     if i%50 == 0:
    #         out = model.predict_segmentation(
    #             inp= os.path.join(test_images_path, image_dir),
    #             out_fname= os.path.join("./Output/", image_dir)
    #         )
    #         # import matplotlib.pyplot as plt
    #         # plt.imshow(out)

    '''
    Test mIoU
    '''
    test_image_list = [os.path.join(test_images_path, i) for i in os.listdir(test_images_path)]
    test_segs_list = [os.path.join(test_segs_path, i) for i in os.listdir(test_segs_path)]

    model.evaluate_segmentation(
        inp_images = test_image_list,
        annotations = test_segs_list
    )

To get output images, you should see example in test.py:

import keras_segmentation
import numpy as np
import cv2
import os

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" 
os.environ["CUDA_VISIBLE_DEVICES"]="2,3" 

DATA_NAME = "VOC"

if DATA_NAME == "VOC":
    train_images_path = "./Datasets/VOC/train/imgs/"
    train_segs_path = "./Datasets/VOC/train/segs/"
    test_images_path = "./Datasets/VOC/test/imgs/"
    test_segs_path = "./Datasets/VOC/test/segs"
    class_num = 21
if DATA_NAME == "CUB":
    train_images_path = "./Datasets/CUB_200_2011/train/imgs/"
    train_segs_path = "./Datasets/CUB_200_2011/train/segs/"
    test_images_path = "./Datasets/CUB_200_2011/test/imgs/"
    test_segs_path = "./Datasets/CUB_200_2011/test/segs"
    class_num = 2
'''
Change model name
'''
model = keras_segmentation.models.unet.resnet50_unet(n_classes=class_num,  input_height=416, input_width=608)
model.load_weights("./tmp/voc_5_10.9")

# load any of the 3 pretrained models

for (i, image_dir) in enumerate(os.listdir(test_images_path)):
    if image_dir == "2011_002114.jpg":
        out = model.predict_segmentation(
            inp= os.path.join(test_images_path, image_dir),
            out_fname= os.path.join("./Output_VOC/", image_dir))

Change model names, checkpoint names to run the corresponding models.

Contact

If you have any problem, please contact me through: lidongyue AT sjtu DOT edu DOT cn