import numpy as np
import gradio as gr
from skimage import color
from sklearn.decomposition import TruncatedSVD
from PIL import Image

def truncated_svd_compress(image, k):
    """Compress the image using Truncated SVD by keeping only the top k singular values."""
    svd = TruncatedSVD(n_components=k)
    U = svd.fit_transform(image)
    S = svd.singular_values_
    Vt = svd.components_
    compressed_image = np.dot(U, np.dot(np.diag(S), Vt))
    return compressed_image

def process_image(image, k):
    """Process the uploaded image, compress it using Truncated SVD, and return the result."""
    # Convert PIL Image to NumPy array
    image_np = np.array(image)
    
    # Convert to grayscale for SVD
    gray_image = color.rgb2gray(image_np)
    
    # Ensure the image is 2D for SVD
    if len(gray_image.shape) == 3:
        gray_image = gray_image[:, :, 0]
    
    # Compress the image using Truncated SVD
    compressed_image = truncated_svd_compress(gray_image, k)
    
    # Normalize the compressed image to the range [0, 255] and convert to uint8
    compressed_image = np.clip(compressed_image, 0, 1)  # Ensure values are within [0, 1]
    compressed_image = (compressed_image * 255).astype(np.uint8)
    
    # Convert compressed image back to PIL Image for Gradio output
    compressed_image_pil = Image.fromarray(compressed_image)
    
    # Ensure the PIL Image is in RGB mode for consistent display
    if compressed_image_pil.mode != 'RGB':
        compressed_image_pil = compressed_image_pil.convert('RGB')
    
    return compressed_image_pil

# Gradio interface
gr_interface = gr.Interface(
    fn=process_image,  # Image processing function
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Slider(1, 100, step=1, value=50, label="Compression Rank")  # Compression rank slider
    ],
    outputs=gr.Image(type="pil", label="Compressed Image"),
    title="Interactive Image Compression using Truncated SVD",
    description="Upload an image and adjust the compression rank to see the compressed version. The app compresses the image while retaining important features."
)

# Launch the Gradio interface
gr_interface.launch()