text_plain_output_1.png
from sklearn.cluster import KMeans

import cv2

import glob

import matplotlib.pyplot as plt

import numpy as np

import os

INPUT_PATH = '../input/cityscapes-image-pairs/cityscapes_data'

BATCH_SIZE = 4
INPUT_IMG_SIZE = 256
OUTPUT_CLASSES = 10
LEARINING_RATE = 0.001

train_files = glob.glob(os.path.join(INPUT_PATH + '/train', '*jpg'))
val_files = glob.glob(os.path.join(INPUT_PATH + '/val', '*jpg'))

fig, axes = plt.subplots(1,2, figsize = (20,5))
for i, ax in enumerate(axes.ravel()):
    img = cv2.imread(train_files[i])
    ax.imshow(img)
    #ax.axis('off')
plt.show()

for i in train_files[:2]:
    org_img = cv2.imread(i)
    img = org_img[:, 0:256, :]
    msk = org_img[:, 256:, :]

kmeans_data = []
for i in train_files[:50]:
    org_img = cv2.imread(i)
    msk = org_img[:, INPUT_IMG_SIZE:, :]
    kmeans_data.append(msk)
kmeans_data = np.array(kmeans_data)
kmeans_data = kmeans_data.reshape(-1, 3)

encoder = KMeans(n_clusters=OUTPUT_CLASSES)
encoder.fit(kmeans_data)

colors = {0: [255, 0, 0], 1: [0, 255, 0], 2: [0, 0, 255], 3: [255, 69, 0], 4: [255, 0, 255], 5: [210, 105, 30], 6: [192, 255, 62], 7: [127, 255, 0], 8: [0, 238, 238], 9: [72, 118, 255]}

for i in train_files[:3]:
    org_img = cv2.imread(i)
    msk = org_img[:, 256:, :]
    test = msk.reshape(-1, 3)
    pred = encoder.predict(test)
    enc_pred = pred.reshape(INPUT_IMG_SIZE, INPUT_IMG_SIZE)
    pred = np.array([colors[p] for p in pred]).reshape(256, 256, 3)
    print('No of classes in encoded mask:', np.unique(enc_pred))
    plt.figure(figsize=(15, 10))
    plt.subplot(1, 2, 1)
    plt.imshow(msk)
    plt.title('Original mask (RGB)')
    plt.subplot(1, 2, 2)
    plt.imshow(pred)
    plt.title('Encoded mask')
    plt.show()