#!/usr/bin/python

import numpy as np
import cv2
import matplotlib.pyplot as plt

def to_rgb(img_list):
    rgb_list = []
    rgb_r = []
    rgb_g = []
    rgb_b = []
    for i in range(len(img_list)):
        rgb_list.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2RGB))
        rgb_r.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2RGB)[:, :, 0])
        rgb_g.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2RGB)[:, :, 1])
        rgb_b.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2RGB)[:, :, 2])
    return rgb_list, rgb_r, rgb_g, rgb_b

def to_hls(img_list):
    hls_list = []
    hls_h = []
    hls_l = []
    hls_s = []
    for i in range(len(img_list)):
        hls_list.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2HLS))
        hls_h.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2HLS)[:, :, 0])
        hls_l.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2HLS)[:, :, 1])
        hls_s.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2HLS)[:, :, 2])
    return hls_list, hls_h, hls_l, hls_s

def to_hsv(img_list):
    hsv_list = []
    hsv_h = []
    hsv_s = []
    hsv_v = []
    for i in range(len(img_list)):
        hsv_list.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2HSV))
        hsv_h.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2HSV)[:, :, 0])
        hsv_s.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2HSV)[:, :, 1])
        hsv_v.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2HSV)[:, :, 2])
    return hsv_list, hsv_h, hsv_s, hsv_v

def to_lab(img_list):
    lab_list = []
    lab_l = []
    lab_a = []
    lab_b = []
    for i in range(len(img_list)):
        lab_list.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2LAB))
        lab_l.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2LAB)[:, :, 0])
        lab_a.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2LAB)[:, :, 1])
        lab_b.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2LAB)[:, :, 2])
    return lab_list, lab_l, lab_a, lab_b

def to_luv(img_list):
    luv_list = []
    luv_l = []
    luv_u = []
    luv_v = []
    for i in range(len(img_list)):
        luv_list.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2LUV))
        luv_l.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2LUV)[:, :, 0])
        luv_u.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2LUV)[:, :, 1])
        luv_v.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2LUV)[:, :, 2])
    return luv_list, luv_l, luv_u, luv_v

def to_xyz(img_list):
    xyz_list = []
    xyz_x = []
    xyz_y = []
    xyz_z = []
    for i in range(len(img_list)):
        xyz_list.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2XYZ))
        xyz_x.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2XYZ)[:, :, 0])
        xyz_y.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2XYZ)[:, :, 1])
        xyz_z.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2XYZ)[:, :, 2])
    return xyz_list, xyz_x, xyz_y, xyz_z

def to_yuv(img_list):
    yuv_list = []
    yuv_y = []
    yuv_u = []
    yuv_v = []
    for i in range(len(img_list)):
        yuv_list.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2YUV))
        yuv_y.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2YUV)[:, :, 0])
        yuv_u.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2YUV)[:, :, 1])
        yuv_v.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2YUV)[:, :, 2])
    return yuv_list, yuv_y, yuv_u, yuv_v

def to_ycrcb(img_list):
    ycrcb_list = []
    ycrcb_y = []
    ycrcb_cr = []
    ycrcb_cb = []
    for i in range(len(img_list)):
        ycrcb_list.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2YCrCb))
        ycrcb_y.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2YCrCb)[:, :, 0])
        ycrcb_cr.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2YCrCb)[:, :, 1])
        ycrcb_cb.append(cv2.cvtColor(img_list[i], cv2.COLOR_BGR2YCrCb)[:, :, 2])
    return ycrcb_list, ycrcb_y, ycrcb_cr, ycrcb_cb


nimgs = 6
img_list = [cv2.imread('../test_images/image0' + str(i) + '.png') for i in range(1, nimgs + 1)]

rgb_list, rgb_r, rgb_g, rgb_b = to_rgb(img_list)
hls_list, hls_h, hls_l, hls_s = to_hls(img_list)
hsv_list, hsv_h, hsv_s, hsv_v = to_hsv(img_list)
lab_list, lab_l, lab_a, lab_b = to_lab(img_list)
luv_list, luv_l, luv_u, luv_v = to_luv(img_list)
xyz_list, xyz_x, xyz_y, xyz_z = to_xyz(img_list)
yuv_list, yuv_y, yuv_u, yuv_v = to_xyz(img_list)
ycrcb_list, ycrcb_y, ycrcb_cr, ycrcb_cb = to_ycrcb(img_list)

