text_plain_output_1.png
from PIL import Image

import cv2

import matplotlib.pyplot as plt

import numpy as np

import os

import seaborn as sns

EPOCHS = 10
BATCH_SIZE = 17
HEIGHT = 256
WIDTH = 256
N_CLASSES = 13

def LoadImage(name, path):
    img = Image.open(os.path.join(path, name))
    img = np.array(img)
    image = img[:, :256]
    mask = img[:, 256:]
    return (image, mask)
def bin_image(mask):
    bins = np.array([20, 40, 60, 80, 100, 120, 140, 160, 180, 200, 220, 240])
    new_mask = np.digitize(mask, bins)
    return new_mask
def getSegmentationArr(image, classes, width=WIDTH, height=HEIGHT):
    seg_labels = np.zeros((height, width, classes))
    img = image[:, :, 0]
    for c in range(classes):
        seg_labels[:, :, c] = (img == c).astype(int)
    return seg_labels
def give_color_to_seg_img(seg, n_classes=N_CLASSES):
    seg_img = np.zeros((seg.shape[0], seg.shape[1], 3)).astype('float')
    colors = sns.color_palette('hls', n_classes)
    for c in range(n_classes):
        segc = seg == c
        seg_img[:, :, 0] += segc * colors[c][0]
        seg_img[:, :, 1] += segc * colors[c][1]
        seg_img[:, :, 2] += segc * colors[c][2]
    return seg_img
classes = 13

train_folder = '../input/cityscapes-image-pairs/cityscapes_data/train'
valid_folder = '../input/cityscapes-image-pairs/cityscapes_data/val'
num_of_training_samples = len(os.listdir(train_folder))
num_of_valid_samples = len(os.listdir(valid_folder))
def DataGenerator(path, batch_size=BATCH_SIZE, classes=N_CLASSES):
    files = os.listdir(path)
    while True:
        for i in range(0, len(files), batch_size):
            batch_files = files[i:i + batch_size]
            imgs = []
            segs = []
            for file in batch_files:
                image, mask = LoadImage(file, path)
                mask_binned = bin_image(mask)
                labels = getSegmentationArr(mask_binned, classes)
                imgs.append(image)
                segs.append(labels)
            yield (np.array(imgs), np.array(segs))

train_gen = DataGenerator(train_folder, batch_size=BATCH_SIZE)
val_gen = DataGenerator(valid_folder, batch_size=BATCH_SIZE)

imgs, segs = next(train_gen)
(imgs.shape, segs.shape)

image = imgs[2]
mask = give_color_to_seg_img(np.argmax(segs[2], axis=-1))
masked_image = cv2.addWeighted(image / 255, 0.5, mask, 0.5, 0)
fig, axs = plt.subplots(1, 3, figsize=(20, 20))
axs[0].imshow(image)
axs[0].set_title('Original Image')
axs[1].imshow(mask)
axs[1].set_title('Segmentation Mask')
axs[2].imshow(masked_image)
axs[2].set_title('Masked Image')
plt.show()