• Stars
    star
    144
  • Rank 255,590 (Top 6 %)
  • Language
    Python
  • License
    MIT License
  • Created almost 7 years ago
  • Updated over 4 years ago

Reviews

There are no reviews yet. Be the first to send feedback to the community and the maintainers!

Repository Details

A simpler way of reading and augmenting image segmentation data into TensorFlow

SemSeg Pipeline

This is an efficient semantic segmentaiton data input pipeline function for Tensorflow 2.

The pipeline is made to be multithreaded, and uses the tf data API to prerocess the image and segmentation masks with augmentations on the CPU.

Augmentation Examples:

Following shows the same image, loaded with the pipeline, note the different augmentations (birghtness, contrast, saturation, cropping, flips) where the masks are changed accordingly. The example image is taken from the PASCAL VOC dataset.

screenshot 2018-10-16 at 22 56 50

Example use:

from dataloader import DataLoader
import tensorflow as tf
import os

IMAGE_DIR_PATH = 'data/training/images'
MASK_DIR_PATH = 'data/training/masks'
BATCH_SIZE = 4

# create list of PATHS
image_paths = [os.path.join(IMAGE_DIR_PATH, x) for x in os.listdir(IMAGE_DIR_PATH) if x.endswith('.png')]
mask_paths = [os.path.join(MASK_DIR_PATH, x) for x in os.listdir(MASK_DIR_PATH) if x.endswith('.png')]

# Where image_paths[0] = 'data/training/images/image_0.png' 
# And mask_paths[0] = 'data/training/masks/mask_0.png'

# Initialize the dataloader object
dataset = DataLoader(image_paths=image_paths,
                     mask_paths=mask_paths,
                     image_size=(256, 256),
                     crop_percent=0.8,
                     channels=(3, 1),
                     augment=True,
                     compose=False,
                     seed=47)

# Parse the images and masks, and return the data in batches, augmented optionally.
dataset = dataset.data_batch(batch_size=BATCH_SIZE,
                             shuffle=True)

# Initialize the data queue
for image, mask in dataset:      
  # Do whatever you want now

Note

You can optionally also specify a color palette for your masks. If your masks are encoded as RGB images and you want to one hot encode them.