'''
Produce an approximation of a color-image by low-rank approximating
each color channel separately.
Displays each of the three color channel approximation and their sum.
Press any key in an active image window to quit.
'''

import numpy as np
import cv2


file = "macaw.jpg"
rank=100


image = cv2.imread(file)

def low_rank_approximation(image, rank=100):
    U, S, Vt = np.linalg.svd(image, full_matrices=False)
    U = U[:, :rank]
    S = np.diag(S[:rank])
    Vt = Vt[:rank, :]
    approx_image = np.dot(U, np.dot(S, Vt))
    approx_image[approx_image < 0] = 0
    approx_image[approx_image > 255] = 255
    return approx_image.astype('uint8')

channels = cv2.split(image)
approx_channels = []
for i, channel in enumerate(channels):
    approx_channels.append(low_rank_approximation(channel, rank=rank))

color_names = ['Red', 'Green', 'Blue']
for i, channel in enumerate(approx_channels):
    color_image = np.zeros(image.shape, dtype=np.uint8)
    color_image[:,:,i] = channel
    cv2.imshow(f"{color_names[i]} Channel", color_image)

merged_image = cv2.merge(approx_channels)
cv2.imshow("Merged", merged_image)

cv2.waitKey(0)
cv2.destroyAllWindows()