# Please install the following libraries before runing code
# pip install numpy matplotlib opencv-contrib-python
import cv2
import numpy as np
import matplotlib.pyplot as plt

# Read the image
I = cv2.imread('image.png')

# Convert the image to grayscale and normalize into [0, 1]
I = np.array(cv2.cvtColor(I, cv2.COLOR_BGR2GRAY) / 255)

# Display original image
fig, axs = plt.subplots(2, 2)
axs[0, 0].imshow(I, cmap='gray')
axs[0, 0].set_title('Original Image')
axs[0, 0].axis('off')

# Perform SVD
U, S, V = np.linalg.svd(I, full_matrices=False)

for i, k in enumerate([5, 10, 50]):
    # Drop singular values after `k`
    S_k = S.copy()
    S_k[k:] = 0

    # Reconstruct
    I_k = U @ np.diag(S_k) @ V
    
    # Display image k
    axs_k = axs[(i+1) // 2, (i+1) % 2]
    axs_k.imshow(I_k, cmap='gray')
    axs_k.set_title(f'Reconstructed Image (k={k})')
    axs_k.axis('off')


plt.show()
# Save image
plt.savefig('result.png')