Synthetic image composting as data-augmentation for semantic segmentation datasets

synthetic.py

from PIL import Image
import os, sys
from scipy.misc import imsave
import cv2
import numpy as np
import random
import argparse

'''
Creates synthetic images with custom backgrounds, using the default images and masks in the dataset.
It can be used to augment semantic segmentaion dataset, used for training deep neural networks.
Args: source image directory, source mask directory, custom background directory, 
destination image directory and destination mask directory.

Sample run: python synthetic_arg.py -iimg PNGImages_128/ -imsk PNGMasks_128/ -ibgd bgd_img/ -oimg syn_img/ -omsk syn_msk/
'''

# Construct argument parser
ap = argparse.ArgumentParser()
ap.add_argument("-iimg","--input_image_path", required=True, help="path to source image directory")
ap.add_argument("-imsk","--input_mask_path", required=True, help="path to  source mask directory")
ap.add_argument("-ibgd","--input_background_path", required=True, help="path to background image directory")
ap.add_argument("-oimg","--output_image_path", required=True, help="path to destination image directory")
ap.add_argument("-omsk","--output_mask_path", required=True, help="path to destination mask image directory")
args= vars(ap.parse_args())

img_path = args["input_image_path"]
msk_path = args["input_mask_path"]

bgd_path = args["input_background_path"]

syn_img  = args["output_image_path"]
syn_msk  = args["output_mask_path"]
 
dirs_img = os.listdir( img_path )
dirs_img.sort()

dirs_msk = os.listdir( msk_path )
dirs_msk.sort()

dirs_bgd = os.listdir( bgd_path )
dirs_bgd.sort()

# Ensure same name for corresponding mask and image
def synthesize():
    for item in dirs_img:
        if os.path.isfile(img_path+item):
            
            # Convert all images to RGB format
            img = Image.open(img_path+item).convert('RGB')
            msk = Image.open(msk_path+item).convert('RGB')
            # Choose a random background
            bgd = Image.open(bgd_path+random.choice(dirs_bgd)).convert('RGB')
            
            # Resize images according to need
            img = img.resize((128,128), Image.ANTIALIAS)
            msk = msk.resize((128,128), Image.ANTIALIAS)
            bgd = bgd.resize((128,128), Image.ANTIALIAS)

            # Normalize images and convert masks to one channel
            img = np.array(img)/255.0
            msk = np.array(msk)[:,:,0].reshape(128,128,1)/255.0
            bgd = np.array(bgd)/255.0
            
            # Create image composite with new background
            synimg = bgd*(1.0-msk) + img*msk
            
            # Save new images and masks
            imsave(syn_img+"syn_"+item, synimg)
            imsave(syn_msk+"syn_"+item, np.squeeze(msk))
          
if __name__ == "__main__":           
    sythesize()