rows = 4
cols = 4

figure, axes = plt.subplots(nrows=rows, ncols=cols, figsize=(8, 8))

axes[0, 0].imshow(rgb_list[0])
axes[0, 0].set_title('RGB', fontsize=30)
axes[0, 1].imshow(rgb_r[0], cmap='gray')
axes[0, 1].set_title('R', fontsize=30)
axes[0, 2].imshow(rgb_g[0], cmap='gray')
axes[0, 2].set_title('G', fontsize=30)
axes[0, 3].imshow(rgb_b[0], cmap='gray')
axes[0, 3].set_title('B', fontsize=30)

axes[1, 0].imshow(hls_list[0])
axes[1, 0].set_title('HLS', fontsize=30)
axes[1, 1].imshow(hls_h[0], cmap='gray')
axes[1, 1].set_title('H', fontsize=30)
axes[1, 2].imshow(hls_l[0], cmap='gray')
axes[1, 2].set_title('L', fontsize=30)
axes[1, 3].imshow(hls_s[0], cmap='gray')
axes[1, 3].set_title('S', fontsize=30)

axes[2, 0].imshow(hsv_list[0])
axes[2, 0].set_title('HSV', fontsize=30)
axes[2, 1].imshow(hsv_h[0], cmap='gray')
axes[2, 1].set_title('H', fontsize=30)
axes[2, 2].imshow(hsv_s[0], cmap='gray')
axes[2, 2].set_title('S', fontsize=30)
axes[2, 3].imshow(hsv_v[0], cmap='gray')
axes[2, 3].set_title('V', fontsize=30)

axes[3, 0].imshow(lab_list[0])
axes[3, 0].set_title('LAB', fontsize=30)
axes[3, 1].imshow(lab_l[0], cmap='gray')
axes[3, 1].set_title('L', fontsize=30)
axes[3, 2].imshow(lab_a[0], cmap='gray')
axes[3, 2].set_title('A', fontsize=30)
axes[3, 3].imshow(lab_b[0], cmap='gray')
axes[3, 3].set_title('B', fontsize=30)
"""
axes[0, 0].imshow(luv_list[0])
axes[0, 0].set_title('LUV', fontsize=30)
axes[0, 1].imshow(luv_l[0], cmap='gray')
axes[0, 1].set_title('L', fontsize=30)
axes[0, 2].imshow(luv_u[0], cmap='gray')
axes[0, 2].set_title('U', fontsize=30)
axes[0, 3].imshow(luv_v[0], cmap='gray')
axes[0, 3].set_title('V', fontsize=30)

axes[1, 0].imshow(xyz_list[0])
axes[1, 0].set_title('XYZ', fontsize=30)
axes[1, 1].imshow(xyz_x[0], cmap='gray')
axes[1, 1].set_title('X', fontsize=30)
axes[1, 2].imshow(xyz_y[0], cmap='gray')
axes[1, 2].set_title('Y', fontsize=30)
axes[1, 3].imshow(xyz_z[0], cmap='gray')
axes[1, 3].set_title('Z', fontsize=30)

axes[2, 0].imshow(yuv_list[0])
axes[2, 0].set_title('YUV', fontsize=30)
axes[2, 1].imshow(yuv_y[0], cmap='gray')
axes[2, 1].set_title('Y', fontsize=30)
axes[2, 2].imshow(yuv_u[0], cmap='gray')
axes[2, 2].set_title('U', fontsize=30)
axes[2, 3].imshow(yuv_v[0], cmap='gray')
axes[2, 3].set_title('V', fontsize=30)

axes[3, 0].imshow(ycrcb_list[0])
axes[3, 0].set_title('YCrCb', fontsize=30)
axes[3, 1].imshow(ycrcb_y[0], cmap='gray')
axes[3, 1].set_title('Y', fontsize=30)
axes[3, 2].imshow(ycrcb_cr[0], cmap='gray')
axes[3, 2].set_title('Cr', fontsize=30)
axes[3, 3].imshow(ycrcb_cb[0], cmap='gray')
axes[3, 3].set_title('Cb', fontsize=30)
"""
for i in range(0, rows):
    for j in range(0, cols):
        axes[i, j].axis('off')

plt.tight_layout()
plt.show()